Spaces:
Running
Running
File size: 6,875 Bytes
38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 2c73ba8 38572a2 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 |
"""
Model Manager for InfiniteTalk
Handles lazy loading and caching of models from HuggingFace Hub
"""
import os
import torch
from huggingface_hub import snapshot_download
from pathlib import Path
import logging
logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)
class ModelManager:
"""Manages model loading and caching"""
def __init__(self, cache_dir=None):
"""
Initialize Model Manager
Args:
cache_dir: Directory for caching models. Defaults to HF_HOME or /data/.huggingface
"""
if cache_dir is None:
cache_dir = os.environ.get("HF_HOME", "/data/.huggingface")
self.cache_dir = Path(cache_dir)
self.cache_dir.mkdir(parents=True, exist_ok=True)
self.models = {}
self.model_paths = {
"wan": None,
"infinitetalk": None,
"wav2vec": None
}
def download_model(self, repo_id, subfolder=None, filename=None):
"""
Download model from HuggingFace Hub with caching
Args:
repo_id: HuggingFace repository ID (e.g., "Kijai/WanVideo_comfy")
subfolder: Optional subfolder within the repository
filename: Optional specific file to download
Returns:
Path to downloaded model directory
"""
try:
logger.info(f"Downloading {repo_id} from HuggingFace Hub...")
download_kwargs = {
"repo_id": repo_id,
"cache_dir": str(self.cache_dir),
"resume_download": True,
}
if subfolder:
download_kwargs["allow_patterns"] = f"{subfolder}/*"
if filename:
download_kwargs["allow_patterns"] = filename
model_path = snapshot_download(**download_kwargs)
if subfolder:
model_path = os.path.join(model_path, subfolder)
logger.info(f"Model downloaded successfully to {model_path}")
return model_path
except Exception as e:
logger.error(f"Error downloading model {repo_id}: {e}")
raise
def get_wan_model_path(self):
"""Get or download Wan2.1 I2V model"""
if self.model_paths["wan"] is None:
logger.info("Downloading Wan2.1-I2V-14B-480P model...")
# This will download the full model - adjust repo_id based on actual HF location
self.model_paths["wan"] = self.download_model(
repo_id="Kijai/WanVideo_comfy",
subfolder="wan2_1_i2v_14B_480P"
)
return self.model_paths["wan"]
def get_infinitetalk_weights_path(self):
"""Get or download InfiniteTalk weights"""
if self.model_paths["infinitetalk"] is None:
logger.info("Downloading InfiniteTalk weights...")
self.model_paths["infinitetalk"] = self.download_model(
repo_id="MeiGen-AI/InfiniteTalk",
subfolder="single"
)
return self.model_paths["infinitetalk"]
def get_wav2vec_model_path(self):
"""Get or download Wav2Vec2 audio encoder"""
if self.model_paths["wav2vec"] is None:
logger.info("Downloading Wav2Vec2 audio encoder...")
self.model_paths["wav2vec"] = self.download_model(
repo_id="TencentGameMate/chinese-wav2vec2-base"
)
return self.model_paths["wav2vec"]
def load_wan_model(self, size="infinitetalk-480", device="cuda", offload_model=True):
"""
Load Wan InfiniteTalk pipeline for inference
Args:
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 InfiniteTalkPipeline
"""
if "wan_pipeline" not in self.models:
import wan
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 InfiniteTalk pipeline from {model_path}...")
# Get configuration for infinitetalk-14B
task = "infinitetalk-14B"
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
)
# Enable memory management for low VRAM if needed
# pipeline.enable_vram_management(num_persistent_param_in_dit=0)
self.models["wan_pipeline"] = pipeline
logger.info("InfiniteTalk pipeline loaded successfully")
return self.models["wan_pipeline"]
def load_audio_encoder(self, device="cuda"):
"""
Load Wav2Vec2 audio encoder
Args:
device: Device to load model on
Returns:
Audio encoder model and feature extractor
"""
if "audio_encoder" not in self.models:
from transformers import Wav2Vec2FeatureExtractor
from src.audio_analysis.wav2vec2 import Wav2Vec2Model
wav2vec_path = self.get_wav2vec_model_path()
logger.info(f"Loading audio encoder from {wav2vec_path}...")
feature_extractor = Wav2Vec2FeatureExtractor.from_pretrained(wav2vec_path)
audio_encoder = Wav2Vec2Model.from_pretrained(wav2vec_path)
audio_encoder.to(device)
audio_encoder.eval()
self.models["audio_encoder"] = (audio_encoder, feature_extractor)
logger.info("Audio encoder loaded successfully")
return self.models["audio_encoder"]
def unload_model(self, model_name):
"""Unload a specific model to free memory"""
if model_name in self.models:
del self.models[model_name]
torch.cuda.empty_cache()
logger.info(f"Unloaded {model_name}")
def clear_all(self):
"""Unload all models"""
self.models.clear()
torch.cuda.empty_cache()
logger.info("All models unloaded")
|