Another layout update

app.py

@@ -58,99 +58,83 @@ if torch.cuda.is_available():
 def get_guide(image):  
   return hed(image,scribble=True)
 
-def generate(image,
-             prompt,
-             num_steps,
-             text_scale,
-             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)
-
-  mask_img=255*CURRENT_IMAGE['mask'].astype(CURRENT_IMAGE['image'].dtype)
-
-  new_image = pipe(
-      prompt,
-      num_inference_steps=num_steps,
-      guidance_scale=text_scale,
-      generator=torch.manual_seed(seed),
-      image=Image.fromarray(CURRENT_IMAGE['image']),
-      control_image=Image.fromarray(CURRENT_IMAGE['guide']),
-      controlnet_conditioning_scale=sketch_scale,
-      mask_image=Image.fromarray(mask_img)
-  ).images#[0]
-
-  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:
-        gr.Markdown('## Cut and Sketch ✂️▶️✏️')
-    
-        with gr.Column() as step_1:
-            gr.Markdown('**Start Here**')
+
+        gr.Markdown('# Cut and Sketch ✂️▶️✏️')
+        with gr.Accordion('Instructions', open=False):
+            gr.Markdown('## Cut ✂️')
             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('## Sketch ✏️')
             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')               
-            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)  
+            gr.Markdown('6. 🔮 Click `Generate` when ready! ') 
+            
+        with gr.Row() as main_blocks:
+            with gr.Column() as step_1:
+                gr.Markdown('### Mask Input')   
+                input_image = gr.Image(source='upload',
+                                      shape=[HEIGHT,WIDTH],
+                                      type='numpy',
+                                     elem_id="image-upload",
+                                     label='Mask Draw (Cut!)',
+                                      tool='sketch',
+                                      brush_radius=80)
+                mask_button = gr.Button(label='Set Mask', value='Set Mask',full_width=False)
+            with gr.Column(visible=False) as step_2:     
+                gr.Markdown('### Sketch Input')         
+                sketch_image = gr.Image(source='upload',
+                                      shape=[HEIGHT,WIDTH],
+                                      type='numpy',
+                                      label='Fill Draw (Sketch!)',
+                                      tool='sketch',
+                                      brush_radius=15)
+                run_button = gr.Button(label='Generate', value='Generate') 
+                prompt = gr.Textbox(label='Prompt')    
+        
+            with gr.Column() as output_step:  
+                gr.Markdown('### Output')   
+                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)  
-              
-            output_image = gr.Gallery(
-                              label="Generated image",
-                              show_label=False,
-                              elem_id="gallery",
-                          )#.style(grid=(1, 2))
-          
+            sketch_scale = gr.Slider(label='Sketch Guidance Scale',
+                                      minimum=0.0,
+                                      maximum=1.0,
+                                      value=1.0,
+                                      step=0.05)                  
+        
         inputs = [
-          sketch_image,
-          prompt,
-          num_steps,
-          text_scale,
-          sketch_scale,
-          seed
+            sketch_image,
+            prompt,
+            num_steps,
+            text_scale,
+            sketch_scale,
+            seed
         ]
         
+        # STEP 1: Set Mask
         def set_mask(image):
             img=image['image'][...,:3]
             mask=1*(image['mask'][...,:3]>0)
@@ -167,11 +151,46 @@ def create_demo(max_images=12, default_num_images=3):
             
             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)}
+            return {input_image : image['image'],
+                sketch_image : vis_image,
+                step_1: gr.update(visible=False),
+                step_2: gr.update(visible=True)
+                }
+        
+        # STEP 2: Generate
+        def generate(image,
+             prompt,
+             num_steps,
+             text_scale,
+             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)
+            
+            mask_img=255*CURRENT_IMAGE['mask'].astype(CURRENT_IMAGE['image'].dtype)
+            
+            new_image = pipe(
+              prompt,
+              num_inference_steps=num_steps,
+              guidance_scale=text_scale,
+              generator=torch.manual_seed(seed),
+              image=Image.fromarray(CURRENT_IMAGE['image']),
+              control_image=Image.fromarray(CURRENT_IMAGE['guide']),
+              controlnet_conditioning_scale=sketch_scale,
+              mask_image=Image.fromarray(mask_img)
+            ).images#[0]
+            
+            return {output_image : new_image,
+                step_1: gr.update(visible=True),
+                step_2: gr.update(visible=False)
+                }
         
-        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
+        mask_button.click(fn=set_mask, inputs=[input_image], outputs=[input_image, sketch_image, step_1,step_2])     
+        run_button.click(fn=generate, inputs=inputs, outputs=[output_image, step_1,step_2])
+        return demo
 
 if __name__ == '__main__':
     demo = create_demo()