Update app.py

clear fixed + reconstruct

app.py

@@ -71,39 +71,44 @@ def sample(zs, wts, prompt_tar="", cfg_scale_tar=15, skip=36, eta = 1):
     return img
 
 
-def reconstruct(tar_prompt, 
-                tar_cfg_scale, 
-                skip, 
-                wts, zs, 
+def reconstruct(tar_prompt,
+                tar_cfg_scale,
+                skip,
+                wts, zs,
                 do_reconstruction,
-                reconstruction, reconstruct_button, hide_reconstruct_button
+                reconstruction,
+                reconstruct_button
                ):
 
-    if do_reconstruction:
-        reconstruction_img = sample(zs.value, wts.value, prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale)
-        reconstruction = gr.State(value=reconstruction_img)
-        do_reconstruction = False
-    return reconstruction.value, reconstruction, do_reconstruction, reconstruct_button.update(visible=False), hide_reconstruct_button.update(visible=True)
+    if reconstruct_button == "Hide Reconstruction":
+      return reconstruction.value, reconstruction, ddpm_edited_image.update(visible=False), do_reconstruction, "Show Reconstruction"
+
+    else:
+      if do_reconstruction:
+          reconstruction_img = sample(zs.value, wts.value, prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale)
+          reconstruction = gr.State(value=reconstruction_img)
+          do_reconstruction = False
+      return reconstruction.value, reconstruction, ddpm_edited_image.update(visible=True), do_reconstruction, "Hide Reconstruction"
+
 
-    
 def load_and_invert(
-                    input_image, 
+                    input_image,
                     do_inversion,
                     seed, randomize_seed,
-                    wts, zs, 
-                    src_prompt ="", 
-                    tar_prompt="", 
+                    wts, zs,
+                    src_prompt ="",
+                    tar_prompt="",
                     steps=100,
                     src_cfg_scale = 3.5,
                     skip=36,
-                    tar_cfg_scale=15, 
+                    tar_cfg_scale=15,
                     progress=gr.Progress(track_tqdm=True)
-                    
+
 ):
 
-    
+
     x0 = load_512(input_image, device=device)
-    
+
     if do_inversion or randomize_seed:
         # invert and retrieve noise maps and latent
         zs_tensor, wts_tensor = invert(x0 =x0 , prompt_src=src_prompt, num_diffusion_steps=steps, cfg_scale_src=src_cfg_scale)
@@ -111,43 +116,40 @@ def load_and_invert(
         zs = gr.State(value=zs_tensor)
         do_inversion = False
 
-    inversion_progress = "Inversion compeleted!"
-    return wts, zs, do_inversion, inversion_progress
+    return wts, zs, do_inversion, inversion_progress.update(visible=False)
 
 ## SEGA ##
-    
+
 def edit(input_image,
-            wts, zs, 
-            tar_prompt, 
+            wts, zs,
+            tar_prompt,
             steps,
             skip,
             tar_cfg_scale,
-            reconstruct_button,
             edit_concept_1,edit_concept_2,edit_concept_3,
             guidnace_scale_1,guidnace_scale_2,guidnace_scale_3,
             warmup_1, warmup_2, warmup_3,
             neg_guidance_1, neg_guidance_2, neg_guidance_3,
             threshold_1, threshold_2, threshold_3):
-                
-                
+
+
     editing_args = dict(
     editing_prompt = [edit_concept_1,edit_concept_2,edit_concept_3],
     reverse_editing_direction = [ neg_guidance_1, neg_guidance_2, neg_guidance_3,],
     edit_warmup_steps=[warmup_1, warmup_2, warmup_3,],
-    edit_guidance_scale=[guidnace_scale_1,guidnace_scale_2,guidnace_scale_3], 
+    edit_guidance_scale=[guidnace_scale_1,guidnace_scale_2,guidnace_scale_3],
     edit_threshold=[threshold_1, threshold_2, threshold_3],
-    edit_momentum_scale=0.3, 
+    edit_momentum_scale=0.3,
     edit_mom_beta=0.6,
     eta=1,)
-    
+
     latnets = wts.value[skip].expand(1, -1, -1, -1)
     sega_out = sem_pipe(prompt=tar_prompt, latents=latnets, guidance_scale = tar_cfg_scale,
-                        num_images_per_prompt=1,  
-                        num_inference_steps=steps, 
+                        num_images_per_prompt=1,
+                        num_inference_steps=steps,
                         use_ddpm=True,  wts=wts.value, zs=zs.value[skip:], **editing_args)
-       
-    # return sega_out.images[0], reconstruct_button.update(visible=True)
-    return sega_out.images[0]
+
+    return sega_out.images[0], reconstruct_button.update(visible=True)
                 
 
 
@@ -215,7 +217,7 @@ def get_example():
 
 ########
 # demo #
-########        
+########
 
 intro = """
 <h1 style="font-weight: 1400; text-align: center; margin-bottom: 7px;">
@@ -234,67 +236,74 @@ For faster inference without waiting in queue, you may duplicate the space and u
 
 help_text = """
 - **Getting Started - edit images with DDPM X SEGA:**
-    
-    The are 3 general setting options you can play with - 
-    
+
+    The are 3 general setting options you can play with -
+
     1. **Pure DDPM Edit -** Describe the desired edited output image in detail
-    2. **Pure SEGA Edit -** Keep the target prompt empty ***or*** with a description of the original image and add editing concepts for Semantic Gudiance editing 
-    3. **Combined -** Describe the desired edited output image in detail and add additional SEGA editing concepts on top 
+    2. **Pure SEGA Edit -** Keep the target prompt empty ***or*** with a description of the original image and add editing concepts for Semantic Gudiance editing
+    3. **Combined -** Describe the desired edited output image in detail and add additional SEGA editing concepts on top
 - **Getting Started - Tips**
-    
+
     While the best approach depends on your editing objective and source image,  we can layout a few guiding tips to use as a starting point -
-    
+
     1. **DDPM** is usually more suited for scene/style changes and major subject changes (for example ) while **SEGA** allows for more fine grained control, changes are more delicate, more suited for adding details (for example facial expressions and attributes, subtle style modifications, object adding/removing)
-    2. The more you describe the scene in the target prompt (both the parts and details you wish to keep the same and those you wish to change), the better the result 
-    3. **Combining DDPM Edit with SEGA -** 
-    Try dividing your editing objective to more significant scene/style/subject changes and detail adding/removing and more moderate changes. Then describe the major changes in a detailed target prompt and add the more fine grained details as SEGA concepts. 
+    2. The more you describe the scene in the target prompt (both the parts and details you wish to keep the same and those you wish to change), the better the result
+    3. **Combining DDPM Edit with SEGA -**
+    Try dividing your editing objective to more significant scene/style/subject changes and detail adding/removing and more moderate changes. Then describe the major changes in a detailed target prompt and add the more fine grained details as SEGA concepts.
     4. **Reconstruction:** Using an empty source prompt + target prompt will lead to a perfect reconstruction
 - **Fidelity vs creativity**:
-    
+
     Bigger values → more fidelity, smaller values → more creativity
-    
-    1. `Skip Steps` 
+
+    1. `Skip Steps`
     2. `Warmup` (SEGA)
     3. `Threshold`  (SEGA)
-    
+
     Bigger values → more creativity, smaller values → more fidelity
-    
+
     1. `Guidance Scale`
     2. `Concept Guidance Scale` (SEGA)
 """
 
 with gr.Blocks(css='style.css') as demo:
-    
+
     def add_concept(sega_concepts_counter):
       if sega_concepts_counter == 1:
         return row2.update(visible=True), row3.update(visible=False), add_concept_button.update(visible=True), 2
       else:
         return row2.update(visible=True), row3.update(visible=True), add_concept_button.update(visible=False), 3
 
-        
+
     def reset_do_inversion():
         do_inversion = True
         return do_inversion
 
+    def reset_do_reconstruction():
+      do_reconstruction = True
+      return  do_reconstruction
+
 
     def update_inversion_progress_visibility(do_inversion):
       if do_inversion:
           return inversion_progress.update(visible=True)
       else:
         return inversion_progress.update(visible=False)
-    
-                
 
-        
+
+
+
+
+
     gr.HTML(intro)
     wts = gr.State()
     zs = gr.State()
-    do_inversion = gr.State(value=True)
     reconstruction = gr.State()
+    do_inversion = gr.State(value=True)
+    do_reconstruction = gr.State(value=True)
     sega_concepts_counter = gr.State(1)
 
-  
-    
+
+
     with gr.Row():
         input_image = gr.Image(label="Input Image", interactive=True)
         ddpm_edited_image = gr.Image(label=f"DDPM Reconstructed Image", interactive=False, visible=False)
@@ -304,7 +313,7 @@ with gr.Blocks(css='style.css') as demo:
         sega_edited_image.style(height=365, width=365)
 
     with gr.Row():
-        inversion_progress = gr.Textbox(visible=False)
+        inversion_progress = gr.Textbox(visible=False, label="Inversion progress")
 
     with gr.Tabs() as tabs:
           with gr.TabItem('1. Describe the desired output', id=0):
@@ -321,34 +330,25 @@ with gr.Blocks(css='style.css') as demo:
               with gr.Row().style(mobile_collapse=False, equal_height=True):
                   neg_guidance_1 = gr.Checkbox(
                       label='Negative Guidance')
-                  warmup_1 = gr.Slider(label='Warmup', minimum=0, maximum=50, 
-                                       value=DEFAULT_WARMUP_STEPS, 
-                                       step=1, interactive=True)
-                  guidnace_scale_1 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30, 
-                                               value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE, 
+
+                  guidnace_scale_1 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30,
+                                               value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,
                                                step=0.5, interactive=True)
-                  threshold_1 = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99, 
-                                          value=DEFAULT_THRESHOLD, steps=0.01, interactive=True)
+
                   edit_concept_1 = gr.Textbox(
                                   label="Edit Concept",
                                   show_label=False,
                                   max_lines=1,
                                   placeholder="Enter your 1st edit prompt",
                               )
-                  
+
               # 2nd SEGA concept
               with gr.Row(visible=False) as row2:
                   neg_guidance_2 = gr.Checkbox(
                       label='Negative Guidance',visible=True)
-                  warmup_2 = gr.Slider(label='Warmup', minimum=0, maximum=50, 
-                                       value=DEFAULT_WARMUP_STEPS, 
-                                       step=1, interactive=True)
-                  guidnace_scale_2 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30, 
-                                               value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE, 
+                  guidnace_scale_2 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30,
+                                               value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,
                                                step=0.5, interactive=True)
-                  threshold_2 = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99, 
-                                          value=DEFAULT_THRESHOLD, 
-                                          steps=0.01, interactive=True)
                   edit_concept_2 = gr.Textbox(
                                   label="Edit Concept",
                                   show_label=False,
@@ -359,34 +359,30 @@ with gr.Blocks(css='style.css') as demo:
               with gr.Row(visible=False) as row3:
                   neg_guidance_3 = gr.Checkbox(
                       label='Negative Guidance',visible=True)
-                  warmup_3 = gr.Slider(label='Warmup', minimum=0, maximum=50, 
-                                       value=DEFAULT_WARMUP_STEPS, step=1, 
-                                       interactive=True)
-                  guidnace_scale_3 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30, 
-                                               value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE, 
+
+                  guidnace_scale_3 = gr.Slider(label='Concept Guidance Scale', minimum=1, maximum=30,
+                                               value=DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE,
                                                step=0.5, interactive=True)
-                  threshold_3 = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99, 
-                                          value=DEFAULT_THRESHOLD, steps=0.01,
-                                          interactive=True)
+
                   edit_concept_3 = gr.Textbox(
                                   label="Edit Concept",
                                   show_label=False,
                                   max_lines=1,
                                   placeholder="Enter your 3rd edit prompt",
                               )
-              
+
               with gr.Row().style(mobile_collapse=False, equal_height=True):
                 add_concept_button = gr.Button("+")
 
-                      
+
     with gr.Row():
         run_button = gr.Button("Edit")
         reconstruct_button = gr.Button("Show Reconstruction", visible=False)
-        undo_button = gr.Button("Undo", visible=False)
-
-        clear_button = gr.ClearButton()
+        clear_button = gr.Button("Clear")
 
     with gr.Accordion("Advanced Options", open=False):
+      with gr.Tabs() as tabs:
+          with gr.TabItem('General options', id=2):
             with gr.Row():
                 with gr.Column():
                     src_prompt = gr.Textbox(lines=1, label="Source Prompt", interactive=True, placeholder="")
@@ -394,53 +390,73 @@ with gr.Blocks(css='style.css') as demo:
                     src_cfg_scale = gr.Number(value=3.5, label=f"Source Guidance Scale", interactive=True)
                     seed = gr.Number(value=0, precision=0, label="Seed", interactive=True)
                     randomize_seed = gr.Checkbox(label='Randomize seed', value=False)
-                with gr.Column():    
+                with gr.Column():
                     skip = gr.Slider(minimum=0, maximum=60, value=36, label="Skip Steps", interactive=True)
-                    tar_cfg_scale = gr.Slider(minimum=7, maximum=30,value=15, label=f"Guidance Scale", interactive=True)  
+                    tar_cfg_scale = gr.Slider(minimum=7, maximum=30,value=15, label=f"Guidance Scale", interactive=True)
+          with gr.TabItem('SEGA options', id=3):
+             # 1st SEGA concept
+              with gr.Row().style(mobile_collapse=False, equal_height=True):
+                  warmup_1 = gr.Slider(label='Warmup', minimum=0, maximum=50,
+                                       value=DEFAULT_WARMUP_STEPS,
+                                       step=1, interactive=True)
+                  threshold_1 = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99,
+                                          value=DEFAULT_THRESHOLD, steps=0.01, interactive=True)
+
+              # 2nd SEGA concept
+              with gr.Row(visible=False) as row2_advanced:
+                  warmup_2 = gr.Slider(label='Warmup', minimum=0, maximum=50,
+                                       value=DEFAULT_WARMUP_STEPS,
+                                       step=1, interactive=True)
+                  threshold_2 = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99,
+                                          value=DEFAULT_THRESHOLD,
+                                          steps=0.01, interactive=True)
+              # 3rd SEGA concept
+              with gr.Row(visible=False) as row3_advanced:
+                  warmup_3 = gr.Slider(label='Warmup', minimum=0, maximum=50,
+                                       value=DEFAULT_WARMUP_STEPS, step=1,
+                                       interactive=True)
+                  threshold_3 = gr.Slider(label='Threshold', minimum=0.5, maximum=0.99,
+                                          value=DEFAULT_THRESHOLD, steps=0.01,
+                                          interactive=True)
+
 
 
-    
 
 
     # with gr.Accordion("Help", open=False):
     #     gr.Markdown(help_text)
-    
+
     caption_button.click(
         fn = caption_image,
         inputs = [input_image],
         outputs = [tar_prompt]
     )
-    
+
     add_concept_button.click(fn = add_concept, inputs=sega_concepts_counter,
                outputs= [row2, row3, add_concept_button, sega_concepts_counter], queue = False)
-    
-    run_button.click(
-        fn = randomize_seed_fn,
-        inputs = [seed, randomize_seed],
-        outputs = [seed], 
-        queue = False).then(fn = update_inversion_progress_visibility, inputs =[do_inversion], outputs=[inversion_progress],queue=False).then(
+
+    run_button.click(fn = update_inversion_progress_visibility, inputs =[do_inversion], outputs=[inversion_progress],queue=False).then(
         fn=load_and_invert,
-        inputs=[input_image, 
+        inputs=[input_image,
                 do_inversion,
                 seed, randomize_seed,
-                wts, zs, 
-                src_prompt, 
-                tar_prompt, 
+                wts, zs,
+                src_prompt,
+                tar_prompt,
                 steps,
                 src_cfg_scale,
                 skip,
-                tar_cfg_scale         
+                tar_cfg_scale
         ],
         outputs=[wts, zs, do_inversion, inversion_progress],
     ).then(fn = update_inversion_progress_visibility, inputs =[do_inversion], outputs=[inversion_progress],queue=False).success(
         fn=edit,
-        inputs=[input_image, 
-                wts, zs, 
-                tar_prompt, 
+        inputs=[input_image,
+                wts, zs,
+                tar_prompt,
                 steps,
                 skip,
                 tar_cfg_scale,
-                reconstruct_button,
                 edit_concept_1,edit_concept_2,edit_concept_3,
                 guidnace_scale_1,guidnace_scale_2,guidnace_scale_3,
                 warmup_1, warmup_2, warmup_3,
@@ -448,66 +464,112 @@ with gr.Blocks(css='style.css') as demo:
                 threshold_1, threshold_2, threshold_3
 
         ],
-        # outputs=[sega_edited_image, reconstruct_button]  
-        outputs=[sega_edited_image]  )
+        outputs=[sega_edited_image, reconstruct_button])
 
 
 
     # Automatically start inverting upon input_image change
     input_image.change(
         fn = reset_do_inversion,
-        outputs = [do_inversion], 
-        queue = False).then(fn = update_inversion_progress_visibility, inputs =[do_inversion], outputs=[inversion_progress],queue=False).then(
+        outputs = [do_inversion],
+        queue = False).then(fn = update_inversion_progress_visibility, inputs =[do_inversion],
+                            outputs=[inversion_progress],queue=False).then(
         fn=load_and_invert,
-        inputs=[input_image, 
+        inputs=[input_image,
                 do_inversion,
                 seed, randomize_seed,
-                wts, zs, 
-                src_prompt, 
-                tar_prompt, 
+                wts, zs,
+                src_prompt,
+                tar_prompt,
                 steps,
                 src_cfg_scale,
                 skip,
-                tar_cfg_scale,          
+                tar_cfg_scale,
         ],
         # outputs=[ddpm_edited_image, wts, zs, do_inversion],
         outputs=[wts, zs, do_inversion, inversion_progress],
-    ).then(fn = update_inversion_progress_visibility, inputs =[do_inversion], outputs=[inversion_progress],queue=False)
+    ).then(fn = update_inversion_progress_visibility, inputs =[do_inversion],
+           outputs=[inversion_progress],queue=False).then(
+              lambda: reconstruct_button.update(visible=False),
+              outputs=[reconstruct_button]).then(
+        fn = reset_do_reconstruction,
+        outputs = [do_reconstruction],
+        queue = False)
 
-    
-    # Repeat inversion when these params are changed:
+
+    # Repeat inversion (and reconstruction) when these params are changed:
     src_prompt.change(
         fn = reset_do_inversion,
-        outputs = [do_inversion], queue = False)
-    
-    steps.change(fn = reset_do_inversion,
-        outputs = [do_inversion], queue = False)
-    
+        outputs = [do_inversion], queue = False).then(
+        fn = reset_do_reconstruction,
+        outputs = [do_reconstruction], queue = False)
+
+    steps.change(
+        fn = reset_do_inversion,
+        outputs = [do_inversion], queue = False).then(
+        fn = reset_do_reconstruction,
+        outputs = [do_reconstruction], queue = False)
+
+
+    src_cfg_scale.change(
+        fn = reset_do_inversion,
+        outputs = [do_inversion], queue = False).then(
+        fn = reset_do_reconstruction,
+        outputs = [do_reconstruction], queue = False)
+
+    # Repeat only reconstruction these params are changed:
+
+    tar_prompt.change(
+        fn = reset_do_reconstruction,
+        outputs = [do_reconstruction], queue = False)
+
+    tar_cfg_scale.change(
+        fn = reset_do_reconstruction,
+        outputs = [do_reconstruction], queue = False)
 
-    src_cfg_scale.change(fn = reset_do_inversion,
-        outputs = [do_inversion], queue = False)
+    skip.change(
+        fn = reset_do_reconstruction,
+        outputs = [do_reconstruction], queue = False)
 
-    components_to_clear = [input_image,ddpm_edited_image,sega_edited_image, do_inversion,
+
+
+    clear_components = [input_image,ddpm_edited_image,ddpm_edited_image,sega_edited_image, do_inversion,
                                    src_prompt, steps, src_cfg_scale, seed,
-                                  tar_prompt, skip, tar_cfg_scale,
+                                  tar_prompt, skip, tar_cfg_scale, reconstruct_button,reconstruct_button,
                                   edit_concept_1, guidnace_scale_1,warmup_1,  threshold_1, neg_guidance_1,
                                   edit_concept_2, guidnace_scale_2,warmup_2,  threshold_2, neg_guidance_2,
-                                  edit_concept_3, guidnace_scale_3,warmup_3,  threshold_3, neg_guidance_3,
-                                  
-                                  ]
-    clear_output_vals = [None, None, None, True,
+                                  edit_concept_3, guidnace_scale_3,warmup_3,  threshold_3, neg_guidance_3,]
+
+    clear_components_output_vals = [None, None,ddpm_edited_image.update(visible=False), None, True,
                      "", DEFAULT_DIFFUSION_STEPS, DEFAULT_SOURCE_GUIDANCE_SCALE, DEFAULT_SEED,
-                     "", DEFAULT_SKIP_STEPS, DEFAULT_TARGET_GUIDANCE_SCALE,
+                     "", DEFAULT_SKIP_STEPS, DEFAULT_TARGET_GUIDANCE_SCALE, reconstruct_button.update(value="Show Reconstruction"),reconstruct_button.update(visible=False),
                      "", DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE, DEFAULT_WARMUP_STEPS, DEFAULT_THRESHOLD, DEFAULT_NEGATIVE_GUIDANCE,
                      "", DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE, DEFAULT_WARMUP_STEPS, DEFAULT_THRESHOLD, DEFAULT_NEGATIVE_GUIDANCE,
                      "", DEFAULT_SEGA_CONCEPT_GUIDANCE_SCALE, DEFAULT_WARMUP_STEPS, DEFAULT_THRESHOLD, DEFAULT_NEGATIVE_GUIDANCE,
-                         ]                            
-    clear_button.click(lambda:clear_output_vals, outputs =components_to_clear)
-    
+                         ]
+
+
+    clear_button.click(lambda: clear_components_output_vals, outputs =clear_components)
+
+    reconstruct_button.click(lambda: ddpm_edited_image.update(visible=True), outputs=[ddpm_edited_image]).then(fn = reconstruct,
+                inputs = [tar_prompt,
+                tar_cfg_scale,
+                skip,
+                wts, zs,
+                do_reconstruction,
+                reconstruction,
+                          reconstruct_button],
+                outputs = [ddpm_edited_image,reconstruction, ddpm_edited_image, do_reconstruction, reconstruct_button])
 
+    randomize_seed.change(
+        fn = randomize_seed_fn,
+        inputs = [seed, randomize_seed],
+        outputs = [seed],
+        queue = False)
+    
     # gr.Examples(
-    #     label='Examples', 
-    #     examples=get_example(), 
+    #     label='Examples',
+    #     examples=get_example(),
     #     inputs=[input_image, src_prompt, tar_prompt, steps,
     #                 # src_cfg_scale,
     #                 skip,
@@ -527,6 +589,7 @@ with gr.Blocks(css='style.css') as demo:
 
 
 
+
 demo.queue()
 demo.launch(share=False)