Update infer.py

infer.py

@@ -122,16 +122,30 @@ def parse_args():
     args = parser.parse_args()
     return args
 
+import torch
+from huggingface_hub import snapshot_download
+from diffusers import WanVideoPipeline
+from transformers import Wav2Vec2Processor, Wav2Vec2Model
+from models import FantasyTalkingAudioConditionModel  # adjust import if needed
+from model_manager import ModelManager  # assuming this exists in your repo
+
+
 def load_models(args):
-    # Load Wan I2V models
+    print("🚀 [Startup] Initializing all models (compile-time preloading)...")
 
+    # --------------------------------------------
+    # STEP 1 — Ensure all model files are cached
+    # --------------------------------------------
     snapshot_download("Wan-AI/Wan2.1-I2V-14B-720P", local_dir="./models/Wan2.1-I2V-14B-720P")
     snapshot_download("facebook/wav2vec2-base-960h", local_dir="./models/wav2vec2-base-960h")
     snapshot_download("acvlab/FantasyTalking", local_dir="./models")
 
+    # --------------------------------------------
+    # STEP 2 — Initialize ModelManager (core loader)
+    # --------------------------------------------
+    print("🔧 Loading Wan I2V model checkpoints via ModelManager...")
+    model_manager = ModelManager(device="cuda")
 
-    model_manager = ModelManager(device="cpu")
-    
     model_manager.load_models(
         [
             [
@@ -147,23 +161,56 @@ def load_models(args):
             f"{args.wan_model_dir}/models_t5_umt5-xxl-enc-bf16.pth",
             f"{args.wan_model_dir}/Wan2.1_VAE.pth",
         ],
-        # torch_dtype=torch.float8_e4m3fn, # You can set `torch_dtype=torch.bfloat16` to disable FP8 quantization.
-        torch_dtype=torch.bfloat16, # You can set `torch_dtype=torch.bfloat16` to disable FP8 quantization.
+        torch_dtype=torch.bfloat16,
+    )
+
+    pipe = WanVideoPipeline.from_model_manager(
+        model_manager, torch_dtype=torch.bfloat16, device="cuda"
     )
-    pipe = WanVideoPipeline.from_model_manager(model_manager, torch_dtype=torch.bfloat16, device="cuda")
 
-    # Load FantasyTalking weights
-    fantasytalking = FantasyTalkingAudioConditionModel(pipe.dit, 768, 2048).to("cuda")
-    fantasytalking.load_audio_processor(args.fantasytalking_model_path, pipe.dit)
-    
-    # You can set `num_persistent_param_in_dit` to a small number to reduce VRAM required.
     pipe.enable_vram_management(num_persistent_param_in_dit=args.num_persistent_param_in_dit)
+    pipe.to("cuda")
+    pipe.eval()
+
+    # --------------------------------------------
+    # STEP 3 — Load FantasyTalking model
+    # --------------------------------------------
+    print("🧠 Loading FantasyTalking model...")
+    fantasytalking = FantasyTalkingAudioConditionModel(pipe.dit, 768, 2048).to("cuda").eval()
+    fantasytalking.load_audio_processor(args.fantasytalking_model_path, pipe.dit)
 
-    # Load wav2vec models
+    # --------------------------------------------
+    # STEP 4 — Load Wav2Vec2 model + processor
+    # --------------------------------------------
+    print("🎙️ Loading Wav2Vec2 model...")
     wav2vec_processor = Wav2Vec2Processor.from_pretrained(args.wav2vec_model_dir)
-    wav2vec = Wav2Vec2Model.from_pretrained(args.wav2vec_model_dir).to("cuda")
+    wav2vec = Wav2Vec2Model.from_pretrained(args.wav2vec_model_dir).to("cuda").eval()
+
+    # --------------------------------------------
+    # STEP 5 — FORCE preload (compile-time warmup)
+    # --------------------------------------------
+    print("🔥 Preloading all models into GPU memory (forcing weight instantiation)...")
+    with torch.no_grad(), torch.autocast("cuda", dtype=torch.bfloat16):
+        # Wav2Vec2 warmup
+        dummy_audio = torch.zeros(1, 16000).to("cuda")
+        _ = wav2vec(dummy_audio)
+
+        # Diffusion UNet warmup
+        dummy_latent = torch.randn(1, pipe.unet.in_channels, 64, 64, device="cuda", dtype=torch.bfloat16)
+        _ = pipe.unet(dummy_latent, 0.5)
+
+        # FantasyTalking warmup
+        try:
+            dummy_feat = torch.randn(1, 256).to("cuda")
+            _ = fantasytalking(dummy_feat)
+        except Exception as e:
+            print(f"⚠️ FantasyTalking warmup skipped: {e}")
+
+    torch.cuda.synchronize()
+    print("✅ [Ready] All models fully loaded and warmed up in GPU memory.")
+
+    return pipe, fantasytalking, wav2vec_processor, wav2vec
 
-    return pipe,fantasytalking,wav2vec_processor,wav2vec