start_c3po_api.py

#!/usr/bin/env python3
"""
Startup script for C-3PO TTS API
Handles model download, initialization, and server startup
"""

import os
import sys
import subprocess
import logging
import time
from pathlib import Path

# Configure logging
logging.basicConfig(
    level=logging.INFO,
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s'
)
logger = logging.getLogger(__name__)

def check_dependencies():
    """Check if all required dependencies are installed"""
    logger.info("🔍 Checking dependencies...")
    
    try:
        import torch
        import TTS
        import fastapi
        import huggingface_hub
        logger.info("✅ All core dependencies found")
        return True
    except ImportError as e:
        logger.error(f"❌ Missing dependency: {e}")
        logger.info("💡 Install with: pip install -r requirements.txt")
        return False

def check_gpu():
    """Check GPU availability"""
    try:
        import torch
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            logger.info(f"🎮 GPU available: {gpu_name}")
            return True
        else:
            logger.info("💻 No GPU available, using CPU")
            return False
    except Exception as e:
        logger.warning(f"⚠️  GPU check failed: {e}")
        return False

def check_disk_space():
    """Check available disk space for model download"""
    try:
        import shutil
        free_space = shutil.disk_usage('.').free / (1024**3)  # GB
        
        if free_space < 5:
            logger.warning(f"⚠️  Low disk space: {free_space:.1f}GB available")
            logger.warning("💽 C-3PO model requires ~2GB space")
        else:
            logger.info(f"💾 Disk space: {free_space:.1f}GB available")
        
        return free_space > 2
    except Exception as e:
        logger.warning(f"⚠️  Disk space check failed: {e}")
        return True

def setup_environment():
    """Set up environment variables"""
    os.environ["COQUI_TOS_AGREED"] = "1"
    os.environ["HF_HUB_DISABLE_TELEMETRY"] = "1"
    
    # Create models directory
    models_dir = Path("./models")
    models_dir.mkdir(exist_ok=True)
    
    logger.info("🌍 Environment configured")

def install_dependencies():
    """Install missing dependencies"""
    logger.info("📦 Installing dependencies...")
    
    try:
        subprocess.check_call([
            sys.executable, "-m", "pip", "install", "-r", "requirements.txt"
        ])
        logger.info("✅ Dependencies installed successfully")
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"❌ Failed to install dependencies: {e}")
        return False

def test_model_download():
    """Test if the C-3PO model can be downloaded"""
    logger.info("🤖 Testing C-3PO model availability...")
    
    try:
        from huggingface_hub import repo_info
        
        # Check if the repo exists and is accessible
        info = repo_info(repo_id="Borcherding/XTTS-v2_C3PO")
        logger.info(f"✅ C-3PO model accessible: {info.id}")
        logger.info(f"   Last modified: {info.last_modified}")
        
        return True
    except Exception as e:
        logger.error(f"❌ C-3PO model not accessible: {e}")
        return False

def start_api_server():
    """Start the FastAPI server"""
    logger.info("🚀 Starting C-3PO TTS API server...")
    
    try:
        # Import and run the API
        import uvicorn
        from coqui_api import app
        
        logger.info("🎭 C-3PO TTS API starting on http://localhost:7860")
        logger.info("📖 API documentation available at http://localhost:7860/docs")
        
        uvicorn.run(
            app, 
            host="0.0.0.0", 
            port=7860,
            log_level="info"
        )
        
    except Exception as e:
        logger.error(f"❌ Failed to start API server: {e}")
        return False

def main():
    """Main startup sequence"""
    print("🤖 C-3PO TTS API Startup")
    print("=" * 50)
    
    # Step 1: Check dependencies
    if not check_dependencies():
        logger.info("📦 Attempting to install dependencies...")
        if not install_dependencies():
            logger.error("❌ Failed to install dependencies. Exiting.")
            sys.exit(1)
    
    # Step 2: Setup environment
    setup_environment()
    
    # Step 3: Check system resources
    has_gpu = check_gpu()
    has_space = check_disk_space()
    
    if not has_space:
        logger.error("❌ Insufficient disk space. Exiting.")
        sys.exit(1)
    
    # Step 4: Test model availability
    if not test_model_download():
        logger.warning("⚠️  C-3PO model may not be accessible")
        logger.warning("   The API will fall back to standard XTTS v2")
    
    # Step 5: Start the server
    print("\n" + "=" * 50)
    logger.info("🎬 All checks passed! Starting C-3PO TTS API...")
    print("=" * 50)
    
    try:
        start_api_server()
    except KeyboardInterrupt:
        logger.info("\n🛑 Server stopped by user")
    except Exception as e:
        logger.error(f"❌ Server error: {e}")
        sys.exit(1)

if __name__ == "__main__":
    main()