[Admin maintenance] Migrate grant to ZeroGPU

app.py

@@ -1,7 +1,7 @@
+import spaces
 import gradio as gr
 
 import torch
-import spaces
 from src.euler_scheduler import MyEulerAncestralDiscreteScheduler
 from diffusers.pipelines.auto_pipeline import AutoPipelineForImage2Image
 from src.sdxl_inversion_pipeline import SDXLDDIMPipeline
@@ -10,98 +10,53 @@ from src.editor import ImageEditorDemo
 
 device = "cuda" if torch.cuda.is_available() else "cpu"
 
-# if torch.cuda.is_available():
-#     torch.cuda.max_memory_allocated(device=device)
-#     pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", torch_dtype=torch.float16, variant="fp16", use_safetensors=True)
-#     pipe.enable_xformers_memory_efficient_attention()
-#     pipe = pipe.to(device)
-# else: 
-#     pipe = DiffusionPipeline.from_pretrained("stabilityai/sdxl-turbo", use_safetensors=True)
-#     pipe = pipe.to(device)
-
-# css = """
-# #col-container-1 {
-#     margin: 0 auto;
-#     max-width: 520px;
-# }
-# #col-container-2 {
-#     margin: 0 auto;
-#     max-width: 520px;
-# }
-# """
-
 if device == "cuda":
     torch.cuda.max_memory_allocated(device=device)
 
 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_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)
 
-if device == "cuda":
-    pipe_inference.enable_xformers_memory_efficient_attention()
-    pipe_inversion.enable_xformers_memory_efficient_attention()
-
 
-
-# with gr.Blocks(css=css) as demo:
-# with gr.Blocks(css="style.css") as demo:
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
     gr.Markdown(f""" # Real Time Editing with GNRI Inversion 🍎⚡️
-    This is a demo for our [paper](https://arxiv.org/abs/2312.12540) **GNRI: Lightning-fast Image Inversion and Editing 
+    This is a demo for our [paper](https://arxiv.org/abs/2312.12540) **GNRI: Lightning-fast Image Inversion and Editing
     for Text-to-Image Diffusion Models. Accepted to ICLR 2025**.
     More details can be found in the [project page](https://barakmam.github.io/rnri.github.io/).
-    
+
     The demo is based on SDXL. Improved results can be achieved with our FLUX version (examples in project page).
     """)
-    inv_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):
-
-        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)
+    def invert_and_edit(input_image, description_prompt, target_prompt,
+                        edit_guidance_scale, num_inference_steps=4,
+                        inversion_max_step=0.6, rnri_iterations=2, rnri_alpha=0.1, rnri_lr=0.2,
+                        progress=gr.Progress(track_tqdm=True)):
+        if input_image is None or not description_prompt or not target_prompt:
+            raise gr.Error("Please provide an input image, description, and edit prompt.")
+
+        num_inference_steps = int(num_inference_steps)
+        edit_guidance_scale = float(edit_guidance_scale)
+        inversion_max_step = float(inversion_max_step)
+        rnri_iterations = int(rnri_iterations)
+        rnri_alpha = float(rnri_alpha)
+        rnri_lr = float(rnri_lr)
 
         config = RunConfig(num_inference_steps=num_inference_steps,
-                           num_inversion_steps=num_inversion_steps,
+                           num_inversion_steps=num_inference_steps,
                            edit_guidance_scale=edit_guidance_scale,
                            inversion_max_step=inversion_max_step)
-        if device == 'cuda':
-            pipe_inference.to('cpu')
-            torch.cuda.empty_cache()
 
-        inversion_state = ImageEditorDemo.invert(pipe_inversion.to(device), input_image, description_prompt, config,
+        inversion_state = ImageEditorDemo.invert(pipe_inversion, input_image, description_prompt, config,
                                                  [rnri_iterations, rnri_alpha, rnri_lr], device)
-        if device == 'cuda':
-            pipe_inversion.to('cpu')
-            torch.cuda.empty_cache()
-            pipe_inference.to(device)
-
-        gr.Info('Input has set!')
-        return inversion_state, "Input has set!"
-
-    @spaces.GPU
-    def edit(inversion_state, target_prompt):
-        if inversion_state is None:
-            raise gr.Error("Set inputs before editing. Progress indication below")
 
-        image = ImageEditorDemo.edit(pipe_inference, target_prompt, inversion_state['latent'], inversion_state['noise'],
+        image = ImageEditorDemo.edit(pipe_inference, target_prompt,
+                                     inversion_state['latent'], inversion_state['noise'],
                                      inversion_state['cfg'], inversion_state['cfg'].edit_guidance_scale)
-
         return image
 
 
@@ -167,9 +122,6 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
                         value=0.2,
                     )
 
-            with gr.Row():
-                is_set_text = gr.Text("", show_label=False)
-
         with gr.Column(elem_id="col-container-2"):
             result = gr.Image(label="Result")
 
@@ -198,57 +150,11 @@ with gr.Blocks(theme=gr.themes.Soft()) as demo:
      Errors may appear due to resource availability by HF.
     Performance may be inferior to the reported in the paper due to hardware limitation.""")
 
-    input_image.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=[inv_state, is_set_text], trigger_mode='once')
-
-    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=[inv_state, is_set_text], trigger_mode='once')
-
-    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=[inv_state, is_set_text], trigger_mode='once')
-    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=[inv_state, is_set_text], trigger_mode='once')
-    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=[inv_state, is_set_text], trigger_mode='once')
-    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=[inv_state, is_set_text], trigger_mode='once')
-    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=[inv_state, is_set_text], trigger_mode='once')
-    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=[inv_state, is_set_text], trigger_mode='once')
-
-    # set_button.click(
-    #     fn=set_pipe,
-    #     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=[inv_state, is_set_text],
-    # )
-
     run_button.click(
-        fn=edit,
-        inputs=[inv_state, target_prompt],
+        fn=invert_and_edit,
+        inputs=[input_image, description_prompt, target_prompt,
+                edit_guidance_scale, num_inference_steps,
+                inversion_max_step, rnri_iterations, rnri_alpha, rnri_lr],
         outputs=[result]
     )
 

requirements.txt

@@ -1,8 +1,5 @@
 accelerate==0.25.0
 diffusers==0.29.0
 invisible_watermark
-torch==2.8.0
 transformers==4.32.1
-xformers
-torchvision==0.23.0
-pyrallis==0.3.1
+pyrallis==0.3.1