update

gradio_demo/app.py

@@ -93,7 +93,7 @@ def show_final_preview(preview_row):
 @torch.no_grad()
 def instantir_restore(
     lq, prompt="", steps=30, cfg_scale=7.0, guidance_end=1.0,
-    creative_restoration=False, seed=3407, height=1024, width=1024, preview_start=0.0, progress=gr.Progress(track_tqdm=True)):
+    creative_restoration=False, seed=3407, height=1024, width=1024, preview_start=0.0, cpu_offload, progress=gr.Progress(track_tqdm=True)):
 
 
 
@@ -139,6 +139,15 @@ def instantir_restore(
         if "previewer" not in pipe.unet.active_adapters():
             pipe.unet.set_adapter('previewer')
 
+    print('optimizing')
+    pipe.vae.enable_slicing()
+    pipe.vae.enable_tiling()
+    if cpu_offload:
+        pipe.enable_model_cpu_offload()
+        pipe.enable_sequential_cpu_offload()
+    print('done')
+
+
     if isinstance(guidance_end, int):
         guidance_end = guidance_end / steps
     elif guidance_end > 1.0:
@@ -248,6 +257,8 @@ with gr.Blocks() as demo:
             preview_start = gr.Slider(label="Preview Start", value=0, minimum=0, maximum=30, step=1)
             prompt = gr.Textbox(label="Restoration prompts (Optional)", placeholder="")
             mode = gr.Checkbox(label="Creative Restoration", value=False)
+            cpu_offload = gr.Checkbox(label="CPU offload", info="If you have a lot of GPU VRAM, uncheck this option for faster generation", value=False)
+
     
             with gr.Row():
                 restore_btn = gr.Button("InstantIR magic!")
@@ -266,7 +277,7 @@ with gr.Blocks() as demo:
     restore_btn.click(
         instantir_restore, inputs=[
             lq_img, prompt, steps, cfg_scale, guidance_end,
-            mode, seed, height, width, preview_start,
+            mode, seed, height, width, preview_start, cpu_offload
         ],
         outputs=[output, pipe_out], api_name="InstantIR"
     )