""" PRODUCTION MODEL DOWNLOADER for OmniAvatar Video Generation This script MUST download the actual models for video generation to work """ import os import subprocess import sys import logging import time from pathlib import Path import requests from urllib.parse import urljoin logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s') logger = logging.getLogger(__name__) class OmniAvatarModelDownloader: """Production-grade model downloader for OmniAvatar video generation""" def __init__(self): self.base_dir = Path.cwd() self.models_dir = self.base_dir / "pretrained_models" # CRITICAL: These models are REQUIRED for video generation self.required_models = { "Wan2.1-T2V-14B": { "repo": "Wan-AI/Wan2.1-T2V-14B", "description": "Base text-to-video generation model", "size": "~28GB", "priority": 1, "essential": True }, "OmniAvatar-14B": { "repo": "OmniAvatar/OmniAvatar-14B", "description": "Avatar LoRA weights and animation model", "size": "~2GB", "priority": 2, "essential": True }, "wav2vec2-base-960h": { "repo": "facebook/wav2vec2-base-960h", "description": "Audio encoder for lip-sync", "size": "~360MB", "priority": 3, "essential": True } } def install_huggingface_cli(self): """Install HuggingFace CLI for model downloads""" logger.info("📦 Installing HuggingFace CLI...") try: subprocess.run([sys.executable, "-m", "pip", "install", "huggingface_hub[cli]"], check=True, capture_output=True) logger.info("✅ HuggingFace CLI installed") return True except subprocess.CalledProcessError as e: logger.error(f"❌ Failed to install HuggingFace CLI: {e}") return False def check_huggingface_cli(self): """Check if HuggingFace CLI is available""" try: result = subprocess.run(["huggingface-cli", "--version"], capture_output=True, text=True) if result.returncode == 0: logger.info("✅ HuggingFace CLI available") return True except FileNotFoundError: pass logger.info("❌ HuggingFace CLI not found, installing...") return self.install_huggingface_cli() def create_model_directories(self): """Create directory structure for models""" logger.info("📁 Creating model directories...") for model_name in self.required_models.keys(): model_dir = self.models_dir / model_name model_dir.mkdir(parents=True, exist_ok=True) logger.info(f"✅ Created: {model_dir}") def download_model_with_cli(self, model_name: str, model_info: dict) -> bool: """Download model using HuggingFace CLI""" local_dir = self.models_dir / model_name # Skip if already downloaded if local_dir.exists() and any(local_dir.iterdir()): logger.info(f"✅ {model_name} already exists, skipping...") return True logger.info(f"📥 Downloading {model_name} ({model_info['size']})...") logger.info(f"📝 {model_info['description']}") cmd = [ "huggingface-cli", "download", model_info["repo"], "--local-dir", str(local_dir), "--local-dir-use-symlinks", "False" ] try: logger.info(f"🚀 Running: {' '.join(cmd)}") result = subprocess.run(cmd, check=True, capture_output=True, text=True) logger.info(f"✅ {model_name} downloaded successfully!") return True except subprocess.CalledProcessError as e: logger.error(f"❌ Failed to download {model_name}: {e.stderr}") return False def download_model_with_git(self, model_name: str, model_info: dict) -> bool: """Fallback: Download model using git clone""" local_dir = self.models_dir / model_name if local_dir.exists() and any(local_dir.iterdir()): logger.info(f"✅ {model_name} already exists, skipping...") return True logger.info(f"📥 Downloading {model_name} with git clone...") # Remove directory if it exists but is empty if local_dir.exists(): local_dir.rmdir() cmd = ["git", "clone", f"https://huggingface.co/{model_info['repo']}", str(local_dir)] try: result = subprocess.run(cmd, check=True, capture_output=True, text=True) logger.info(f"✅ {model_name} downloaded with git!") return True except subprocess.CalledProcessError as e: logger.error(f"❌ Git clone failed for {model_name}: {e.stderr}") return False def verify_downloads(self) -> bool: """Verify all required models are downloaded""" logger.info("🔍 Verifying model downloads...") all_present = True for model_name in self.required_models.keys(): model_dir = self.models_dir / model_name if model_dir.exists() and any(model_dir.iterdir()): file_count = len(list(model_dir.rglob("*"))) logger.info(f"✅ {model_name}: {file_count} files found") else: logger.error(f"❌ {model_name}: Missing or empty") all_present = False return all_present def download_all_models(self) -> bool: """Download all required models for video generation""" logger.info("🎬 DOWNLOADING OMNIAVATAR MODELS FOR VIDEO GENERATION") logger.info("=" * 60) logger.info("⚠️ This will download approximately 30GB of models") logger.info("🎯 These models are REQUIRED for avatar video generation") logger.info("") # Check prerequisites if not self.check_huggingface_cli(): logger.error("❌ Cannot proceed without HuggingFace CLI") return False # Create directories self.create_model_directories() # Download each model success_count = 0 for model_name, model_info in self.required_models.items(): logger.info(f"\n📦 Processing {model_name} (Priority {model_info['priority']})...") # Try HuggingFace CLI first success = self.download_model_with_cli(model_name, model_info) # Fallback to git if CLI fails if not success: logger.info("🔄 Trying git clone fallback...") success = self.download_model_with_git(model_name, model_info) if success: success_count += 1 logger.info(f"✅ {model_name} download completed") else: logger.error(f"❌ {model_name} download failed") if model_info["essential"]: logger.error("🚨 This model is ESSENTIAL for video generation!") # Verify all downloads if self.verify_downloads(): logger.info("\n🎉 ALL OMNIAVATAR MODELS DOWNLOADED SUCCESSFULLY!") logger.info("🎬 Avatar video generation is now FULLY ENABLED!") logger.info("💡 Restart your application to activate video generation") return True else: logger.error("\n❌ Model download incomplete") logger.error("🎯 Video generation will not work without all required models") return False def main(): """Main function to download OmniAvatar models""" downloader = OmniAvatarModelDownloader() try: success = downloader.download_all_models() if success: print("\n🎬 OMNIAVATAR VIDEO GENERATION READY!") print("✅ All models downloaded successfully") print("🚀 Your app can now generate avatar videos!") return 0 else: print("\n❌ MODEL DOWNLOAD FAILED") print("🎯 Video generation will not work") print("💡 Please check the error messages above") return 1 except KeyboardInterrupt: print("\n⏹️ Download cancelled by user") return 1 except Exception as e: print(f"\n💥 Unexpected error: {e}") return 1 if __name__ == "__main__": sys.exit(main())