Upload 2 files

app.py

@@ -14,21 +14,6 @@ except Exception as e:
     print(f"GPU warmup failed: {e}")
 os.environ["GRADIO_TEMP_DIR"] = "./tmp"
 
-try:
-    # You can define minimal dummy args here for initialization
-    class DummyArgs:
-        wan_model_dir = "./models/Wan2.1-I2V-14B-720P"
-        fantasytalking_model_path = "./models/fantasytalking_model.ckpt"
-        wav2vec_model_dir = "./models/wav2vec2-base-960h"
-
-    args = DummyArgs()
-    pipe, fantasytalking, wav2vec_processor, wav2vec = load_models(args)
-    print("✅ Models loaded successfully.")
-except Exception as e:
-    print(f"❌ Error loading models: {e}")
-    pipe = fantasytalking = wav2vec_processor = wav2vec = None
-    raise e  # fail fast on startup if models can't load
-
 pipe,fantasytalking,wav2vec_processor,wav2vec = None,None,None,None
 @spaces.GPU(duration=1200)
 def generate_video(
@@ -68,7 +53,17 @@ def generate_video(
         seed=seed,
     )
 
-    
+    try:
+        global pipe, fantasytalking, wav2vec_processor, wav2vec  
+        if pipe is None:
+            pipe,fantasytalking,wav2vec_processor,wav2vec = load_models(args)
+        output_path=main(
+            args,pipe,fantasytalking,wav2vec_processor,wav2vec
+        )
+        return output_path  # Ensure the output path is returned
+    except Exception as e:
+        print(f"Error during processing: {str(e)}")
+        raise gr.Error(f"Error during processing: {str(e)}")
 
 
 def create_args(

infer.py

@@ -122,30 +122,16 @@ 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):
-    print("🚀 [Startup] Initializing all models (compile-time preloading)...")
+    # Load Wan I2V models
 
-    # --------------------------------------------
-    # 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(
         [
             [
@@ -161,56 +147,23 @@ 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.bfloat16,
-    )
-
-    pipe = WanVideoPipeline.from_model_manager(
-        model_manager, torch_dtype=torch.bfloat16, device="cuda"
+        # 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.
     )
+    pipe = WanVideoPipeline.from_model_manager(model_manager, torch_dtype=torch.bfloat16, device="cuda")
 
-    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()
+    # 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)
 
-    # --------------------------------------------
-    # STEP 4 — Load Wav2Vec2 model + processor
-    # --------------------------------------------
-    print("🎙️ Loading Wav2Vec2 model...")
+    # Load wav2vec models
     wav2vec_processor = Wav2Vec2Processor.from_pretrained(args.wav2vec_model_dir)
-    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
+    wav2vec = Wav2Vec2Model.from_pretrained(args.wav2vec_model_dir).to("cuda")
 
+    return pipe,fantasytalking,wav2vec_processor,wav2vec