preload

app.py

@@ -102,6 +102,48 @@ def load_video_frames(video_path: str, source_frames: int):
 
     return input_video, original_height, original_width
 
+
+def preload_models(controller, default_model_path, default_lora_name, acc_lora_path):
+    """
+    Preload base model and LoRAs before launching the app to avoid first-run latency.
+    """
+    controller.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+    # Ensure tracking flags exist
+    if not hasattr(controller, "_active_lora_path"):
+        controller._active_lora_path = None
+    if not hasattr(controller, "_acc_lora_active"):
+        controller._acc_lora_active = False
+
+    try:
+        print(f"[preload] Loading base model: {default_model_path}")
+        controller.update_diffusion_transformer(default_model_path)
+        controller.update_base_model(default_model_path)
+
+        print(f"[preload] Loading VideoCoF LoRA: {default_lora_name}")
+        controller.update_lora_model(default_lora_name)
+        if controller.lora_model_path and controller.lora_model_path != "none":
+            controller.pipeline = merge_lora(
+                controller.pipeline,
+                controller.lora_model_path,
+                multiplier=1.0,
+                device=controller.device,
+            )
+            controller._active_lora_path = controller.lora_model_path
+
+        if acc_lora_path and os.path.exists(acc_lora_path):
+            print(f"[preload] Loading Acceleration LoRA: {acc_lora_path}")
+            controller.pipeline = merge_lora(
+                controller.pipeline, acc_lora_path, multiplier=1.0, device=controller.device
+            )
+            controller._acc_lora_active = True
+        else:
+            print(f"[preload] Acceleration LoRA not found at {acc_lora_path}")
+    except Exception as e:
+        print(f"[preload] Warning: preload failed: {e}")
+    finally:
+        torch.cuda.empty_cache()
+
 class VideoCoF_Controller(Wan_Controller):
     @spaces.GPU(duration=300)
     @timer
@@ -317,24 +359,33 @@ def ui(GPU_memory_mode, scheduler_dict, config_path, compile_dit, weight_dtype):
             try:
                 from huggingface_hub import snapshot_download, hf_hub_download
                 print("Downloading Wan2.1-T2V-14B weights...")
-                snapshot_download(repo_id="Wan-AI/Wan2.1-T2V-14B", local_dir="Wan-AI/Wan2.1-T2V-14B")
+                default_model_path = "Wan-AI/Wan2.1-T2V-14B"
+                snapshot_download(repo_id=default_model_path, local_dir=default_model_path)
                 
                 os.makedirs("models/Personalized_Model", exist_ok=True)
                 
                 print("Downloading VideoCoF weights...")
-                hf_hub_download(repo_id="XiangpengYang/VideoCoF", filename="videocof.safetensors", local_dir="models/Personalized_Model")
+                default_lora_name = "videocof.safetensors"
+                hf_hub_download(repo_id="XiangpengYang/VideoCoF", filename=default_lora_name, local_dir="models/Personalized_Model")
                 
                 print("Downloading FusionX Acceleration LoRA...")
-                hf_hub_download(repo_id="MonsterMMORPG/Wan_GGUF", filename="Wan2.1_Text_to_Video_14B_FusionX_LoRA.safetensors", local_dir="models/Personalized_Model")
+                acc_lora_filename = "Wan2.1_Text_to_Video_14B_FusionX_LoRA.safetensors"
+                hf_hub_download(repo_id="MonsterMMORPG/Wan_GGUF", filename=acc_lora_filename, local_dir="models/Personalized_Model")
                 
             except Exception as e:
                 print(f"Warning: Failed to pre-download weights: {e}")
 
-            base_model_dropdown, lora_model_dropdown, lora_alpha_slider, _ = create_finetune_models_checkpoints(controller, visible=False, default_lora="videocof.safetensors")
+            base_model_dropdown, lora_model_dropdown, lora_alpha_slider, _ = create_finetune_models_checkpoints(
+                controller, visible=False, default_lora="videocof.safetensors"
+            )
             
             # Set default LoRA alpha to 1.0 (matching inference.py)
             lora_alpha_slider.value = 1.0
 
+        # Preload heavy weights and LoRAs before launching the UI to avoid first-run latency.
+        acc_lora_path = os.path.join("models", "Personalized_Model", "Wan2.1_Text_to_Video_14B_FusionX_LoRA.safetensors")
+        preload_models(controller, default_model_path, "videocof.safetensors", acc_lora_path)
+
         with gr.Column(variant="panel"):
             prompt_textbox, negative_prompt_textbox = create_prompts(prompt="Remove the young man with short black hair wearing black shirt on the left.")