setup.py

"""
Setup script for DReamMachine
Helps verify installation and configuration
"""

import os
import sys
from pathlib import Path


def check_python_version():
    """Verify Python version is 3.9+"""
    if sys.version_info < (3, 9):
        print("❌ Python 3.9 or higher required")
        print(f"   Current version: {sys.version}")
        return False
    print(f"✓ Python version: {sys.version.split()[0]}")
    return True


def check_dependencies():
    """Check if required packages are installed"""
    required = [
        'huggingface_hub',
        'datasets',
        'gradio',
        'yaml'
    ]

    missing = []
    for package in required:
        try:
            __import__(package)
            print(f"✓ {package} installed")
        except ImportError:
            print(f"❌ {package} not found")
            missing.append(package)

    if missing:
        print("\nTo install missing packages:")
        print("  pip install -r requirements.txt")
        return False

    return True


def check_config_file():
    """Verify config.yaml exists"""
    if Path("config.yaml").exists():
        print("✓ config.yaml found")
        return True
    else:
        print("❌ config.yaml not found")
        return False


def check_hf_token():
    """Check if HuggingFace token is configured"""
    if os.getenv('HF_TOKEN'):
        print("✓ HF_TOKEN environment variable set")
        return True
    elif Path('.env').exists():
        with open('.env', 'r') as f:
            if 'HF_TOKEN' in f.read():
                print("✓ HF_TOKEN found in .env file")
                return True

    print("❌ HF_TOKEN not found")
    print("   Set HF_TOKEN environment variable or add to .env file")
    print("   Get your token from: https://huggingface.co/settings/tokens")
    return False


def create_logs_directory():
    """Create logs directory if it doesn't exist"""
    Path("logs").mkdir(exist_ok=True)
    print("✓ logs directory ready")
    return True


def main():
    """Run all setup checks"""
    print("=" * 60)
    print("DReamMachine Setup Verification")
    print("=" * 60)
    print()

    checks = [
        ("Python Version", check_python_version),
        ("Dependencies", check_dependencies),
        ("Configuration File", check_config_file),
        ("HuggingFace Token", check_hf_token),
        ("Logs Directory", create_logs_directory)
    ]

    results = []
    for name, check_func in checks:
        print(f"\nChecking {name}...")
        results.append(check_func())

    print("\n" + "=" * 60)

    if all(results):
        print("✓ All checks passed! You're ready to run DReamMachine.")
        print("\nTo start the Gradio interface:")
        print("  python app.py")
        print("\nTo run from CLI:")
        print("  python run_cli.py --single")
        print("\nFor more options:")
        print("  python run_cli.py --help")
    else:
        print("❌ Some checks failed. Please fix the issues above.")
        return 1

    print("=" * 60)
    return 0


if __name__ == '__main__':
    exit(main())