Update app.py

app.py

@@ -35,9 +35,9 @@ def caption_image(input_image):
     generated_caption = blip_processor.batch_decode(generated_ids, skip_special_tokens=True)[0]
     return generated_caption, generated_caption
 
-def sample(zs, wts, attention_store, prompt_tar="", cfg_scale_tar=15, skip=36, eta=1):
+def sample(zs, wts, attention_store, text_cross_attention_maps, prompt_tar="", cfg_scale_tar=15, skip=36, eta=1):
     latents = wts[-1].expand(1, -1, -1, -1)
-    img, attention_store = pipe(
+    img, attention_store, text_cross_attention_maps = pipe(
         prompt=prompt_tar,
         init_latents=latents,
         guidance_scale=cfg_scale_tar,
@@ -45,10 +45,10 @@ def sample(zs, wts, attention_store, prompt_tar="", cfg_scale_tar=15, skip=36, e
         # num_inference_steps=steps,
         # use_ddpm=True,
         # wts=wts.value,
-        attention_store = attention_store,
+        attention_store = attention_store, text_cross_attention_maps=text_cross_attention_maps,
         zs=zs,
     )
-    return img.images[0], attention_store
+    return img.images[0], attention_store, text_cross_attention_maps
 
 
 def reconstruct(
@@ -59,6 +59,7 @@ def reconstruct(
     wts,
     zs,
     attention_store,
+    text_cross_attention_maps,
     do_reconstruction,
     reconstruction,
     reconstruct_button,
@@ -79,8 +80,8 @@ def reconstruct(
             ):  # if image caption was not changed, run actual reconstruction
                 tar_prompt = ""
             latents = wts[-1].expand(1, -1, -1, -1)
-            reconstruction, attention_store = sample(
-                zs, wts, attention_store=attention_store, prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale
+            reconstruction, attention_store, text_cross_attention_maps = sample(
+                zs, wts, attention_store=attention_store, text_cross_attention_maps=text_cross_attention_maps,prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale
             )
             do_reconstruction = False
         return (
@@ -130,7 +131,7 @@ def load_and_invert(
 ## SEGA ##
 
 def edit(input_image,
-            wts, zs, attention_store,
+            wts, zs, attention_store, text_cross_attention_maps,
             tar_prompt,
             image_caption,
             steps,
@@ -197,27 +198,27 @@ def edit(input_image,
       )
 
       latnets = wts[-1].expand(1, -1, -1, -1)
-      sega_out, attention_store = pipe(prompt=tar_prompt, 
+      sega_out, attention_store, text_cross_attention_maps = pipe(prompt=tar_prompt, 
                           init_latents=latnets, 
                           guidance_scale = tar_cfg_scale,
                           # num_images_per_prompt=1,
                           # num_inference_steps=steps,
                           # use_ddpm=True,  
                           # wts=wts.value, 
-                          zs=zs, attention_store=attention_store, **editing_args)
+                          zs=zs, attention_store=attention_store, text_cross_attention_maps=text_cross_attention_maps, **editing_args)
       
-      return sega_out.images[0], gr.update(visible=True), do_reconstruction, reconstruction, wts, zs, attention_store, do_inversion, show_share_button
+      return sega_out.images[0], gr.update(visible=True), do_reconstruction, reconstruction, wts, zs, attention_store, text_cross_attention_maps, do_inversion, show_share_button
     
     
     else: # if sega concepts were not added, performs regular ddpm sampling
       
       if do_reconstruction: # if ddpm sampling wasn't computed
-          pure_ddpm_img, attention_store = sample(zs, wts, attention_store=attention_store, prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale)
+          pure_ddpm_img, attention_store, text_cross_attention_maps = sample(zs, wts, attention_store=attention_store, text_cross_attention_maps=text_cross_attention_maps, prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale)
           reconstruction = pure_ddpm_img
           do_reconstruction = False
-          return pure_ddpm_img, gr.update(visible=False), do_reconstruction, reconstruction, wts, zs, attention_store, do_inversion, show_share_button
+          return pure_ddpm_img, gr.update(visible=False), do_reconstruction, reconstruction, wts, zs, attention_store, text_cross_attention_maps, do_inversion, show_share_button
       
-      return reconstruction, gr.update(visible=False), do_reconstruction, reconstruction, wts, zs, attention_store, do_inversion, show_share_button
+      return reconstruction, gr.update(visible=False), do_reconstruction, reconstruction, wts, zs, attention_store, text_cross_attention_maps, do_inversion, show_share_button
         
 
 def randomize_seed_fn(seed, is_random):
@@ -461,6 +462,7 @@ with gr.Blocks(css="style.css") as demo:
     wts = gr.State()
     zs = gr.State()
     attention_store=gr.State()
+    text_cross_attention_maps = gr.State()
     reconstruction = gr.State()
     do_inversion = gr.State(value=True)
     do_reconstruction = gr.State(value=True)
@@ -697,6 +699,7 @@ with gr.Blocks(css="style.css") as demo:
         fn=edit,
         inputs=[input_image,
                 wts, zs, attention_store,
+                text_cross_attention_maps,
                 tar_prompt,
                 image_caption,
                 steps,
@@ -716,7 +719,7 @@ with gr.Blocks(css="style.css") as demo:
 
 
         ],
-        outputs=[sega_edited_image, reconstruct_button, do_reconstruction, reconstruction, wts, zs,attention_store, do_inversion, share_btn_container])
+        outputs=[sega_edited_image, reconstruct_button, do_reconstruction, reconstruction, wts, zs,attention_store, text_cross_attention_maps, do_inversion, share_btn_container])
     # .success(fn=update_gallery_display, inputs= [prev_output_image, sega_edited_image], outputs = [gallery, gallery, prev_output_image])