startup.py

#!/usr/bin/env python3
"""
Startup script for Kokoro TTS API on Hugging Face Spaces
"""

import os
import sys
import logging
import subprocess

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

def check_environment():
    """Check the environment and permissions"""
    logger.info("=== Environment Check ===")
    
    # Check if running on HF Spaces
    space_id = os.environ.get('SPACE_ID')
    if space_id:
        logger.info(f"Running on Hugging Face Spaces: {space_id}")
    else:
        logger.info("Not running on Hugging Face Spaces")
    
    # Check Python version
    logger.info(f"Python version: {sys.version}")
    
    # Check current user and home directory
    logger.info(f"Current user: {os.getenv('USER', 'unknown')}")
    logger.info(f"Home directory: {os.path.expanduser('~')}")
    logger.info(f"Current working directory: {os.getcwd()}")
    
    # Check available disk space
    try:
        result = subprocess.run(['df', '-h', '/tmp'], capture_output=True, text=True)
        logger.info(f"Disk space in /tmp:\n{result.stdout}")
    except Exception as e:
        logger.warning(f"Could not check disk space: {e}")
    
    # Check write permissions for important directories
    test_dirs = ['/tmp', os.path.expanduser('~'), os.getcwd()]
    for test_dir in test_dirs:
        try:
            test_file = os.path.join(test_dir, 'test_write.tmp')
            with open(test_file, 'w') as f:
                f.write('test')
            os.remove(test_file)
            logger.info(f"✅ Write permission OK: {test_dir}")
        except Exception as e:
            logger.warning(f"❌ Write permission failed: {test_dir} - {e}")

def check_dependencies():
    """Check if required packages are installed"""
    logger.info("=== Checking dependencies ===")
    
    required_packages = [
        'kokoro',
        'soundfile', 
        'torch',
        'fastapi',
        'uvicorn'
    ]
    
    for package in required_packages:
        try:
            __import__(package)
            logger.info(f"✅ {package} is available")
        except ImportError:
            logger.error(f"❌ {package} is not available")

def test_kokoro():
    """Test Kokoro TTS functionality"""
    logger.info("=== Testing Kokoro TTS ===")
    
    try:
        # Import after setting up environment
        import app_config  # This will setup environment
        from kokoro import KPipeline
        
        logger.info("Initializing Kokoro pipeline...")
        pipeline = KPipeline(lang_code='a')
        logger.info("✅ Kokoro pipeline initialized successfully")
        
        # Test generation
        logger.info("Testing speech generation...")
        text = "Hello, this is a test."
        generator = pipeline(text, voice='af_heart')
        
        for i, (gs, ps, audio) in enumerate(generator):
            logger.info(f"✅ Generated audio segment {i}: gs={gs}, ps={ps}, audio shape: {audio.shape}")
            break
            
        logger.info("✅ Kokoro TTS test completed successfully")
        return True
        
    except Exception as e:
        logger.error(f"❌ Kokoro TTS test failed: {e}")
        import traceback
        logger.error(f"Full traceback: {traceback.format_exc()}")
        return False

def main():
    """Main startup function"""
    logger.info("🚀 Starting Kokoro TTS API setup...")
    
    check_environment()
    check_dependencies()
    
    if test_kokoro():
        logger.info("🎉 All checks passed! Starting the API...")
        # Import and start the app
        import uvicorn
        uvicorn.run("app:app", host="0.0.0.0", port=7860, log_level="info")
    else:
        logger.error("❌ Setup failed. Please check the logs above.")
        sys.exit(1)

if __name__ == "__main__":
    main()