Update app.py

app.py

@@ -10,6 +10,7 @@ import torch
 import gradio as gr
 from PIL import Image
 from huggingface_hub import hf_hub_download
+from comfy import model_management
 
 def hf_hub_download_local(repo_id, filename, local_dir, **kwargs):
     downloaded_path = hf_hub_download(repo_id=repo_id, filename=filename, **kwargs)
@@ -199,6 +200,17 @@ MODELS_AND_NODES["VAEDecode"] = NODE_CLASS_MAPPINGS["VAEDecode"]()
 MODELS_AND_NODES["CreateVideo"] = NODE_CLASS_MAPPINGS["CreateVideo"]()
 MODELS_AND_NODES["SaveVideo"] = NODE_CLASS_MAPPINGS["SaveVideo"]()
 
+print("Pre-loading main models onto GPU...")
+model_loaders = [
+    MODELS_AND_NODES["clip"],
+    MODELS_AND_NODES["vae"],
+    MODELS_AND_NODES["model_low_noise"],  # This is the UNET + LoRA
+    MODELS_AND_NODES["model_high_noise"], # This is the other UNET + LoRA
+    MODELS_AND_NODES["clip_vision"],
+]
+model_management.load_models_gpu([
+    loader[0].patcher if hasattr(loader[0], 'patcher') else loader[0] for loader in model_loaders
+])
 print("All models loaded successfully!")
 
 # --- Main Video Generation Logic ---