download_models_production.py

"""
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("SUCCESS: HuggingFace CLI installed")
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"ERROR: 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("SUCCESS: HuggingFace CLI available")
                return True
        except FileNotFoundError:
            pass
        
        logger.info("ERROR: 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"SUCCESS: 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"SUCCESS: {model_name} already exists, skipping...")
            return True
        
        logger.info(f"📥 Downloading {model_name} ({model_info['size']})...")
        logger.info(f"[INFO] {model_info['description']}")
        
        cmd = [
            "huggingface-cli", "download",
            model_info["repo"],
            "--local-dir", str(local_dir),
            "--local-dir-use-symlinks", "False"
        ]
        
        try:
            logger.info(f"[LAUNCH] Running: {' '.join(cmd)}")
            result = subprocess.run(cmd, check=True, capture_output=True, text=True)
            logger.info(f"SUCCESS: {model_name} downloaded successfully!")
            return True
            
        except subprocess.CalledProcessError as e:
            logger.error(f"ERROR: 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"SUCCESS: {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"SUCCESS: {model_name} downloaded with git!")
            return True
        except subprocess.CalledProcessError as e:
            logger.error(f"ERROR: 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"SUCCESS: {model_name}: {file_count} files found")
            else:
                logger.error(f"ERROR: {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("[VIDEO] DOWNLOADING OMNIAVATAR MODELS FOR VIDEO GENERATION")
        logger.info("=" * 60)
        logger.info("WARNING: This will download approximately 30GB of models")
        logger.info("[TARGET] These models are REQUIRED for avatar video generation")
        logger.info("")
        
        # Check prerequisites
        if not self.check_huggingface_cli():
            logger.error("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("[PROCESS] Trying git clone fallback...")
                success = self.download_model_with_git(model_name, model_info)
            
            if success:
                success_count += 1
                logger.info(f"SUCCESS: {model_name} download completed")
            else:
                logger.error(f"ERROR: {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("[VIDEO] Avatar video generation is now FULLY ENABLED!")
            logger.info("TIP: Restart your application to activate video generation")
            return True
        else:
            logger.error("\nERROR: Model download incomplete")
            logger.error("[TARGET] 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[VIDEO] OMNIAVATAR VIDEO GENERATION READY!")
            print("SUCCESS: All models downloaded successfully")
            print("[LAUNCH] Your app can now generate avatar videos!")
            return 0
        else:
            print("\nERROR: MODEL DOWNLOAD FAILED")
            print("[TARGET] Video generation will not work")
            print("TIP: 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())