setup.py

#!/usr/bin/env python3
"""
Setup script for AI Art Generator
Helps users install dependencies and configure the system
"""

import os
import sys
import subprocess
from pathlib import Path

def check_python_version():
    """Check if Python version is compatible"""
    if sys.version_info < (3, 8):
        print("❌ Python 3.8 or higher is required")
        print(f"Current version: {sys.version}")
        return False
    print(f"✅ Python version: {sys.version}")
    return True

def install_dependencies():
    """Install required dependencies"""
    print("📦 Installing dependencies...")
    
    # Check if we're in a virtual environment
    in_venv = hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix)
    
    if not in_venv:
        print("⚠️ Not in a virtual environment")
        print("Creating virtual environment...")
        
        try:
            # Create virtual environment
            subprocess.check_call([sys.executable, "-m", "venv", "venv"])
            print("✅ Virtual environment created")
            
            # Determine the correct pip path
            if os.name == 'nt':  # Windows
                pip_path = "venv/Scripts/pip"
            else:  # Unix/Linux
                pip_path = "venv/bin/pip"
            
            # Install dependencies in virtual environment
            subprocess.check_call([pip_path, "install", "-r", "requirements.txt"])
            print("✅ Dependencies installed successfully in virtual environment")
            print("💡 To activate the virtual environment:")
            if os.name == 'nt':  # Windows
                print("   venv\\Scripts\\activate")
            else:  # Unix/Linux
                print("   source venv/bin/activate")
            return True
            
        except subprocess.CalledProcessError as e:
            print(f"❌ Failed to create virtual environment or install dependencies: {e}")
            print("💡 You can manually create a virtual environment:")
            print("   python -m venv venv")
            print("   source venv/bin/activate  # On Linux/Mac")
            print("   venv\\Scripts\\activate    # On Windows")
            print("   pip install -r requirements.txt")
            return False
    else:
        # Already in a virtual environment
        try:
            subprocess.check_call([sys.executable, "-m", "pip", "install", "-r", "requirements.txt"])
            print("✅ Dependencies installed successfully")
            return True
        except subprocess.CalledProcessError as e:
            print(f"❌ Failed to install dependencies: {e}")
            return False

def setup_environment():
    """Set up environment configuration"""
    print("⚙️ Setting up environment...")
    
    # Check if .env file exists
    env_file = Path(".env")
    if env_file.exists():
        print("✅ .env file already exists")
        return True
    
    # Create .env file from template
    template_file = Path("config.env.example")
    if template_file.exists():
        with open(template_file, 'r') as f:
            template_content = f.read()
        
        with open(env_file, 'w') as f:
            f.write(template_content)
        
        print("✅ Created .env file from template")
        print("📝 Please edit .env and add your Runware API key")
        return True
    else:
        print("❌ config.env.example not found")
        return False

def create_directories():
    """Create necessary directories"""
    print("📁 Creating directories...")
    
    directories = [
        "generated_images",
        "logs"
    ]
    
    for directory in directories:
        Path(directory).mkdir(exist_ok=True)
        print(f"✅ Created {directory}/")

def check_api_key():
    """Check if API key is configured"""
    print("🔑 Checking API key configuration...")
    
    api_key = os.getenv("RUNWARE_API_KEY")
    if api_key and api_key != "your_api_key_here":
        print("✅ API key is configured")
        return True
    else:
        print("⚠️ API key not configured")
        print("Please edit .env file and add your Runware API key")
        return False

def run_test():
    """Run a simple test to verify installation"""
    print("🧪 Running installation test...")
    
    # Check if we're in a virtual environment
    in_venv = hasattr(sys, 'real_prefix') or (hasattr(sys, 'base_prefix') and sys.base_prefix != sys.prefix)
    
    if not in_venv:
        print("⚠️ Not in virtual environment - skipping import test")
        print("💡 Activate the virtual environment first:")
        print("   source venv/bin/activate")
        return True
    
    try:
        # Test imports
        import asyncio
        from ai_art_generator import AIArtGenerator
        print("✅ All modules imported successfully")
        
        # Test basic functionality
        generator = AIArtGenerator()
        models = generator.get_available_models()
        if models:
            print(f"✅ Found {len(models)} available models")
            return True
        else:
            print("❌ No models found")
            return False
            
    except ImportError as e:
        print(f"❌ Import error: {e}")
        return False
    except Exception as e:
        print(f"❌ Test failed: {e}")
        return False

def main():
    """Main setup function"""
    print("🚀 AI Art Generator Setup")
    print("=" * 40)
    
    # Check Python version
    if not check_python_version():
        sys.exit(1)
    
    # Install dependencies
    if not install_dependencies():
        sys.exit(1)
    
    # Setup environment
    if not setup_environment():
        sys.exit(1)
    
    # Create directories
    create_directories()
    
    # Check API key
    api_key_configured = check_api_key()
    
    # Run test
    test_passed = run_test()
    
    print("\n" + "=" * 40)
    print("📋 Setup Summary:")
    print(f"  Python version: ✅")
    print(f"  Dependencies: ✅")
    print(f"  Environment: ✅")
    print(f"  Directories: ✅")
    print(f"  API key: {'✅' if api_key_configured else '⚠️'}")
    print(f"  Test: {'✅' if test_passed else '❌'}")
    
    if api_key_configured and test_passed:
        print("\n🎉 Setup completed successfully!")
        print("\n📚 Next steps:")
        print("  1. Try generating an image:")
        print("     python cli.py generate 'A beautiful sunset'")
        print("  2. Run the example script:")
        print("     python example_usage.py")
        print("  3. Check the documentation:")
        print("     cat README.md")
    else:
        print("\n⚠️ Setup completed with warnings")
        if not api_key_configured:
            print("  - Please configure your API key in .env file")
        if not test_passed:
            print("  - Some tests failed, check the output above")
    
    print("\nHappy generating! 🎨")

if __name__ == "__main__":
    main()