clean code

app.py

@@ -56,42 +56,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     Image editing using RNRI for inversion demonstrates significant speed-up and improved quality compared to previous state-of-the-art methods.
     Take a look at the [project page](https://barakmam.github.io/rnri.github.io/).
     """)
-    editor_state = gr.State()
-
-
-    # @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')
-
-            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)
-        image_editor = ImageEditorDemo(pipe_inversion, pipe_inference, input_image,
-                                       description_prompt, config, device,
-                                       [rnri_iterations, rnri_alpha, rnri_lr])
-        return image_editor, "Input has set!"
+    inv_state = gr.State()
 
 
     @spaces.GPU
-    def set_pipe(inversion_state, input_image, description_prompt, edit_guidance_scale, num_inference_steps=4,
+    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):
 
         if input_image is None or not description_prompt:
@@ -131,29 +100,10 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         gr.Info('Input has set!')
         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":
-        #       image = editor.to(device).edit(target_prompt)
-        #  else:
-        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. Progress indication below")
-        # if device == "cuda":
-        #       image = editor.to(device).edit(target_prompt)
-        #  else:
-
-        # if device == 'cuda':
-        #     torch.cuda.empty_cache()
 
         print(f"#### 5 #### pipe_inversion.device: {pipe_inversion.device}")
         print(f"#### 6 #### pipe_inference.device: {pipe_inference.device}")
@@ -161,11 +111,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
         image = ImageEditorDemo.edit(pipe_inference, target_prompt, inversion_state['latent'], inversion_state['noise'],
                                      inversion_state['cfg'], inversion_state['cfg'].edit_guidance_scale)
 
-        # if device == 'cuda':
-            # pipe_inference.to('cpu')
-            # torch.cuda.empty_cache()
-
-
         print(f"#### 7 #### pipe_inversion.device: {pipe_inversion.device}")
         print(f"#### 8 #### pipe_inference.device: {pipe_inference.device}")
         return image
@@ -262,57 +207,57 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
     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,
+                       inputs=[input_image, description_prompt, edit_guidance_scale, num_inference_steps,
                                num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha, rnri_lr],
-                       outputs=[editor_state, is_set_text])
+                       outputs=[inv_state, is_set_text])
 
-    description_prompt.change(set_pipe, inputs=[editor_state, input_image, description_prompt, edit_guidance_scale,
+    description_prompt.change(set_pipe, inputs=[input_image, description_prompt, edit_guidance_scale,
                                                 num_inference_steps,
                                                 num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha,
                                                 rnri_lr],
-                              outputs=[editor_state, is_set_text])
+                              outputs=[inv_state, is_set_text])
 
-    edit_guidance_scale.change(set_pipe, inputs=[editor_state, input_image, description_prompt, edit_guidance_scale,
+    edit_guidance_scale.change(set_pipe, inputs=[input_image, description_prompt, edit_guidance_scale,
                                                  num_inference_steps,
                                                  num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha,
                                                  rnri_lr],
-                               outputs=[editor_state, is_set_text])
-    num_inference_steps.change(set_pipe, inputs=[editor_state, input_image, description_prompt, edit_guidance_scale,
+                               outputs=[inv_state, is_set_text])
+    num_inference_steps.change(set_pipe, inputs=[input_image, description_prompt, edit_guidance_scale,
                                                  num_inference_steps,
                                                  num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha,
                                                  rnri_lr],
-                               outputs=[editor_state, is_set_text])
-    inversion_max_step.change(set_pipe, inputs=[editor_state, input_image, description_prompt, edit_guidance_scale,
+                               outputs=[inv_state, is_set_text])
+    inversion_max_step.change(set_pipe, inputs=[input_image, description_prompt, edit_guidance_scale,
                                                 num_inference_steps,
                                                 num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha,
                                                 rnri_lr],
-                              outputs=[editor_state, is_set_text])
-    rnri_iterations.change(set_pipe, inputs=[editor_state, input_image, description_prompt, edit_guidance_scale,
+                              outputs=[inv_state, is_set_text])
+    rnri_iterations.change(set_pipe, inputs=[input_image, description_prompt, edit_guidance_scale,
                                              num_inference_steps,
                                              num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha,
                                              rnri_lr],
-                           outputs=[editor_state, is_set_text])
-    rnri_alpha.change(set_pipe, inputs=[editor_state, input_image, description_prompt, edit_guidance_scale,
+                           outputs=[inv_state, is_set_text])
+    rnri_alpha.change(set_pipe, inputs=[input_image, description_prompt, edit_guidance_scale,
                                         num_inference_steps,
                                         num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha,
                                         rnri_lr],
-                      outputs=[editor_state, is_set_text])
-    rnri_lr.change(set_pipe, inputs=[editor_state, input_image, description_prompt, edit_guidance_scale,
+                      outputs=[inv_state, is_set_text])
+    rnri_lr.change(set_pipe, inputs=[input_image, description_prompt, edit_guidance_scale,
                                      num_inference_steps,
                                      num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha,
                                      rnri_lr],
-                   outputs=[editor_state, is_set_text])
+                   outputs=[inv_state, is_set_text])
 
     # set_button.click(
     #     fn=set_pipe,
-    #     inputs=[editor_state, input_image, description_prompt, edit_guidance_scale, num_inference_steps,
+    #     inputs=[inv_state, input_image, description_prompt, edit_guidance_scale, num_inference_steps,
     #             num_inference_steps, inversion_max_step, rnri_iterations, rnri_alpha, rnri_lr],
-    #     outputs=[editor_state, is_set_text],
+    #     outputs=[inv_state, is_set_text],
     # )
 
     run_button.click(
         fn=edit,
-        inputs=[editor_state, target_prompt],
+        inputs=[inv_state, target_prompt],
         outputs=[result]
     )