update

gradio_demo/app.py

@@ -117,7 +117,8 @@ 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}")
-    pipe.to(device=device, dtype=torch_dtype)
+    if not cpu_offload:
+        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)
 
@@ -125,10 +126,11 @@ def instantir_restore(
     print("Loading checkpoint...")
     aggregator_state_dict = torch.load(
         f"{instantir_path}/aggregator.pt",
-#        map_location="cpu"
+        map_location="cpu"
     )
     pipe.aggregator.load_state_dict(aggregator_state_dict, strict=True)
-    pipe.aggregator.to(device=device, dtype=torch_dtype)
+    if not cpu_offload:
+        pipe.aggregator.to(device=device, dtype=torch_dtype)
 
     print("******loaded")