runs without errors

.gitignore

@@ -6,3 +6,4 @@
 *.iml
 out
 gen
+*.pyc

app.py

@@ -1,7 +1,5 @@
 import gradio as gr
-import numpy as np
-import random
-from diffusers import DiffusionPipeline
+
 import torch
 from src.euler_scheduler import MyEulerAncestralDiscreteScheduler
 from diffusers.pipelines.auto_pipeline import AutoPipelineForImage2Image
@@ -11,15 +9,13 @@ from src.editor import ImageEditorDemo
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-
 scheduler_class = MyEulerAncestralDiscreteScheduler
 
-
 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)
+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)
 
 
 # if torch.cuda.is_available():
@@ -32,104 +28,110 @@ pipe_inversion.scheduler_inference  = scheduler_class.from_config(pipe_inference
 #     pipe = pipe.to(device)
 
 
-def infer(input_image, description_prompt, target_prompt, guidance_scale, num_inference_steps=4, num_inversion_steps=4, inversion_max_step=0.6):
+def infer(input_image, description_prompt, target_prompt, edit_guidance_scale, num_inference_steps=4,
+          num_inversion_steps=4,
+          inversion_max_step=0.6):
     config = RunConfig(num_inference_steps=num_inference_steps,
                        num_inversion_steps=num_inversion_steps,
-                       guidance_scale=guidance_scale,
+                       edit_guidance_scale=edit_guidance_scale,
                        inversion_max_step=inversion_max_step)
-    
-    editor = ImageEditorDemo(pipe_inversion, pipe_inference, input_image, description_prompt, config)
+
+    editor = ImageEditorDemo(pipe_inversion, pipe_inference, input_image, description_prompt, config, device)
 
     image = editor.edit(target_prompt)
     return image
 
+
 examples = [
     "Astronaut in a jungle, cold color palette, muted colors, detailed, 8k",
     "An astronaut riding a green horse",
     "A delicious ceviche cheesecake slice",
 ]
 
-css="""
-#col-container {
-    margin: 0 auto;
-    max-width: 520px;
-}
-"""
+# css = """
+# #col-container-1 {
+#     margin: 0 auto;
+#     max-width: 520px;
+# }
+# #col-container-2 {
+#     margin: 0 auto;
+#     max-width: 520px;
+# }
+# """
 
 if torch.cuda.is_available():
     power_device = "GPU"
 else:
     power_device = "CPU"
 
-with gr.Blocks(css=css) as demo:
-
+# with gr.Blocks(css=css) as demo:
+with gr.Blocks() as demo:
     gr.Markdown(f"""
-    # RNRI briel and links on device: {power_device}.
+    This is a demo for our [paper]("https://arxiv.org/abs/2312.12540") **RNRI: Regularized Newton Raphson Inversion for Text-to-Image Diffusion Models**.
+    Image editing using our RNRI for inversion demonstrates significant speed-up and improved quality compared to previous state-of-the-art methods.
+    RNRI can be applied to a variety of diffusion models, including SDXL, DDIM, and others.
+    Take a look at our [project page]("https://barakmam.github.io/rnri.github.io/").
     """)
-    with gr.Column(elem_id="col-container"):
-
-        with gr.Row():
-            input_image = gr.Image(label="Input image", sources=['upload', 'webcam', 'clipboard'], type="pil")
-        
-        with gr.Row():
-            
-            description_prompt = gr.Text(
-                label="Image description",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your image description",
-                container=False,
-            )
-          
-        
-        with gr.Row():
-  
-            target_prompt = gr.Text(
-                label="Edit prompt",
-                show_label=False,
-                max_lines=1,
-                placeholder="Enter your edit prompt",
-                container=False,
-            )
-
-        
-        with gr.Accordion("Advanced Settings", open=False):
-            
+    with gr.Row():
+        with gr.Column(elem_id="col-container-1"):
+            with gr.Row():
+                input_image = gr.Image(label="Input image", sources=['upload', 'webcam', 'clipboard'], type="pil")
+
             with gr.Row():
-                
-                guidance_scale = gr.Slider(
-                    label="Guidance scale",
-                    minimum=0.0,
-                    maximum=10.0,
-                    step=0.1,
-                    value=1.2,
+                description_prompt = gr.Text(
+                    label="Image description",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter your image description",
+                    container=False,
                 )
-                
-                num_inference_steps = gr.Slider(
-                    label="Number of RNRI iterations",
-                    minimum=1,
-                    maximum=12,
-                    step=1,
-                    value=4,
+
+            with gr.Row():
+                target_prompt = gr.Text(
+                    label="Edit prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="Enter your edit prompt",
+                    container=False,
                 )
 
-        
-        with gr.Row():
-            run_button = gr.Button("Edit", scale=0)
+            with gr.Accordion("Advanced Settings", open=False):
+                with gr.Row():
+                    edit_guidance_scale = gr.Slider(
+                        label="Guidance scale",
+                        minimum=0.0,
+                        maximum=10.0,
+                        step=0.1,
+                        value=1.2,
+                    )
+
+                    num_inference_steps = gr.Slider(
+                        label="Number of RNRI iterations",
+                        minimum=1,
+                        maximum=12,
+                        step=1,
+                        value=4,
+                    )
 
-    with gr.Column(elem_id="col-container"):
-    
-        result = gr.Image(label="Result", show_label=False)
-    
-        # gr.Examples(
-        #     examples = examples,
-        #     inputs = [prompt]
-        # )
+            with gr.Row():
+                run_button = gr.Button("Edit", scale=1)
+
+        with gr.Column(elem_id="col-container-2"):
+            result = gr.Image(label="Result")
+
+            # gr.Examples(
+            #     examples = examples,
+            #     inputs = [prompt]
+            # )
 
     run_button.click(
-        fn = infer,
-        inputs = [input_image, description_prompt, target_prompt, guidance_scale, num_inference_steps, num_inference_steps],
-        outputs = [result]
+        fn=infer,
+        inputs=[input_image, description_prompt, target_prompt, edit_guidance_scale, num_inference_steps,
+                num_inference_steps],
+        outputs=[result]
     )
 
-demo.queue().launch()
+demo.queue().launch()
+
+# im = infer(input_image, description_prompt, target_prompt, edit_guidance_scale, num_inference_steps=4, num_inversion_steps=4,
+#           inversion_max_step=0.6)

src/config.py

@@ -9,7 +9,8 @@ class RunConfig:
 
     num_inversion_steps: int = 100
 
-    guidance_scale: float = 0.0
+    inversion_guidance_scale: float = 0.0
+    edit_guidance_scale: float = 1.2
 
     inversion_max_step: float = 1.0
 

src/editor.py

@@ -35,7 +35,7 @@ def load_im_into_format_from_path(im_path):
 
 
 class ImageEditorDemo:
-    def __init__(self, pipe_inversion, pipe_inference, input_image, description_prompt, cfg):
+    def __init__(self, pipe_inversion, pipe_inference, input_image, description_prompt, cfg, device):
         self.pipe_inversion = pipe_inversion
         self.pipe_inference = pipe_inference
         self.original_image = load_im_into_format_from_path(input_image).convert("RGB")
@@ -44,7 +44,7 @@ class ImageEditorDemo:
         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("cuda:0"), generator=g_cpu) for i
+        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.scheduler.set_noise_list(noise)
         pipe_inference.scheduler.set_noise_list(noise)
@@ -55,10 +55,10 @@ class ImageEditorDemo:
         self.pipe_inversion.cfg = cfg
         self.pipe_inference.cfg = cfg
         self.inv_hp = [2, 0.1, 0.2]
-        self.edit_cfg = 1.2
+        self.edit_cfg = cfg.edit_guidance_scale
 
-        self.pipe_inference.to("cuda")
-        self.pipe_inversion.to("cuda")
+        self.pipe_inference.to(device)
+        self.pipe_inversion.to(device)
 
         self.last_latent = self.invert(self.original_image, description_prompt)
         self.original_latent = self.last_latent
@@ -68,7 +68,7 @@ class ImageEditorDemo:
                              num_inversion_steps=self.cfg.num_inversion_steps,
                              num_inference_steps=self.cfg.num_inference_steps,
                              image=init_image,
-                             guidance_scale=self.cfg.guidance_scale,
+                             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,

src/sdxl_inversion_pipeline.py

@@ -304,7 +304,10 @@ class SDXLDDIMPipeline(StableDiffusionXLImg2ImgPipeline):
 
     def get_timestamp_dist(self, z_0, timesteps):
         timesteps = timesteps.to(z_0.device)
-        sigma = self.scheduler.sigmas.cuda()[:-1][self.scheduler.timesteps == timesteps]
+        if "cuda" in str(z_0.device):
+            sigma = self.scheduler.sigmas.cuda()[:-1][self.scheduler.timesteps == timesteps]
+        else:
+            sigma = self.scheduler.sigmas[:-1][self.scheduler.timesteps == timesteps]
         z_0 = z_0.reshape(-1, 1)
 
         def gaussian_pdf(x):