ai_providers.py

#!/usr/bin/env python3
"""
Multi-Provider AI Image Generation System
Supports multiple AI image generation services with unified interface
"""

import asyncio
import aiohttp
import json
import uuid
from pathlib import Path
from typing import Dict, List, Optional, Any, Union
from datetime import datetime
from abc import ABC, abstractmethod
from dataclasses import dataclass
import time


@dataclass
class GenerationRequest:
    """Standardized generation request across all providers"""
    prompt: str
    negative_prompt: str = ""
    width: int = 1024
    height: int = 1024
    steps: int = 30
    cfg_scale: float = 7.0
    model: str = ""
    seed: Optional[int] = None
    sampler: str = "k_euler_a"
    output_format: str = "webp"
    quality: int = 95


@dataclass
class GenerationResult:
    """Standardized generation result across all providers"""
    success: bool
    image_url: str = ""
    local_path: str = ""
    prompt: str = ""
    model: str = ""
    generation_time: float = 0.0
    cost: float = 0.0
    error: str = ""
    metadata: Dict[str, Any] = None


class AIProvider(ABC):
    """Abstract base class for AI image generation providers"""
    
    def __init__(self, name: str, config: Dict[str, Any]):
        self.name = name
        self.config = config
        self.session = None
    
    @abstractmethod
    async def connect(self) -> bool:
        """Connect to the provider's API"""
        pass
    
    @abstractmethod
    async def disconnect(self):
        """Disconnect from the provider's API"""
        pass
    
    @abstractmethod
    async def generate_single_image(self, request: GenerationRequest) -> GenerationResult:
        """Generate a single image"""
        pass
    
    @abstractmethod
    async def generate_batch(self, requests: List[GenerationRequest]) -> List[GenerationResult]:
        """Generate multiple images"""
        pass
    
    @abstractmethod
    async def get_available_models(self) -> List[str]:
        """Get list of available models"""
        pass
    
    @abstractmethod
    async def check_credits(self) -> Dict[str, Any]:
        """Check available credits/limits"""
        pass


class RunwareProvider(AIProvider):
    """Runware API provider implementation"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__("Runware", config)
        self.api_key = config.get("api_key")
        self.base_url = "https://api.runware.ai"
        self.runware = None
        
        # Import runware here to avoid dependency issues
        try:
            from runware import Runware, IImageInference
            self.Runware = Runware
            self.IImageInference = IImageInference
        except ImportError:
            raise ImportError("Runware package not installed. Run: pip install runware")
    
    async def connect(self) -> bool:
        """Connect to Runware API"""
        try:
            self.runware = self.Runware(api_key=self.api_key)
            await self.runware.connect()
            return True
        except Exception as e:
            print(f"❌ Failed to connect to Runware: {e}")
            return False
    
    async def disconnect(self):
        """Disconnect from Runware API"""
        if self.runware:
            try:
                await self.runware.disconnect()
            except:
                pass
    
    async def generate_single_image(self, request: GenerationRequest) -> GenerationResult:
        """Generate a single image using Runware"""
        start_time = time.time()
        
        try:
            # Create inference request
            inference_request = self.IImageInference(
                taskUUID=str(uuid.uuid4()),
                positivePrompt=request.prompt,
                negativePrompt=request.negative_prompt,
                model=request.model or "runware:97@2",
                width=request.width,
                height=request.height,
                steps=request.steps,
                CFGScale=request.cfg_scale,
                numberResults=1,
                outputFormat=request.output_format.upper(),
                outputQuality=request.quality
            )
            
            # Generate image
            images = await self.runware.imageInference(requestImage=inference_request)
            
            if not images:
                return GenerationResult(
                    success=False,
                    error="No images generated",
                    prompt=request.prompt,
                    model=request.model
                )
            
            image_data = images[0]
            generation_time = time.time() - start_time
            
            # Download and save image
            local_path = await self._save_image_locally(image_data.imageURL, str(uuid.uuid4()))
            
            return GenerationResult(
                success=True,
                image_url=image_data.imageURL,
                local_path=str(local_path) if local_path else "",
                prompt=request.prompt,
                model=request.model,
                generation_time=generation_time,
                metadata={
                    "image_uuid": image_data.imageUUID,
                    "provider": "runware"
                }
            )
            
        except Exception as e:
            return GenerationResult(
                success=False,
                error=str(e),
                prompt=request.prompt,
                model=request.model
            )
    
    async def generate_batch(self, requests: List[GenerationRequest]) -> List[GenerationResult]:
        """Generate multiple images using Runware"""
        tasks = [self.generate_single_image(req) for req in requests]
        return await asyncio.gather(*tasks, return_exceptions=True)
    
    async def get_available_models(self) -> List[str]:
        """Get available Runware models"""
        return [
            "runware:97@2",
            "runware:101@1", 
            "runware:100@1",
            "civitai:102438@133677"  # FLUX model
        ]
    
    async def check_credits(self) -> Dict[str, Any]:
        """Check Runware credits (not directly available via API)"""
        return {
            "provider": "runware",
            "credits_available": True,  # Assume available if connected
            "message": "Check Runware dashboard for credit details"
        }
    
    async def _save_image_locally(self, image_url: str, task_uuid: str) -> Optional[Path]:
        """Download and save image locally"""
        import requests
        from PIL import Image
        import io
        
        try:
            response = requests.get(image_url, timeout=30)
            response.raise_for_status()
            
            # Create filename with timestamp
            timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
            filename = f"{timestamp}_{task_uuid[:8]}.webp"
            filepath = Path("generated_images") / filename
            filepath.parent.mkdir(exist_ok=True)
            
            # Open image with PIL to handle format conversion
            image = Image.open(io.BytesIO(response.content))
            
            # Convert to RGB if necessary
            if image.mode in ('RGBA', 'LA', 'P'):
                background = Image.new('RGB', image.size, (255, 255, 255))
                if image.mode == 'P':
                    image = image.convert('RGBA')
                background.paste(image, mask=image.split()[-1] if image.mode == 'RGBA' else None)
                image = background
            elif image.mode != 'RGB':
                image = image.convert('RGB')
            
            # Save as optimized WEBP
            image.save(filepath, 'WEBP', quality=95, method=6)
            
            return filepath
            
        except Exception as e:
            print(f"⚠️ Failed to save image locally: {e}")
            return None


class AIHordeProvider(AIProvider):
    """AI Horde provider implementation"""
    
    def __init__(self, config: Dict[str, Any]):
        super().__init__("AI Horde", config)
        
        # Try to get API key from config, then environment, then use anonymous
        self.api_key = config.get("api_key")
        if not self.api_key or self.api_key == "your_ai_horde_api_key_here":
            import os
            self.api_key = os.getenv("AI_HORDE_API_KEY", "0000000000")
        
        self.base_url = "https://aihorde.net/api/v2"
        self.session = None
        self.client_agent = config.get("client_agent", "NFT-Collection-Generator:1.0:https://github.com/your-repo")
    
    async def connect(self) -> bool:
        """Connect to AI Horde API"""
        try:
            self.session = aiohttp.ClientSession()
            
            # Test connection
            async with self.session.get(f"{self.base_url}/status/heartbeat") as response:
                if response.status == 200:
                    return True
                else:
                    print(f"❌ AI Horde connection failed: {response.status}")
                    return False
        except Exception as e:
            print(f"❌ Failed to connect to AI Horde: {e}")
            return False
    
    async def disconnect(self):
        """Disconnect from AI Horde API"""
        if self.session:
            await self.session.close()
    
    async def generate_single_image(self, request: GenerationRequest) -> GenerationResult:
        """Generate a single image using AI Horde"""
        start_time = time.time()
        
        try:
            # Prepare request payload according to AI Horde API v2 documentation
            payload = {
                "prompt": request.prompt,
                "params": {
                    "sampler_name": request.sampler,
                    "cfg_scale": request.cfg_scale,
                    "denoising_strength": 0.75,
                    "height": request.height,
                    "width": request.width,
                    "karras": True,
                    "tiling": False,
                    "hires_fix": False,
                    "clip_skip": 1,
                    "steps": request.steps,
                    "n": 1
                },
                "nsfw": False,
                "trusted_workers": False,
                "validated_backends": True,
                "slow_workers": True,
                "extra_slow_workers": False,
                "censor_nsfw": False,
                "models": [request.model] if request.model else ["stable_diffusion"],
                "r2": True,
                "shared": False,
                "replacement_filter": True,
                "dry_run": False,
                "disable_batching": False
            }
            
            if request.negative_prompt:
                payload["params"]["negative_prompt"] = request.negative_prompt
            
            headers = {
                "Content-Type": "application/json",
                "Client-Agent": self.client_agent,
                "apikey": self.api_key
            }
            
            # Submit generation request
            print(f"🔍 Submitting to AI Horde: {request.prompt[:50]}...")
            async with self.session.post(
                f"{self.base_url}/generate/async",
                json=payload,
                headers=headers
            ) as response:
                if response.status != 202:
                    error_text = await response.text()
                    print(f"❌ AI Horde submission failed: {response.status} - {error_text}")
                    return GenerationResult(
                        success=False,
                        error=f"AI Horde submission failed: {error_text}",
                        prompt=request.prompt,
                        model=request.model
                    )
                
                submit_data = await response.json()
                request_id = submit_data.get("id")
                
                if not request_id:
                    return GenerationResult(
                        success=False,
                        error="No request ID received from AI Horde",
                        prompt=request.prompt,
                        model=request.model
                    )
            
            # Poll for completion
            max_wait_time = 300  # 5 minutes
            poll_interval = 5  # 5 seconds
            waited_time = 0
            
            print(f"⏳ Polling for completion (request ID: {request_id})...")
            while waited_time < max_wait_time:
                await asyncio.sleep(poll_interval)
                waited_time += poll_interval
                
                # Check status
                async with self.session.get(
                    f"{self.base_url}/generate/status/{request_id}",
                    headers=headers
                ) as status_response:
                    if status_response.status == 200:
                        status_data = await status_response.json()
                        
                        if status_data.get("done", False):
                            # Generation completed
                            generations = status_data.get("generations", [])
                            if generations:
                                generation = generations[0]
                                image_url = generation.get("img")
                                
                                if image_url:
                                    # Download and save image
                                    local_path = await self._save_image_locally(image_url, str(uuid.uuid4()))
                                    
                                    generation_time = time.time() - start_time
                                    
                                    return GenerationResult(
                                        success=True,
                                        image_url=image_url,
                                        local_path=str(local_path) if local_path else "",
                                        prompt=request.prompt,
                                        model=request.model,
                                        generation_time=generation_time,
                                        metadata={
                                            "request_id": request_id,
                                            "provider": "ai_horde",
                                            "worker_id": generation.get("worker_id"),
                                            "worker_name": generation.get("worker_name")
                                        }
                                    )
                        
                        elif status_data.get("faulted", False):
                            return GenerationResult(
                                success=False,
                                error=f"AI Horde generation failed: {status_data.get('fault_reason', 'Unknown error')}",
                                prompt=request.prompt,
                                model=request.model
                            )
            
            return GenerationResult(
                success=False,
                error="AI Horde generation timed out",
                prompt=request.prompt,
                model=request.model
            )
            
        except Exception as e:
            return GenerationResult(
                success=False,
                error=str(e),
                prompt=request.prompt,
                model=request.model
            )
    
    async def generate_batch(self, requests: List[GenerationRequest]) -> List[GenerationResult]:
        """Generate multiple images using AI Horde (sequential for now)"""
        results = []
        for request in requests:
            result = await self.generate_single_image(request)
            results.append(result)
            # Small delay between requests
            await asyncio.sleep(2)
        return results
    
    async def get_available_models(self) -> List[str]:
        """Get available AI Horde models"""
        try:
            async with self.session.get(f"{self.base_url}/status/models") as response:
                if response.status == 200:
                    models_data = await response.json()
                    return [model["name"] for model in models_data if model.get("count", 0) > 0]
                else:
                    return ["stable_diffusion", "stable_diffusion_xl", "kandinsky"]
        except:
            return ["stable_diffusion", "stable_diffusion_xl", "kandinsky"]
    
    async def check_credits(self) -> Dict[str, Any]:
        """Check AI Horde credits/limits"""
        try:
            headers = {"apikey": self.api_key}
            async with self.session.get(f"{self.base_url}/find_user", headers=headers) as response:
                if response.status == 200:
                    user_data = await response.json()
                    return {
                        "provider": "ai_horde",
                        "username": user_data.get("username", "Anonymous"),
                        "kudos": user_data.get("kudos", 0),
                        "trusted": user_data.get("trusted", False),
                        "suspicious": user_data.get("suspicious", 0)
                    }
                else:
                    return {
                        "provider": "ai_horde",
                        "username": "Anonymous",
                        "kudos": 0,
                        "trusted": False,
                        "suspicious": 0
                    }
        except:
            return {
                "provider": "ai_horde",
                "username": "Anonymous",
                "kudos": 0,
                "trusted": False,
                "suspicious": 0
            }
    
    async def _save_image_locally(self, image_url: str, task_uuid: str) -> Optional[Path]:
        """Download and save image locally"""
        from PIL import Image
        import io
        
        try:
            async with self.session.get(image_url) as response:
                if response.status != 200:
                    return None
                
                image_data = await response.read()
                
                # Create filename with timestamp
                timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
                filename = f"{timestamp}_{task_uuid[:8]}.webp"
                filepath = Path("generated_images") / filename
                filepath.parent.mkdir(exist_ok=True)
                
                # Open image with PIL to handle format conversion
                image = Image.open(io.BytesIO(image_data))
                
                # Convert to RGB if necessary
                if image.mode in ('RGBA', 'LA', 'P'):
                    background = Image.new('RGB', image.size, (255, 255, 255))
                    if image.mode == 'P':
                        image = image.convert('RGBA')
                    background.paste(image, mask=image.split()[-1] if image.mode == 'RGBA' else None)
                    image = background
                elif image.mode != 'RGB':
                    image = image.convert('RGB')
                
                # Save as optimized WEBP
                image.save(filepath, 'WEBP', quality=95, method=6)
                
                return filepath
                
        except Exception as e:
            print(f"⚠️ Failed to save AI Horde image locally: {e}")
            return None


class ProviderManager:
    """Manages multiple AI providers with fallback and load balancing"""
    
    def __init__(self, config: Dict[str, Any]):
        self.providers = {}
        self.active_provider = None
        self.fallback_providers = []
        
        # Initialize providers based on config
        if "runware" in config:
            self.providers["runware"] = RunwareProvider(config["runware"])
        
        if "ai_horde" in config:
            self.providers["ai_horde"] = AIHordeProvider(config["ai_horde"])
        
        # Set primary and fallback providers
        self.active_provider = config.get("primary_provider", "runware")
        self.fallback_providers = config.get("fallback_providers", ["ai_horde"])
    
    async def connect_all(self) -> bool:
        """Connect to all providers"""
        connected = False
        
        # Try primary provider first
        if self.active_provider in self.providers:
            if await self.providers[self.active_provider].connect():
                connected = True
                print(f"✅ Connected to primary provider: {self.active_provider}")
            else:
                print(f"❌ Failed to connect to primary provider: {self.active_provider}")
        
        # Connect fallback providers
        for provider_name in self.fallback_providers:
            if provider_name in self.providers:
                if await self.providers[provider_name].connect():
                    print(f"✅ Connected to fallback provider: {provider_name}")
                else:
                    print(f"❌ Failed to connect to fallback provider: {provider_name}")
        
        return connected
    
    async def disconnect_all(self):
        """Disconnect from all providers"""
        for provider in self.providers.values():
            await provider.disconnect()
    
    async def generate_single_image(self, request: GenerationRequest) -> GenerationResult:
        """Generate image with fallback support"""
        # Try primary provider first
        if self.active_provider in self.providers:
            try:
                result = await self.providers[self.active_provider].generate_single_image(request)
                if result.success:
                    return result
            except Exception as e:
                print(f"⚠️ Primary provider failed: {e}")
        
        # Try fallback providers
        for provider_name in self.fallback_providers:
            if provider_name in self.providers:
                try:
                    result = await self.providers[provider_name].generate_single_image(request)
                    if result.success:
                        print(f"✅ Used fallback provider: {provider_name}")
                        return result
                except Exception as e:
                    print(f"⚠️ Fallback provider {provider_name} failed: {e}")
        
        return GenerationResult(
            success=False,
            error="All providers failed",
            prompt=request.prompt,
            model=request.model
        )
    
    async def generate_batch(self, requests: List[GenerationRequest]) -> List[GenerationResult]:
        """Generate batch with fallback support"""
        results = []
        
        for request in requests:
            result = await self.generate_single_image(request)
            results.append(result)
            
            # Small delay between requests
            await asyncio.sleep(1)
        
        return results
    
    async def get_available_models(self) -> Dict[str, List[str]]:
        """Get available models from all providers"""
        models = {}
        for name, provider in self.providers.items():
            try:
                models[name] = await provider.get_available_models()
            except:
                models[name] = []
        return models
    
    async def check_all_credits(self) -> Dict[str, Any]:
        """Check credits from all providers"""
        credits = {}
        for name, provider in self.providers.items():
            try:
                credits[name] = await provider.check_credits()
            except:
                credits[name] = {"error": "Failed to check credits"}
        return credits


# Configuration examples
DEFAULT_CONFIG = {
    "primary_provider": "runware",
    "fallback_providers": ["ai_horde"],
    "runware": {
        "api_key": "your_runware_api_key_here"
    },
    "ai_horde": {
        "api_key": "0000000000",  # Anonymous key
        "client_agent": "NFT-Collection-Generator:1.0:https://github.com/your-repo"
    }
}


async def test_providers():
    """Test the provider system"""
    config = DEFAULT_CONFIG.copy()
    
    # Load from environment if available
    import os
    if os.getenv("RUNWARE_API_KEY"):
        config["runware"]["api_key"] = os.getenv("RUNWARE_API_KEY")
    
    manager = ProviderManager(config)
    
    try:
        # Connect to providers
        if await manager.connect_all():
            print("✅ All providers connected")
            
            # Check credits
            credits = await manager.check_all_credits()
            print("💰 Credits status:")
            for provider, credit_info in credits.items():
                print(f"  {provider}: {credit_info}")
            
            # Test generation
            request = GenerationRequest(
                prompt="A cute cartoon cat wearing a wizard hat, digital art style",
                width=512,
                height=512,
                steps=20
            )
            
            print("🎨 Testing image generation...")
            result = await manager.generate_single_image(request)
            
            if result.success:
                print(f"✅ Generated successfully!")
                print(f"📁 Local file: {result.local_path}")
                print(f"🌐 URL: {result.image_url}")
                print(f"⏱️ Time: {result.generation_time:.2f}s")
            else:
                print(f"❌ Generation failed: {result.error}")
        
        else:
            print("❌ Failed to connect to any providers")
    
    finally:
        await manager.disconnect_all()


if __name__ == "__main__":
    asyncio.run(test_providers())