update

gradio_demo/app.py

@@ -117,8 +117,7 @@ def instantir_restore(
     print(f"use lora alpha {lora_alpha}")
     lora_alpha = pipe.prepare_previewers(lcm_repo_id, use_lcm=True)
     print(f"use lora alpha {lora_alpha}")
-    if not cpu_offload:
-        pipe.to(device=device, dtype=torch_dtype)
+    pipe.to(device=device, dtype=torch_dtype)
     pipe.scheduler = DDPMScheduler.from_pretrained(sdxl_repo_id, subfolder="scheduler")
     lcm_scheduler = LCMSingleStepScheduler.from_config(pipe.scheduler.config)
 
@@ -126,15 +125,10 @@ def instantir_restore(
     print("Loading checkpoint...")
     aggregator_state_dict = torch.load(
         f"{instantir_path}/aggregator.pt",
-#        map_location = device
-#        map_location = device if not cpu_offload else "cpu"
         map_location="cpu"
     )
     pipe.aggregator.load_state_dict(aggregator_state_dict, strict=True)
-    if cpu_offload:
-        pipe.aggregator.to(device="cpu")
-    else:
-        pipe.aggregator.to(device=device, dtype=torch_dtype)
+    pipe.aggregator.to(device=device, dtype=torch_dtype)
 
     print("******loaded")
 
@@ -145,12 +139,8 @@ def instantir_restore(
         if "previewer" not in pipe.unet.active_adapters():
             pipe.unet.set_adapter('previewer')
 
-    print('optimizing')
-#    pipe.enable_vae_tiling()
-    if cpu_offload:
-        pipe.enable_model_cpu_offload()
-#        pipe.enable_sequential_cpu_offload()
-    print('done')
+#    if cpu_offload:
+#        pipe.enable_model_cpu_offload()
 
 
     if isinstance(guidance_end, int):
@@ -168,10 +158,6 @@ def instantir_restore(
     else:
         lq = [resize_img(lq.convert("RGB"), size=None)]
    
-    #if cpu_offload:
-    #    generator = torch.Generator().manual_seed(seed)
-    #else:
-    #    generator = torch.Generator(device=device).manual_seed(seed)
     generator = torch.Generator(device=device).manual_seed(seed)
     timesteps = [
         i * (1000//steps) + pipe.scheduler.config.steps_offset for i in range(0, steps)
@@ -266,7 +252,7 @@ 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)
+            cpu_offload = gr.Checkbox(label="CPU offload", info="If you have a lot of GPU VRAM, uncheck this option for faster generation", value=False, visible=False)
 
     
             with gr.Row():