Update app.py

app.py

@@ -188,132 +188,132 @@ with gr.Blocks() as demo:
     with gr.Row():
         gr.Markdown("""# D-Edit""")
 
-    if 1:
-        with gr.Row():
-            with gr.Column():
-                canvas = gr.Image(value = None, type="numpy",  label="Show Mask", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
-                example_inps = [['./img.png'],['./img2.png'],['./img3.png'],['./img4.png']]
-                gr.Examples(examples=example_inps, inputs=[canvas],
-                            label='examples', cache_examples='lazy', outputs=[],
-                            fn=change_image)
-                gr.Markdown(f"Each image must first undergo segmentation. Afterwards, you can modify the \n mask ID and the prompt for image editing, then proceed with the editing process. \n The link of D-edit paper: [https://arxiv.org/abs/2403.04880v2](https://arxiv.org/abs/2403.04880v2), [https://huggingface.co/papers/2403.04880](https://huggingface.co/papers/2403.04880)")
 
-            with gr.Column():
-                result_info0 = gr.Text(label="Response")
-                segment_button  = gr.Button("Step 1. Run segmentation")
-                flag = gr.State(False)
+    with gr.Row():
+        with gr.Column():
+            canvas = gr.Image(value = None, type="numpy",  label="Show Mask", show_label=True, height=LENGTH, width=LENGTH, interactive=True)
+            example_inps = [['./img.png'],['./img2.png'],['./img3.png'],['./img4.png']]
+            gr.Examples(examples=example_inps, inputs=[canvas],
+                        label='examples', cache_examples='lazy', outputs=[],
+                        fn=change_image)
+            gr.Markdown(f"Each image must first undergo segmentation. Afterwards, you can modify the \n mask ID and the prompt for image editing, then proceed with the editing process. \n The link of D-edit paper: [https://arxiv.org/abs/2403.04880v2](https://arxiv.org/abs/2403.04880v2), [https://huggingface.co/papers/2403.04880](https://huggingface.co/papers/2403.04880)")
+
+        with gr.Column():
+            result_info0 = gr.Text(label="Response")
+            segment_button  = gr.Button("Step 1. Run segmentation")
+            flag = gr.State(False)
+
+            # mask_np_list_updated.value = copy.deepcopy(mask_np_list.value) #!!
+            mask_np_list_updated = mask_np_list
+            gr.Markdown("""<p style="text-align: center; font-size: 20px">Edit Mask (Do not change it during the editing process)</p>""")
+            slider =  gr.Slider(0, 20, step=1, label = 'mask id',  visible=False)
+            label = gr.Text(label='label')
 
-                # mask_np_list_updated.value = copy.deepcopy(mask_np_list.value) #!!
-                mask_np_list_updated = mask_np_list
-                gr.Markdown("""<p style="text-align: center; font-size: 20px">Edit Mask (Do not change it during the editing process)</p>""")
-                slider =  gr.Slider(0, 20, step=1, label = 'mask id',  visible=False)
-                label = gr.Text(label='label')
+        
+            
 
+            result_info = gr.Text(label="Response")
             
+            opt_flag = gr.State(0)
+            gr.Markdown("""<p style="text-align: center; font-size: 20px">Optimization settings</p>""")
+            with gr.Accordion(label="Advanced settings", open=False):
+                num_tokens = gr.Number(value="5", label="num tokens to represent each object", interactive= True)
+                num_tokens_global = num_tokens
+                embedding_learning_rate = gr.Textbox(value="0.00025", label="Embedding optimization: Learning rate", interactive= True )
+                max_emb_train_steps =  gr.Number(value="6", label="embedding optimization: Training steps", interactive= True )
                 
-
-                result_info = gr.Text(label="Response")
+                diffusion_model_learning_rate = gr.Textbox(value="0.0002", label="UNet Optimization: Learning rate", interactive= True )
+                max_diffusion_train_steps = gr.Number(value="28", label="UNet Optimization: Learning rate: Training steps", interactive= True )
                 
-                opt_flag = gr.State(0)
-                gr.Markdown("""<p style="text-align: center; font-size: 20px">Optimization settings</p>""")
+                train_batch_size = gr.Number(value="20", label="Batch size", interactive= True )
+                gradient_accumulation_steps=gr.Number(value="2", label="Gradient accumulation", interactive= True )
+            
+            def run_optimization_wrapper (
+                    mask_np_list,
+                    mask_label_list,
+                    image,
+                    opt_flag,                     
+                    num_tokens,
+                    embedding_learning_rate , 
+                    max_emb_train_steps , 
+                    diffusion_model_learning_rate , 
+                    max_diffusion_train_steps,
+                    train_batch_size,
+                    gradient_accumulation_steps,
+            ):
+                try:
+                    run_optimization = partial(
+                        run_main,  
+                        mask_np_list=mask_np_list, 
+                        mask_label_list=mask_label_list,
+                        image_gt=np.array(image),
+                        num_tokens=int(num_tokens),
+                        embedding_learning_rate = float(embedding_learning_rate), 
+                        max_emb_train_steps = int(max_emb_train_steps), 
+                        diffusion_model_learning_rate= float(diffusion_model_learning_rate), 
+                        max_diffusion_train_steps = int(max_diffusion_train_steps),
+                        train_batch_size=int(train_batch_size),
+                        gradient_accumulation_steps=int(gradient_accumulation_steps)
+                    )
+                    run_optimization()
+                    gr.Info("Optimization Finished! Move to the next step.")
+                    return "Optimization finished! Move to the next step."#,gr.Button("Step 3. Run Editing",interactive = True)
+                except Exception as e:
+                    print(e)
+                    gr.Error("e")
+                    return "Error: use a smaller batch size or try latter."#,gr.Button("Step 3. Run Editing",interactive = False)
+
+
+
+    if 1:
+        with gr.Row():
+            with gr.Column():
+                canvas_text_edit = gr.Image(value = None, type = "pil", label="Editing results", show_label=True,visible = True)
+                # canvas_text_edit = gr.Gallery(label = "Edited results")
+            
+            with gr.Column():
+                gr.Markdown("""<p style="text-align: center; font-size: 20px">Editing setting</p>""")
+                tgt_prompt =  gr.Textbox(value="text prompt", label="Editing: Text prompt", interactive= True )
                 with gr.Accordion(label="Advanced settings", open=False):
-                    num_tokens = gr.Number(value="5", label="num tokens to represent each object", interactive= True)
-                    num_tokens_global = num_tokens
-                    embedding_learning_rate = gr.Textbox(value="0.00025", label="Embedding optimization: Learning rate", interactive= True )
-                    max_emb_train_steps =  gr.Number(value="6", label="embedding optimization: Training steps", interactive= True )
-                    
-                    diffusion_model_learning_rate = gr.Textbox(value="0.0002", label="UNet Optimization: Learning rate", interactive= True )
-                    max_diffusion_train_steps = gr.Number(value="28", label="UNet Optimization: Learning rate: Training steps", interactive= True )
-                    
-                    train_batch_size = gr.Number(value="20", label="Batch size", interactive= True )
-                    gradient_accumulation_steps=gr.Number(value="2", label="Gradient accumulation", interactive= True )
+                    slider2 = gr.Slider(0, 20, step=1, label = 'mask id',  visible=False)
+                    guidance_scale = gr.Textbox(value="5", label="Editing: CFG guidance scale", interactive= True )
+                    num_sampling_steps = gr.Number(value="20", label="Editing: Sampling steps", interactive= True )
+                    edge_thickness = gr.Number(value="10", label="Editing: Edge thickness", interactive= True )
+                    strength = gr.Textbox(value="0.5", label="Editing: Mask strength", interactive= True )
                 
-                def run_optimization_wrapper (
+                add_button  = gr.Button("Step 2. Run Editing",interactive = True)
+                def run_edit_text_wrapper(
                         mask_np_list,
                         mask_label_list,
                         image,
-                        opt_flag,                     
                         num_tokens,
-                        embedding_learning_rate , 
-                        max_emb_train_steps , 
-                        diffusion_model_learning_rate , 
-                        max_diffusion_train_steps,
-                        train_batch_size,
-                        gradient_accumulation_steps,
+                        guidance_scale,
+                        num_sampling_steps ,
+                        strength ,
+                        edge_thickness,
+                        tgt_prompt ,
+                        tgt_index
                 ):
-                    try:
-                        run_optimization = partial(
-                            run_main,  
-                            mask_np_list=mask_np_list, 
-                            mask_label_list=mask_label_list,
-                            image_gt=np.array(image),
-                            num_tokens=int(num_tokens),
-                            embedding_learning_rate = float(embedding_learning_rate), 
-                            max_emb_train_steps = int(max_emb_train_steps), 
-                            diffusion_model_learning_rate= float(diffusion_model_learning_rate), 
-                            max_diffusion_train_steps = int(max_diffusion_train_steps),
-                            train_batch_size=int(train_batch_size),
-                            gradient_accumulation_steps=int(gradient_accumulation_steps)
-                        )
-                        run_optimization()
-                        gr.Info("Optimization Finished! Move to the next step.")
-                        return "Optimization finished! Move to the next step."#,gr.Button("Step 3. Run Editing",interactive = True)
-                    except Exception as e:
-                        print(e)
-                        gr.Error("e")
-                        return "Error: use a smaller batch size or try latter."#,gr.Button("Step 3. Run Editing",interactive = False)
-
-
-
-        
-            with gr.Row():
-                with gr.Column():
-                    canvas_text_edit = gr.Image(value = None, type = "pil", label="Editing results", show_label=True,visible = True)
-                    # canvas_text_edit = gr.Gallery(label = "Edited results")
-                
-                with gr.Column():
-                    gr.Markdown("""<p style="text-align: center; font-size: 20px">Editing setting</p>""")
-                    tgt_prompt =  gr.Textbox(value="text prompt", label="Editing: Text prompt", interactive= True )
-                    with gr.Accordion(label="Advanced settings", open=False):
-                        slider2 = gr.Slider(0, 20, step=1, label = 'mask id',  visible=False)
-                        guidance_scale = gr.Textbox(value="5", label="Editing: CFG guidance scale", interactive= True )
-                        num_sampling_steps = gr.Number(value="20", label="Editing: Sampling steps", interactive= True )
-                        edge_thickness = gr.Number(value="10", label="Editing: Edge thickness", interactive= True )
-                        strength = gr.Textbox(value="0.5", label="Editing: Mask strength", interactive= True )
-                    
-                    add_button  = gr.Button("Step 2. Run Editing",interactive = True)
-                    def run_edit_text_wrapper(
-                            mask_np_list,
-                            mask_label_list,
-                            image,
-                            num_tokens,
-                            guidance_scale,
-                            num_sampling_steps ,
-                            strength ,
-                            edge_thickness,
-                            tgt_prompt ,
-                            tgt_index
-                    ):
-                            
-                        run_edit_text = partial(
-                            run_main,
-                            mask_np_list=mask_np_list, 
-                            mask_label_list=mask_label_list,
-                            image_gt=np.array(image),
-                            load_trained=True,
-                            text=True,
-                            num_tokens = int(num_tokens_global.value),
-                            guidance_scale = float(guidance_scale),
-                            num_sampling_steps = int(num_sampling_steps),
-                            strength = float(strength),
-                            edge_thickness = int(edge_thickness),
-                            num_imgs = 1,
-                            tgt_prompt = tgt_prompt,
-                            tgt_index = int(tgt_index)
-                        )
-                        run_edit_text()
-                        gr.Info('Image editing completed.')
-                        return load_pil_img()
+                        
+                    run_edit_text = partial(
+                        run_main,
+                        mask_np_list=mask_np_list, 
+                        mask_label_list=mask_label_list,
+                        image_gt=np.array(image),
+                        load_trained=True,
+                        text=True,
+                        num_tokens = int(num_tokens_global.value),
+                        guidance_scale = float(guidance_scale),
+                        num_sampling_steps = int(num_sampling_steps),
+                        strength = float(strength),
+                        edge_thickness = int(edge_thickness),
+                        num_imgs = 1,
+                        tgt_prompt = tgt_prompt,
+                        tgt_index = int(tgt_index)
+                    )
+                    run_edit_text()
+                    gr.Info('Image editing completed.')
+                    return load_pil_img()