Changed layout in app.py

app.py

@@ -56,25 +56,7 @@ if torch.cuda.is_available():
 # Functions
 
 def get_guide(image):  
-    return hed(image,scribble=True)
-
-def set_mask(image):
-    img=image['image'][...,:3]
-    mask=1*(image['mask'][...,:3]>0)
-    # save vars
-    CURRENT_IMAGE['image']=img
-    CURRENT_IMAGE['mask']=mask
-
-    guide=get_guide(img)
-    CURRENT_IMAGE['guide']=np.array(guide)
-    guide=255-np.asarray(guide)  
-
-    seg_img = guide*(1-mask) + mask*192
-    preview = img * (seg_img==255)
-
-    vis_image=(preview/2).astype(seg_img.dtype) + seg_img * (seg_img!=255)
-
-    return vis_image
+  return hed(image,scribble=True)
 
 def generate(image,
              prompt,
@@ -83,14 +65,14 @@ def generate(image,
              sketch_scale,
              seed):
 
-    sketch=(255*(image['mask'][...,:3]>0)).astype(CURRENT_IMAGE['image'].dtype)
-    mask=CURRENT_IMAGE['mask']
-
-    CURRENT_IMAGE['guide']=(CURRENT_IMAGE['guide']*(mask==0) + sketch*(mask!=0)).astype(CURRENT_IMAGE['image'].dtype)
+  sketch=(255*(image['mask'][...,:3]>0)).astype(CURRENT_IMAGE['image'].dtype)
+  mask=CURRENT_IMAGE['mask']
+  
+  CURRENT_IMAGE['guide']=(CURRENT_IMAGE['guide']*(mask==0) + sketch*(mask!=0)).astype(CURRENT_IMAGE['image'].dtype)
 
-    mask_img=255*CURRENT_IMAGE['mask'].astype(CURRENT_IMAGE['image'].dtype)
+  mask_img=255*CURRENT_IMAGE['mask'].astype(CURRENT_IMAGE['image'].dtype)
 
-    new_image = pipe(
+  new_image = pipe(
       prompt,
       num_inference_steps=num_steps,
       guidance_scale=text_scale,
@@ -99,72 +81,68 @@ def generate(image,
       control_image=Image.fromarray(CURRENT_IMAGE['guide']),
       controlnet_conditioning_scale=sketch_scale,
       mask_image=Image.fromarray(mask_img)
-    ).images
+  ).images#[0]
 
-    return new_image
+  return new_image
 
 def create_demo(max_images=12, default_num_images=3):
-  
-    with gr.Blocks(theme=gr.themes.Default(font=[gr.themes.GoogleFont("IBM Plex Mono"), "ui-monospace","monospace"])
-                  ) as demo:
+    with gr.Blocks(theme=gr.themes.Default(font=[gr.themes.GoogleFont("IBM Plex Mono"), "ui-monospace","monospace"])) as demo:
         gr.Markdown('## Cut and Sketch ✂️▶️✏️')
-        gr.Markdown('**Usage**')
-        gr.Markdown('1. Upload your image to the left window')
-        gr.Markdown('2. Draw the mask in the left window (Cut ✂️)')
-        gr.Markdown('3. Click `Set Mask`')
-        gr.Markdown('4. In the right window, sketch a replacement object! (Sketch ✏️)')
-        gr.Markdown('5. (You can also provide a text prompt if you want)')
-        gr.Markdown('6. 🔮 Click Generate! ')
-        
-        prompt = gr.Textbox(label='Prompt')
-
-        with gr.Row():
-            with gr.Column():
-                with gr.Row():
-                    input_image = gr.Image(source='upload',
-                                            shape=[HEIGHT,WIDTH],
-                                            type='numpy',
-                                          label='Mask Draw',
-                                            tool='sketch',
-                                            brush_radius=70)
-                    sketch_image = gr.Image(source='upload',
-                                            shape=[HEIGHT,WIDTH],
-                                            type='numpy',
-                                          label='Fill Draw',
-                                            tool='sketch',
-                                            brush_radius=15)
-                with gr.Row():
-                    mask_button = gr.Button(label='Set Mask', value='Set Mask')
-                    run_button = gr.Button(label='Generate', value='Generate')
+    
+        with gr.Column() as step_1:
+            gr.Markdown('**Start Here**')
+            gr.Markdown('1. Upload your image below')
+            gr.Markdown('2. **Draw the mask** for the region you want changed (Cut ✂️)')
+            input_image = gr.Image(source='upload',
+                                    shape=[HEIGHT,WIDTH],
+                                    type='numpy',
+                                  label='Mask Draw (Cut!)',
+                                    tool='sketch',
+                                    brush_radius=80)
+            gr.Markdown('3. Click `Set Mask` when it is ready!')
+            mask_button = gr.Button(label='Set Mask', value='Set Mask')
+        with gr.Column(visible=False) as step_2:
+            gr.Markdown('4. Now, you can **sketch a replacement** object! (Sketch ✏️)')
+            sketch_image = gr.Image(source='upload',
+                                          shape=[HEIGHT,WIDTH],
+                                          type='numpy',
+                                        label='Fill Draw (Sketch!)',
+                                          tool='sketch',
+                                          brush_radius=20)
+            gr.Markdown('5. (You can also provide a **text prompt** if you want)')
+            prompt = gr.Textbox(label='Prompt')   
+            gr.Markdown('6. 🔮 Click `Generate` when ready! ')   
+            run_button = gr.Button(label='Generate', value='Generate')         
             output_image = gr.Gallery(
-                        label="Generated images",
-                        show_label=False,
-                        elem_id="gallery",
-                    )
-
-        with gr.Accordion('Advanced options', open=False):
-            num_steps = gr.Slider(label='Steps',
-                                minimum=1,
-                                maximum=100,
-                                value=20,
-                                step=1)
-            text_scale = gr.Slider(label='Text Guidance Scale',
-                                      minimum=0.1,
-                                      maximum=30.0,
-                                      value=7.5,
-                                      step=0.1)
-            seed = gr.Slider(label='Seed',
-                            minimum=-1,
-                            maximum=2147483647,
-                            step=1,
-                            randomize=True)  
-
-            sketch_scale = gr.Slider(label='Sketch Guidance Scale',
-                                      minimum=0.0,
-                                      maximum=1.0,
-                                      value=1.0,
-                                      step=0.05)                  
-
+                              label="Generated images",
+                              visible=False,
+                              show_label=False,
+                              elem_id="gallery",
+                          )#.style(grid=(1, 2))
+        
+            with gr.Accordion('Advanced options', open=False):
+                num_steps = gr.Slider(label='Steps',
+                                      minimum=1,
+                                      maximum=100,
+                                      value=20,
+                                      step=1)
+                text_scale = gr.Slider(label='Text Guidance Scale',
+                                            minimum=0.1,
+                                            maximum=30.0,
+                                            value=7.5,
+                                            step=0.1)
+                seed = gr.Slider(label='Seed',
+                                  minimum=-1,
+                                  maximum=2147483647,
+                                  step=1,
+                                  randomize=True)  
+            
+                sketch_scale = gr.Slider(label='Sketch Guidance Scale',
+                                            minimum=0.0,
+                                            maximum=1.0,
+                                            value=1.0,
+                                            step=0.05)                  
+          
         inputs = [
           sketch_image,
           prompt,
@@ -173,10 +151,28 @@ def create_demo(max_images=12, default_num_images=3):
           sketch_scale,
           seed
         ]
-
-        mask_button.click(fn=set_mask, inputs=input_image, outputs=sketch_image)     
+        
+        def set_mask(image):
+            img=image['image'][...,:3]
+            mask=1*(image['mask'][...,:3]>0)
+            # save vars
+            CURRENT_IMAGE['image']=img
+            CURRENT_IMAGE['mask']=mask
+            
+            guide=get_guide(img)
+            CURRENT_IMAGE['guide']=np.array(guide)
+            guide=255-np.asarray(guide)  
+            
+            seg_img = guide*(1-mask) + mask*192
+            preview = img * (seg_img==255)
+            
+            vis_image=(preview/2).astype(seg_img.dtype) + seg_img * (seg_img!=255)
+            
+            return {input_image : image['image'], sketch_image : vis_image, step_2: gr.update(visible=True)}
+        
+        mask_button.click(fn=set_mask, inputs=[input_image], outputs=[input_image, sketch_image,step_2])     
         run_button.click(fn=generate, inputs=inputs, outputs=output_image)
-        return demo
+    return demo
 
 if __name__ == '__main__':
     demo = create_demo()