Upload utils/model_loader.py with huggingface_hub

utils/model_loader.py

@@ -104,54 +104,58 @@ class ModelManager:
             )
         return self.model_paths["wav2vec"]
 
-    def load_wan_model(self, size="infinitetalk-480", device="cuda"):
+    def load_wan_model(self, size="infinitetalk-480", device="cuda", offload_model=True):
         """
-        Load Wan model for inference
+        Load Wan InfiniteTalk pipeline for inference
 
         Args:
-            size: Model size configuration
+            size: Model size configuration (infinitetalk-480 or infinitetalk-720)
             device: Device to load model on
+            offload_model: Whether to offload model to CPU between forwards
 
         Returns:
-            Loaded model
+            Loaded InfiniteTalkPipeline
         """
-        if "wan_model" not in self.models:
+        if "wan_pipeline" not in self.models:
             import wan
-            from wan.configs import SIZE_CONFIGS, WAN_CONFIGS
+            from wan.configs import WAN_CONFIGS
 
             model_path = self.get_wan_model_path()
             infinitetalk_path = self.get_infinitetalk_weights_path()
+            infinitetalk_weights = os.path.join(infinitetalk_path, "infinitetalk.safetensors")
 
-            logger.info(f"Loading Wan model from {model_path}...")
+            logger.info(f"Loading InfiniteTalk pipeline from {model_path}...")
 
-            # Initialize model based on InfiniteTalk's approach
+            # Get configuration for infinitetalk-14B
             task = "infinitetalk-14B"
-            args_dict = {
-                "ckpt_dir": model_path,
-                "infinitetalk_dir": os.path.join(infinitetalk_path, "infinitetalk.safetensors"),
-                "task": task,
-                "size": size,
-                "sample_steps": 40,
-                "sample_shift": 7 if size == "infinitetalk-480" else 11,
-            }
-
-            # Create a simple namespace object for args
-            class Args:
-                def __init__(self, **kwargs):
-                    self.__dict__.update(kwargs)
-
-            args = Args(**args_dict)
+            cfg = WAN_CONFIGS[task]
+
+            # Create InfiniteTalk pipeline
+            # This matches the initialization in generate_infinitetalk.py
+            pipeline = wan.InfiniteTalkPipeline(
+                config=cfg,
+                checkpoint_dir=model_path,
+                quant_dir=None,  # No quantization for now
+                device_id=device if isinstance(device, int) else 0,
+                rank=0,  # Single GPU
+                t5_fsdp=False,
+                dit_fsdp=False,
+                use_usp=False,
+                t5_cpu=False,
+                lora_dir=None,
+                lora_scales=None,
+                quant=None,
+                dit_path=None,
+                infinitetalk_dir=infinitetalk_weights
+            )
 
-            # Load model (simplified - actual loading would use wan.load_model())
-            # This is a placeholder - actual implementation would call the wan library
-            model = wan.WanModel(args)
-            model.to(device)
-            model.eval()
+            # Enable memory management for low VRAM if needed
+            # pipeline.enable_vram_management(num_persistent_param_in_dit=0)
 
-            self.models["wan_model"] = model
-            logger.info("Wan model loaded successfully")
+            self.models["wan_pipeline"] = pipeline
+            logger.info("InfiniteTalk pipeline loaded successfully")
 
-        return self.models["wan_model"]
+        return self.models["wan_pipeline"]
 
     def load_audio_encoder(self, device="cuda"):
         """