Update pipeline_semantic_stable_diffusion_img2img_solver.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, prompt_tar="", cfg_scale_tar=15, skip=36, eta=1):
+def sample(zs, wts, attention_store, prompt_tar="", cfg_scale_tar=15, skip=36, eta=1):
     latents = wts[-1].expand(1, -1, -1, -1)
-    img = pipe(
+    img, attention_store = pipe(
         prompt=prompt_tar,
         init_latents=latents,
         guidance_scale=cfg_scale_tar,
@@ -45,9 +45,10 @@ def sample(zs, wts, prompt_tar="", cfg_scale_tar=15, skip=36, eta=1):
         # num_inference_steps=steps,
         # use_ddpm=True,
         # wts=wts.value,
+        attention_store = attention_store,
         zs=zs,
-    ).images[0]
-    return img
+    )
+    return img.images[0], attention_store
 
 
 def reconstruct(
@@ -57,6 +58,7 @@ def reconstruct(
     skip,
     wts,
     zs,
+    attention_store,
     do_reconstruction,
     reconstruction,
     reconstruct_button,
@@ -77,8 +79,8 @@ def reconstruct(
             ):  # if image caption was not changed, run actual reconstruction
                 tar_prompt = ""
             latents = wts[-1].expand(1, -1, -1, -1)
-            reconstruction = sample(
-                zs, wts, prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale
+            reconstruction, attention_store = sample(
+                zs, wts, attention_store=attention_store, prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale
             )
             do_reconstruction = False
         return (
@@ -128,7 +130,7 @@ def load_and_invert(
 ## SEGA ##
 
 def edit(input_image,
-            wts, zs,
+            wts, zs, attention_store,
             tar_prompt,
             image_caption,
             steps,
@@ -195,27 +197,27 @@ def edit(input_image,
       )
 
       latnets = wts[-1].expand(1, -1, -1, -1)
-      sega_out = pipe(prompt=tar_prompt, 
+      sega_out, attention_store = 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, **editing_args)
+                          zs=zs, attention_store=attention_store, **editing_args)
       
-      return sega_out.images[0], gr.update(visible=True), do_reconstruction, reconstruction, wts, zs, do_inversion, show_share_button
+      return sega_out.images[0], gr.update(visible=True), do_reconstruction, reconstruction, wts, zs, attention_store, 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 = sample(zs, wts, prompt_tar=tar_prompt, skip=skip, cfg_scale_tar=tar_cfg_scale)
+          pure_ddpm_img, attention_store = sample(zs, wts, attention_store=attention_store, 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, do_inversion, show_share_button
+          return pure_ddpm_img, 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, do_inversion, show_share_button
+      return reconstruction, gr.update(visible=False), do_reconstruction, reconstruction, wts, zs, attention_store, do_inversion, show_share_button
         
 
 def randomize_seed_fn(seed, is_random):
@@ -458,6 +460,7 @@ with gr.Blocks(css="style.css") as demo:
     gr.HTML(intro)
     wts = gr.State()
     zs = gr.State()
+    attention_store=gr.State()
     reconstruction = gr.State()
     do_inversion = gr.State(value=True)
     do_reconstruction = gr.State(value=True)
@@ -693,7 +696,7 @@ with gr.Blocks(css="style.css") as demo:
     run_button.click(
         fn=edit,
         inputs=[input_image,
-                wts, zs,
+                wts, zs, attention_store,
                 tar_prompt,
                 image_caption,
                 steps,
@@ -713,7 +716,7 @@ with gr.Blocks(css="style.css") as demo:
 
 
         ],
-        outputs=[sega_edited_image, reconstruct_button, do_reconstruction, reconstruction, wts, zs, do_inversion, share_btn_container])
+        outputs=[sega_edited_image, reconstruct_button, do_reconstruction, reconstruction, wts, zs,attention_store, do_inversion, share_btn_container])
     # .success(fn=update_gallery_display, inputs= [prev_output_image, sega_edited_image], outputs = [gallery, gallery, prev_output_image])
 
 

pipeline_semantic_stable_diffusion_img2img_solver.py

@@ -499,6 +499,7 @@ class SemanticStableDiffusionImg2ImgPipeline_DPMSolver(DiffusionPipeline):
             verbose=True,
             use_cross_attn_mask: bool = False,
             # Attention store (just for visualization purposes)
+            attention_store = None,
             attn_store_steps: Optional[List[int]] = [],
             store_averaged_over_steps: bool = True,
             use_intersect_mask: bool = False,
@@ -771,8 +772,8 @@ class SemanticStableDiffusionImg2ImgPipeline_DPMSolver(DiffusionPipeline):
             timesteps = timesteps[-zs.shape[0]:]
 
         if use_cross_attn_mask:
-            self.attention_store = AttentionStore(average=store_averaged_over_steps, batch_size=batch_size)
-            self.prepare_unet(self.attention_store, PnP=False)
+            attention_store = AttentionStore(average=store_averaged_over_steps, batch_size=batch_size)
+            self.prepare_unet(attention_store, PnP=False)
         # 5. Prepare latent variables
         num_channels_latents = self.unet.config.in_channels
         latents = self.prepare_latents(
@@ -917,8 +918,8 @@ class SemanticStableDiffusionImg2ImgPipeline_DPMSolver(DiffusionPipeline):
                             noise_guidance_edit_tmp = noise_guidance_edit_tmp * user_mask
 
                         if use_cross_attn_mask:
-                            out = self.attention_store.aggregate_attention(
-                                attention_maps=self.attention_store.step_store,
+                            out = attention_store.aggregate_attention(
+                                attention_maps=attention_store.step_store,
                                 prompts=self.text_cross_attention_maps,
                                 res=16,
                                 from_where=["up", "down"],
@@ -1080,7 +1081,7 @@ class SemanticStableDiffusionImg2ImgPipeline_DPMSolver(DiffusionPipeline):
                 store_step = i in attn_store_steps
                 if store_step:
                     print(f"storing attention for step {i}")
-                self.attention_store.between_steps(store_step)
+                attention_store.between_steps(store_step)
 
             # call the callback, if provided
             if callback is not None and i % callback_steps == 0:
@@ -1102,9 +1103,9 @@ class SemanticStableDiffusionImg2ImgPipeline_DPMSolver(DiffusionPipeline):
         image = self.image_processor.postprocess(image, output_type=output_type, do_denormalize=do_denormalize)
 
         if not return_dict:
-            return (image, has_nsfw_concept)
+            return (image, has_nsfw_concept), attention_store
 
-        return SemanticStableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept)
+        return SemanticStableDiffusionPipelineOutput(images=image, nsfw_content_detected=has_nsfw_concept),  attention_store
 
     def encode_text(self, prompts):
         text_inputs = self.tokenizer(