update

gradio_demo/app.py

@@ -63,41 +63,6 @@ else:
     device = "cpu"
     torch_dtype = torch.float32
 
-# Load pretrained models.
-print("Initializing pipeline...")
-pipe = InstantIRPipeline.from_pretrained(
-    sdxl_repo_id,
-    torch_dtype=torch_dtype,
-)
-
-# Image prompt projector.
-print("Loading LQ-Adapter...")
-load_adapter_to_pipe(
-    pipe,
-    f"{instantir_path}/adapter.pt",
-    dinov2_repo_id,
-)
-
-# Prepare previewer
-lora_alpha = pipe.prepare_previewers(instantir_path)
-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)
-pipe.scheduler = DDPMScheduler.from_pretrained(sdxl_repo_id, subfolder="scheduler")
-lcm_scheduler = LCMSingleStepScheduler.from_config(pipe.scheduler.config)
-
-# Load weights.
-print("Loading checkpoint...")
-aggregator_state_dict = torch.load(
-    f"{instantir_path}/aggregator.pt",
-    map_location="cpu"
-)
-pipe.aggregator.load_state_dict(aggregator_state_dict, strict=True)
-pipe.aggregator.to(device=device, dtype=torch_dtype)
-
-print("******loaded")
-
 MAX_SEED = np.iinfo(np.int32).max
 MAX_IMAGE_SIZE = 1024
 
@@ -129,6 +94,44 @@ def show_final_preview(preview_row):
 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)):
+
+
+
+    # Load pretrained models.
+    print("Initializing pipeline...")
+    pipe = InstantIRPipeline.from_pretrained(
+        sdxl_repo_id,
+        torch_dtype=torch_dtype,
+    )
+
+    # Image prompt projector.
+    print("Loading LQ-Adapter...")
+    load_adapter_to_pipe(
+        pipe,
+        f"{instantir_path}/adapter.pt",
+        dinov2_repo_id,
+    )
+
+    # Prepare previewer
+    lora_alpha = pipe.prepare_previewers(instantir_path)
+    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)
+    pipe.scheduler = DDPMScheduler.from_pretrained(sdxl_repo_id, subfolder="scheduler")
+    lcm_scheduler = LCMSingleStepScheduler.from_config(pipe.scheduler.config)
+
+    # Load weights.
+    print("Loading checkpoint...")
+    aggregator_state_dict = torch.load(
+        f"{instantir_path}/aggregator.pt",
+        map_location="cpu"
+    )
+    pipe.aggregator.load_state_dict(aggregator_state_dict, strict=True)
+    pipe.aggregator.to(device=device, dtype=torch_dtype)
+
+    print("******loaded")
+
     if creative_restoration:
         if "lcm" not in pipe.unet.active_adapters():
             pipe.unet.set_adapter('lcm')
@@ -177,6 +180,7 @@ def instantir_restore(
     for i, preview_img in enumerate(out[1]):
         preview_img.append(f"preview_{i}")
 
+    del pipe
     gc.collect()
     print(f"TORCH={torch}")
     if torch.cuda.is_available():