1. Save in state only inversion results

app.py

@@ -35,8 +35,9 @@ if device == "cuda":
 
 scheduler_class = MyEulerAncestralDiscreteScheduler
 
-pipe_inversion = SDXLDDIMPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)  # .to('cpu')
-pipe_inference = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)  # .to('cpu')
+pipe_inversion = SDXLDDIMPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True).to(device)
+pipe_inference = AutoPipelineForImage2Image.from_pretrained("stabilityai/sdxl-turbo",
+                                                            use_safetensors=True).to(device)
 pipe_inference.scheduler = scheduler_class.from_config(pipe_inference.scheduler.config)
 pipe_inversion.scheduler = scheduler_class.from_config(pipe_inversion.scheduler.config)
 pipe_inversion.scheduler_inference = scheduler_class.from_config(pipe_inference.scheduler.config)
@@ -57,31 +58,27 @@ with gr.Blocks(css="style.css") as demo:
     editor_state = gr.State()
 
 
-    @spaces.GPU
-    def set_pipe(image_editor, input_image, description_prompt, edit_guidance_scale, num_inference_steps=4,
-                 num_inversion_steps=4, inversion_max_step=0.6, rnri_iterations=2, rnri_alpha=0.1, rnri_lr=0.2):
-        
+    # @spaces.GPU
+    def set_pipe1(image_editor, input_image, description_prompt, edit_guidance_scale, num_inference_steps=4,
+                  num_inversion_steps=4, inversion_max_step=0.6, rnri_iterations=2, rnri_alpha=0.1, rnri_lr=0.2):
+
         if device == 'cuda':
-            if image_editor is not None:
-                image_editor = image_editor.to('cpu')
+            # if image_editor is not None:
+            #     image_editor = image_editor.to('cpu')
 
             torch.cuda.empty_cache()
-            
+
         if input_image is None or not description_prompt:
             return None, "Please set all inputs."
 
-        print('### description_prompt ', description_prompt)
-        print('### input_image.size ', input_image.size)
-        
-        if isinstance(num_inference_steps, str): num_inference_steps = int(num_inference_steps) 
-        if isinstance(num_inversion_steps, str): num_inversion_steps = int(num_inversion_steps) 
-        if isinstance(edit_guidance_scale, str): edit_guidance_scale = float(edit_guidance_scale)  
-        if isinstance(inversion_max_step, str): inversion_max_step = float(inversion_max_step) 
-        if isinstance(rnri_iterations, str): rnri_iterations = int(rnri_iterations)  
+        if isinstance(num_inference_steps, str): num_inference_steps = int(num_inference_steps)
+        if isinstance(num_inversion_steps, str): num_inversion_steps = int(num_inversion_steps)
+        if isinstance(edit_guidance_scale, str): edit_guidance_scale = float(edit_guidance_scale)
+        if isinstance(inversion_max_step, str): inversion_max_step = float(inversion_max_step)
+        if isinstance(rnri_iterations, str): rnri_iterations = int(rnri_iterations)
         if isinstance(rnri_alpha, str): rnri_alpha = float(rnri_alpha)
-        if isinstance(rnri_lr, str): rnri_lr = float(rnri_lr) 
-        
-        
+        if isinstance(rnri_lr, str): rnri_lr = float(rnri_lr)
+
         config = RunConfig(num_inference_steps=num_inference_steps,
                            num_inversion_steps=num_inversion_steps,
                            edit_guidance_scale=edit_guidance_scale,
@@ -92,8 +89,37 @@ with gr.Blocks(css="style.css") as demo:
         return image_editor, "Input has set!"
 
 
-    @spaces.GPU
-    def edit(editor, target_prompt):
+    # @spaces.GPU
+    def set_pipe(inversion_state, input_image, description_prompt, edit_guidance_scale, num_inference_steps=4,
+                 num_inversion_steps=4, inversion_max_step=0.6, rnri_iterations=2, rnri_alpha=0.1, rnri_lr=0.2):
+        if device == 'cuda':
+            # if image_editor is not None:
+            #     image_editor = image_editor.to('cpu')
+
+            torch.cuda.empty_cache()
+
+        if input_image is None or not description_prompt:
+            return None, "Please set all inputs."
+
+        if isinstance(num_inference_steps, str): num_inference_steps = int(num_inference_steps)
+        if isinstance(num_inversion_steps, str): num_inversion_steps = int(num_inversion_steps)
+        if isinstance(edit_guidance_scale, str): edit_guidance_scale = float(edit_guidance_scale)
+        if isinstance(inversion_max_step, str): inversion_max_step = float(inversion_max_step)
+        if isinstance(rnri_iterations, str): rnri_iterations = int(rnri_iterations)
+        if isinstance(rnri_alpha, str): rnri_alpha = float(rnri_alpha)
+        if isinstance(rnri_lr, str): rnri_lr = float(rnri_lr)
+
+        config = RunConfig(num_inference_steps=num_inference_steps,
+                           num_inversion_steps=num_inversion_steps,
+                           edit_guidance_scale=edit_guidance_scale,
+                           inversion_max_step=inversion_max_step)
+        inversion_state = ImageEditorDemo.invert(pipe_inversion, input_image, description_prompt, config,
+                                                 [rnri_iterations, rnri_alpha, rnri_lr], device)
+        return inversion_state, "Input has set!"
+
+
+    # @spaces.GPU
+    def edit1(editor, target_prompt):
         if editor is None:
             raise gr.Error("Set inputs before editing.")
         # if device == "cuda":
@@ -102,6 +128,19 @@ with gr.Blocks(css="style.css") as demo:
         image = editor.edit(target_prompt)
         return image
 
+
+    # @spaces.GPU
+    def edit(inversion_state, target_prompt):
+        if inversion_state is None:
+            raise gr.Error("Set inputs before editing.")
+        # if device == "cuda":
+        #       image = editor.to(device).edit(target_prompt)
+        #  else:
+        image = ImageEditorDemo.edit(pipe_inference, target_prompt, inversion_state['latent'], inversion_state['noise'],
+                                     inversion_state['cfg'], inversion_state['cfg'].edit_guidance_scale)
+        return image
+
+
     with gr.Row():
         with gr.Column(elem_id="col-container-1"):
             with gr.Row():
@@ -165,7 +204,6 @@ with gr.Blocks(css="style.css") as demo:
                     )
 
             with gr.Row():
-            #     set_.button = gr.Button("Set input image & description & settings", scale=1)
                 is_set_text = gr.Text("", show_label=False)
 
         with gr.Column(elem_id="col-container-2"):
@@ -191,6 +229,8 @@ with gr.Blocks(css="style.css") as demo:
                             inversion_max_step, rnri_iterations, rnri_alpha, rnri_lr],
                 )
 
+    gr.Markdown(f"""Disclaimer: Performance may be inferior to the reported in the paper due to hardware limitation.
+    """)
     input_image.change(set_pipe,
                        inputs=[editor_state, input_image, description_prompt, edit_guidance_scale, num_inference_steps,
                                num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha, rnri_lr],

src/editor.py

@@ -11,9 +11,11 @@ from diffusers.utils.torch_utils import randn_tensor
 def inversion_callback(pipe, step, timestep, callback_kwargs):
     return callback_kwargs
 
+
 def inference_callback(pipe, step, timestep, callback_kwargs):
     return callback_kwargs
 
+
 def center_crop(im):
     width, height = im.size  # Get dimensions
     min_dim = min(width, height)
@@ -36,67 +38,62 @@ def load_im_into_format_from_path(im_path):
 
 class ImageEditorDemo:
     def __init__(self, pipe_inversion, pipe_inference, input_image, description_prompt, cfg, device, inv_hp):
-        self.pipe_inversion = pipe_inversion
-        self.pipe_inference = pipe_inference
         self.original_image = load_im_into_format_from_path(input_image).convert("RGB")
-        self.load_image = True
+
+
+        # self.pipe_inversion = self.pipe_inversion.to(device)
+        # self.last_latent = self.invert(pipe_inversion, self.original_image, description_prompt)
+
+        # if device  == 'cuda':
+        # after the inversion, we can move the inversion model to the CPU
+        # self.pipe_inversion = self.pipe_inversion.to('cpu')
+
+        # self.pipe_inference = self.pipe_inference.to(device)
+
+    @staticmethod
+    def invert(pipe_inversion, init_image, base_prompt, cfg, inv_hp, device):
+        init_image = load_im_into_format_from_path(init_image).convert("RGB")
         g_cpu = torch.Generator().manual_seed(7865)
-        img_size = (512,512)
+        img_size = (512, 512)
         VQAE_SCALE = 8
         latents_size = (1, 4, img_size[0] // VQAE_SCALE, img_size[1] // VQAE_SCALE)
         noise = [randn_tensor(latents_size, dtype=torch.float16, device=torch.device(device), generator=g_cpu) for i
                  in range(cfg.num_inversion_steps)]
+        pipe_inversion.cfg = cfg
         pipe_inversion.scheduler.set_noise_list(noise)
         pipe_inversion.scheduler_inference.set_noise_list(noise)
         pipe_inversion.set_progress_bar_config(disable=True)
-        self.cfg = cfg
-        self.pipe_inversion.cfg = cfg
-        self.pipe_inference.cfg = cfg
-        self.inv_hp = inv_hp  # [2, 0.1, 0.2]
-        self.edit_cfg = cfg.edit_guidance_scale
-
-        self.pipe_inversion = self.pipe_inversion.to(device)
-        self.last_latent = self.invert(self.original_image, description_prompt)
-        self.original_latent = self.last_latent
-
-        # if device  == 'cuda':
-            # after the inversion, we can move the inversion model to the CPU
-        self.pipe_inversion = self.pipe_inversion.to('cpu')
 
+        res = pipe_inversion(prompt=base_prompt,
+                             num_inversion_steps=cfg.num_inversion_steps,
+                             num_inference_steps=cfg.num_inference_steps,
+                             image=init_image,
+                             guidance_scale=cfg.inversion_guidance_scale,
+                             strength=cfg.inversion_max_step,
+                             denoising_start=1.0 - cfg.inversion_max_step,
+                             inv_hp=inv_hp)[0][0]
+        return {"latent": res, "noise": noise, "cfg": cfg}
+
+    @staticmethod
+    def edit(pipe_inference, target_prompt, last_latent, noise, cfg, edit_cfg):
+        pipe_inference.cfg = cfg
         pipe_inference.scheduler.set_noise_list(noise)
         pipe_inference.set_progress_bar_config(disable=True)
-        self.pipe_inference = self.pipe_inference.to(device)
-
-    def invert(self, init_image, base_prompt):
-        res = self.pipe_inversion(prompt=base_prompt,
-                             num_inversion_steps=self.cfg.num_inversion_steps,
-                             num_inference_steps=self.cfg.num_inference_steps,
-                             image=init_image,
-                             guidance_scale=self.cfg.inversion_guidance_scale,
-                             callback_on_step_end=inversion_callback,
-                             strength=self.cfg.inversion_max_step,
-                             denoising_start=1.0 - self.cfg.inversion_max_step,
-                             inv_hp=self.inv_hp)[0][0]
-        return res
-
-    def edit(self, target_prompt):
-        image = self.pipe_inference(prompt=target_prompt,
-                            num_inference_steps=self.cfg.num_inference_steps,
-                            negative_prompt="",
-                            callback_on_step_end=inference_callback,
-                            image=self.last_latent,
-                            strength=self.cfg.inversion_max_step,
-                            denoising_start=1.0 - self.cfg.inversion_max_step,
-                            guidance_scale=self.edit_cfg).images[0]
+        image = pipe_inference(prompt=target_prompt,
+                               num_inference_steps=cfg.num_inference_steps,
+                               negative_prompt="",
+                               image=last_latent,
+                               strength=cfg.inversion_max_step,
+                               denoising_start=1.0 - cfg.inversion_max_step,
+                               guidance_scale=edit_cfg).images[0]
         return image
 
-    def to(self, device):
-        self.pipe_inference = self.pipe_inference.to(device)
-        self.pipe_inversion = self.pipe_inversion.to(device)
-        self.last_latent = self.last_latent.to(device)
-        self.original_latent = self.original_latent.to(device)
-        
-        self.pipe_inversion.scheduler.set_noise_list_device(device)
-        self.pipe_inference.scheduler.set_noise_list_device(device)
-        self.pipe_inversion.scheduler_inference.set_noise_list_device(device)
-        return self
+    # def to(self, device):
+    #     self.pipe_inference = self.pipe_inference.to(device)
+    #     self.pipe_inversion = self.pipe_inversion.to(device)
+    #     self.last_latent = self.last_latent.to(device)
+    #
+    #     self.pipe_inversion.scheduler.set_noise_list_device(device)
+    #     self.pipe_inference.scheduler.set_noise_list_device(device)
+    #     self.pipe_inversion.scheduler_inference.set_noise_list_device(device)
+    #     return self