major update:

1. set inputs once and edit on the same inversion.
2. RNRI hyperparams to user

app.py

@@ -27,23 +27,13 @@ pipe_inversion.scheduler_inference = scheduler_class.from_config(pipe_inference.
 #     pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
 #     pipe = pipe.to(device)
 
-@spaces.GPU
-def set_pipe(input_image, description_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,
-                       edit_guidance_scale=edit_guidance_scale,
-                       inversion_max_step=inversion_max_step)
-
-    return ImageEditorDemo(pipe_inversion, pipe_inference, input_image, description_prompt, config, device)
-
     # return editor
     # image = editor.edit(target_prompt)
     # return image
 
 def edit(editor, target_prompt):
     if editor is None:
-        return "Please set image & description before editing."
+        raise gr.Error("Set inputs before editing.")
     return editor.edit(target_prompt)
 
 examples = [
@@ -75,6 +65,20 @@ with gr.Blocks(css="style.css") as demo:
     Image editing using our RNRI for inversion demonstrates significant speed-up and improved quality compared to previous state-of-the-art methods.
     Take a look at our [project page](https://barakmam.github.io/rnri.github.io/).
     """)
+    editor_state = gr.State()
+
+    @spaces.GPU
+    def set_pipe(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):
+        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)
+        return ImageEditorDemo(pipe_inversion, pipe_inference, input_image,
+                               description_prompt, config, device,
+                               [rnri_iterations, rnri_alpha, rnri_lr]), "Input has set!"
+
+
     with gr.Row():
         with gr.Column(elem_id="col-container-1"):
             with gr.Row():
@@ -90,19 +94,6 @@ with gr.Blocks(css="style.css") as demo:
                     container=False,
                 )
 
-            with gr.Row():
-                set_button = gr.Button("Set input image & description", scale=1)
-
-            editor_state = gr.State()
-            with gr.Row():
-                target_prompt = gr.Text(
-                    label="Edit prompt",
-                    info = "Enter your edit prompt",
-                    show_label=False,
-                    max_lines=1,
-                    placeholder="an oreo cake on a table",
-                    container=False,
-                )
 
             with gr.Accordion("Advanced Settings", open=False):
                 with gr.Row():
@@ -115,7 +106,7 @@ with gr.Blocks(css="style.css") as demo:
                     )
 
                     num_inference_steps = gr.Slider(
-                        label="Number of RNRI iterations",
+                        label="Inference steps",
                         minimum=1,
                         maximum=12,
                         step=1,
@@ -130,12 +121,52 @@ with gr.Blocks(css="style.css") as demo:
                         value=0.6,
                     )
 
+                    rnri_iterations = gr.Slider(
+                        label="RNRI iterations",
+                        minimum=0,
+                        maximum=5,
+                        step=1,
+                        value=2,
+                    )
+                    rnri_alpha = gr.Slider(
+                        label="RNRI alpha",
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.05,
+                        value=0.1,
+                    )
+                    rnri_lr = gr.Slider(
+                        label="RNRI learning rate",
+                        minimum=0.0,
+                        maximum=1.0,
+                        step=0.05,
+                        value=0.2,
+                    )
+
             with gr.Row():
-                run_button = gr.Button("Edit", scale=1)
+                set_button = gr.Button("Set input image & description & settings", scale=1)
+
+                is_set_text = gr.Text("", show_label=False)
+
+                # Create a loading indicator
+                loading_indicator = gr.Markdown(value="⏳ Processing...", visible=False)
 
         with gr.Column(elem_id="col-container-2"):
             result = gr.Image(label="Result")
 
+            with gr.Row():
+                target_prompt = gr.Text(
+                    label="Edit prompt",
+                    info = "Enter your edit prompt",
+                    show_label=False,
+                    max_lines=1,
+                    placeholder="an oreo cake on a table",
+                    container=False,
+                )
+
+            with gr.Row():
+                run_button = gr.Button("Edit", scale=1)
+
             # gr.Examples(
             #     examples = examples,
             #     inputs = [prompt]
@@ -144,8 +175,8 @@ with gr.Blocks(css="style.css") as demo:
     set_button.click(
         fn=set_pipe,
         inputs=[input_image, description_prompt, edit_guidance_scale, num_inference_steps,
-                num_inference_steps, inversion_max_step],
-        outputs=editor_state
+                num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha, rnri_lr],
+        outputs=[editor_state, is_set_text],
     )
 
     run_button.click(

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, device):
+    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")
@@ -54,7 +54,7 @@ class ImageEditorDemo:
         self.cfg = cfg
         self.pipe_inversion.cfg = cfg
         self.pipe_inference.cfg = cfg
-        self.inv_hp = [2, 0.1, 0.2]
+        self.inv_hp = inv_hp  # [2, 0.1, 0.2]
         self.edit_cfg = cfg.edit_guidance_scale
 
         self.pipe_inference.to(device)