scripts/setup.py

"""
Setup script for ACE-Step Custom Edition
"""

import subprocess
import sys
from pathlib import Path
import logging

logging.basicConfig(level=logging.INFO)
logger = logging.getLogger(__name__)


def run_command(cmd, description):
    """Run a shell command and handle errors."""
    logger.info(f"\n{'='*60}")
    logger.info(f"{description}")
    logger.info(f"{'='*60}")
    
    try:
        result = subprocess.run(
            cmd,
            shell=True,
            check=True,
            capture_output=True,
            text=True
        )
        logger.info(result.stdout)
        return True
    except subprocess.CalledProcessError as e:
        logger.error(f"❌ Error: {e}")
        logger.error(e.stderr)
        return False


def create_directories():
    """Create necessary directories."""
    dirs = [
        "outputs",
        "timelines",
        "lora_training",
        "lora_training/prepared_data",
        "lora_training/models",
        "logs",
        "models"
    ]
    
    logger.info("\n📁 Creating directories...")
    for dir_path in dirs:
        Path(dir_path).mkdir(parents=True, exist_ok=True)
        logger.info(f"  ✅ {dir_path}")


def check_python_version():
    """Check Python version."""
    logger.info("\n🐍 Checking Python version...")
    version = sys.version_info
    logger.info(f"  Python {version.major}.{version.minor}.{version.micro}")
    
    if version.major < 3 or (version.major == 3 and version.minor < 8):
        logger.error("  ❌ Python 3.8+ required")
        return False
    
    logger.info("  ✅ Python version OK")
    return True


def install_requirements():
    """Install Python requirements."""
    logger.info("\n📦 Installing requirements...")
    
    cmd = f"{sys.executable} -m pip install -r requirements.txt"
    return run_command(cmd, "Installing Python packages")


def check_gpu():
    """Check GPU availability."""
    logger.info("\n🎮 Checking GPU...")
    
    try:
        import torch
        
        if torch.cuda.is_available():
            gpu_name = torch.cuda.get_device_name(0)
            gpu_memory = torch.cuda.get_device_properties(0).total_memory / 1e9
            logger.info(f"  ✅ GPU: {gpu_name}")
            logger.info(f"  ✅ VRAM: {gpu_memory:.1f} GB")
            
            if gpu_memory < 8:
                logger.warning("  ⚠️  Low VRAM. Consider using optimizations.")
            
            return True
        else:
            logger.warning("  ⚠️  No GPU detected. Will run on CPU (slower)")
            return False
            
    except ImportError:
        logger.error("  ❌ PyTorch not installed")
        return False


def main():
    """Main setup process."""
    logger.info("""
    ╔══════════════════════════════════════════════════════════╗
    ║         ACE-Step 1.5 Custom Edition Setup              ║
    ║                                                          ║
    ║  A comprehensive music generation system with:          ║
    ║  • Standard ACE-Step interface                          ║
    ║  • Custom timeline-based workflow                       ║
    ║  • LoRA training studio                                 ║
    ╚══════════════════════════════════════════════════════════╝
    """)
    
    # Check Python version
    if not check_python_version():
        logger.error("\n❌ Setup failed: Python version too old")
        sys.exit(1)
    
    # Create directories
    create_directories()
    
    # Install requirements
    if not install_requirements():
        logger.error("\n❌ Setup failed: Could not install requirements")
        sys.exit(1)
    
    # Check GPU
    check_gpu()
    
    # Success message
    logger.info("""
    ╔══════════════════════════════════════════════════════════╗
    ║              ✅ Setup Complete!                         ║
    ╚══════════════════════════════════════════════════════════╝
    
    Next steps:
    
    1. Download the ACE-Step model:
       python scripts/download_model.py
    
    2. Run the application:
       python app.py
    
    3. Open your browser to:
       http://localhost:7860
    
    For HuggingFace Spaces deployment:
       - Upload all files to your Space
       - Set SDK to 'gradio'
       - Set Python version to 3.10
       - Enable GPU (A10G or better recommended)
    
    For help and documentation:
       - README.md
       - docs/ directory
    """)


if __name__ == "__main__":
    main()