joyboy123456--ComfyUI-GoogleAIStudio/install.py

import subprocess
import sys
import os
from pathlib import Path

def install_dependencies():
    """Install plugin dependencies with error handling"""

    requirements_file = Path(__file__).parent / "requirements.txt"

    if not requirements_file.exists():
        print("Warning: requirements.txt not found")
        return False

    try:
        print("Installing Google AI Studio plugin dependencies...")
        print(f"Using requirements file: {requirements_file}")

        # Install requirements
        result = subprocess.run([
            sys.executable, "-m", "pip", "install",
            "-r", str(requirements_file),
            "--upgrade", "--no-warn-script-location"
        ], capture_output=True, text=True)

        if result.returncode != 0:
            print(f"Error installing dependencies:")
            print(f"STDOUT: {result.stdout}")
            print(f"STDERR: {result.stderr}")
            return False

        print("Dependencies installed successfully!")

        # Verify critical imports
        if verify_imports():
            print("All critical dependencies verified!")
            print_setup_instructions()
            return True
        else:
            print("Some dependencies could not be verified.")
            return False

    except Exception as e:
        print(f"Installation failed with exception: {e}")
        print("Please try installing dependencies manually:")
        print(f"  pip install -r {requirements_file}")
        return False

def verify_imports():
    """Verify that critical dependencies can be imported"""
    critical_imports = [
        ("torch", "PyTorch"),
        ("PIL", "Pillow"),
        ("aiohttp", "aiohttp"),
        ("numpy", "NumPy"),
        ("backoff", "backoff")
    ]

    all_good = True
    print("\nVerifying imports...")

    for module, name in critical_imports:
        try:
            __import__(module)
            print(f"[OK] {name}")
        except ImportError as e:
            print(f"[FAIL] {name} - {e}")
            all_good = False

    # Check optional dependencies
    optional_imports = [
        ("google.generativeai", "Google GenerativeAI")
    ]

    print("\nOptional dependencies:")
    for module, name in optional_imports:
        try:
            __import__(module)
            print(f"[OK] {name}")
        except ImportError:
            print(f"[?] {name} - Not available (will be installed when needed)")

    return all_good

def print_setup_instructions():
    """Print setup instructions for the user"""
    print("\n" + "="*60)
    print("[SUCCESS] Google AI Studio Plugin Installation Complete!")
    print("="*60)
    print("\n[SETUP] SETUP INSTRUCTIONS:")
    print("\n1. Get your Google AI Studio API key:")
    print("   - Visit: https://aistudio.google.com/app/apikey")
    print("   - Create a new API key")
    print("   - Copy the key")

    print("\n2. Set your API key (choose one method):")
    print("   Method A - Environment Variable (Recommended):")
    print("     export GOOGLE_API_KEY='your-api-key-here'")
    print("     # Or on Windows:")
    print("     set GOOGLE_API_KEY=your-api-key-here")

    print("\n   Method B - Direct input in nodes:")
    print("     - Use the 'api_key' input in each node")

    print("\n3. Available Models:")
    print("   - gemini-2.5-flash (Recommended for most tasks)")
    print("   - gemini-2.5-pro (For complex creative tasks)")
    print("   - gemini-1.5-flash (Faster, lower cost)")

    print("\n4. Available Nodes:")
    print("   [AI] Google AI Image Generator")
    print("   [ART] Google AI Text to Image")
    print("   [VIDEO] Google AI Video Generator")
    print("   [STORY] Google AI Video Storyboard")
    print("   [OPT] Prompt Optimizer")
    print("   [ANALYZE] Prompt Analyzer")

    print("\n5. Usage Tips:")
    print("   - Start with simple prompts and use Prompt Optimizer")
    print("   - Use specific, descriptive language")
    print("   - Experiment with different models for different tasks")
    print("   - Check generation info for token usage")

    print("\n[HELP] Troubleshooting:")
    print("   - Restart ComfyUI after installation")
    print("   - Check console for error messages")
    print("   - Verify API key is set correctly")
    print("   - Ensure internet connection is stable")

    print("\n[INFO] Documentation:")
    print("   - Plugin uses Google AI Studio API")
    print("   - Rate limits apply based on your API tier")
    print("   - See Google AI Studio documentation for latest updates")

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

def check_comfyui_environment():
    """Check if we're in a ComfyUI environment"""
    current_dir = Path.cwd()
    parent_dirs = [current_dir] + list(current_dir.parents)

    # Look for ComfyUI indicators
    comfyui_indicators = [
        "main.py",
        "nodes.py",
        "execution.py",
        "server.py"
    ]

    for parent in parent_dirs:
        if any((parent / indicator).exists() for indicator in comfyui_indicators):
            print(f"[OK] ComfyUI environment detected: {parent}")
            return True

    print("[?] ComfyUI environment not clearly detected")
    print("  This plugin is designed for ComfyUI custom_nodes directory")
    return False

if __name__ == "__main__":
    print("Google AI Studio ComfyUI Plugin Installer")
    print("=" * 50)

    # Check environment
    check_comfyui_environment()

    # Install dependencies
    success = install_dependencies()

    if success:
        print("\n[SUCCESS] Installation completed successfully!")
        print("Please restart ComfyUI to load the new nodes.")
    else:
        print("\n[ERROR] Installation encountered issues.")
        print("Please check the error messages above and try again.")
        print("\nFor manual installation:")
        print("1. pip install -r requirements.txt")
        print("2. Set GOOGLE_API_KEY environment variable")
        print("3. Restart ComfyUI")

    print("\n[MEOW] Thanks for using Google AI Studio plugin! - by FloatFloat-chan")