fix inversion

gradio_app.py

@@ -18,7 +18,6 @@ This demo supports both generated images and real images. To modify a real image
 '''
 
 stable, stable_config = setup(LPMConfig())
-stable_for_inversion, _ = setup(LPMConfig())
 
 def main_pipeline(
         prompt: str,
@@ -48,7 +47,7 @@ def main_pipeline(
             real_image_path="" if input_image is None else input_image
         )
 
-        result_images, result_proxy_words = main(stable, stable_config, stable_for_inversion, args)
+        result_images, result_proxy_words = main(stable, stable_config, args)
         result_images = [im.permute(1, 2, 0).cpu().numpy() for im in result_images]
         result_images = [(im * 255).astype(np.uint8) for im in result_images]
         result_images = [Image.fromarray(im) for im in result_images]

main.py

@@ -1,13 +1,14 @@
 import json
 import os
+from dataclasses import dataclass, field
+from typing import List
+
 import pyrallis
 import torch
-from dataclasses import dataclass, field
 from torch.utils.data import DataLoader
 from torchvision.transforms import ToTensor
 from torchvision.utils import save_image
 from tqdm import tqdm
-from typing import List
 
 from src.diffusion_model_wrapper import DiffusionModelWrapper, get_stable_diffusion_model, get_stable_diffusion_config, \
     generate_original_image
@@ -34,7 +35,7 @@ def setup(args):
     return ldm_stable, ldm_stable_config
 
 
-def main(ldm_stable, ldm_stable_config, ldm_stable_inversion, args):
+def main(ldm_stable, ldm_stable_config, args):
 
     similar_words, prompts, another_prompts = get_proxy_prompts(args, ldm_stable)
     exp_path = save_args_dict(args, similar_words)
@@ -44,7 +45,8 @@ def main(ldm_stable, ldm_stable_config, ldm_stable_inversion, args):
     uncond_embeddings = None
 
     if args.real_image_path != "":
-        x_t, uncond_embeddings = invert_image(args, ldm_stable_inversion, ldm_stable_config, prompts, exp_path)
+        ldm_stable, ldm_stable_config = setup(args)
+        x_t, uncond_embeddings = invert_image(args, ldm_stable, ldm_stable_config, prompts, exp_path)
 
     image, x_t, orig_all_latents, orig_mask, average_attention = generate_original_image(args, ldm_stable, ldm_stable_config, prompts, x_t, uncond_embeddings)
     save_image(ToTensor()(image[0]), f"{exp_path}/{similar_words[0]}.jpg")

src/diffusion_model_wrapper.py

@@ -1,13 +1,13 @@
-import torch
-import numpy as np
 from typing import Optional, List
 
+import numpy as np
+import torch
+from cv2 import dilate
 from diffusers import DDIMScheduler, StableDiffusionPipeline
 from tqdm import tqdm
-from cv2 import dilate
 
-from src.attention_utils import show_cross_attention
 from src.attention_based_segmentation import Segmentor
+from src.attention_utils import show_cross_attention
 from src.prompt_to_prompt_controllers import DummyController, AttentionStore
 
 
@@ -136,7 +136,7 @@ class DiffusionModelWrapper:
                 if self.enbale_attn_controller_changes:
                     attn = self.controller(attn, is_cross, place_in_unet)
                 
-                if is_cross and context[1] is not None and self.prompt_mixing is not None:
+                if is_cross and self.prompt_mixing is not None and context[1] is not None:
                     attn = self.prompt_mixing.get_cross_attn(self, self.diff_step, attn, place_in_unet, batch_size)
 
                 if not is_cross and (not self.model_config["low_resource"] or not self.uncond_pred) and self.prompt_mixing is not None: