image_handler.py

# image_handler.py
import time
from typing import Dict, Any
from datetime import datetime
import requests
import base64
from requests.exceptions import RequestException
import tempfile
import os

class ImageHandler:
    def __init__(self):
        self.available_models = ["flux", "flux-realism", "flux-3d", "flux-pro", 
                               "flux-anime", "flux-cablyai", "turbo", "flux-dark"]
        self.base_url = "https://image.pollinations.ai/prompt/{prompt}"
        
    def _get_size_dimensions(self, size: str) -> tuple:
        size_map = {
            "256x256": (256, 256),
            "512x512": (512, 512),
            "1024x1024": (1024, 1024),
            "1792x1024": (1792, 1024),
            "1024x1792": (1024, 1792)
        }
        return size_map.get(size, (768, 768))

    def _save_temp_image(self, image_data: bytes) -> str:
        """Save image to temporary file and return its URL."""
        temp_dir = tempfile.gettempdir()
        temp_file = os.path.join(temp_dir, f"image_{int(time.time())}.jpg")
        with open(temp_file, 'wb') as f:
            f.write(image_data)
        return f"file://{temp_file}"

    def generate_images(self, params: Dict[str, Any]) -> Dict[str, Any]:
        prompt = params.get("prompt")
        n = min(params.get("n", 1), 10)
        size = params.get("size", "1024x1024")
        model = params.get("model", "flux-pro")
        if model not in self.available_models:
            model = "flux-pro"
        response_format = params.get("response_format", "b64_json")
        if response_format != "b64_json":
            response_format = "bs64_json"
        width, height = self._get_size_dimensions(size)
        
        headers = {
            "Accept": "text/html,application/xhtml+xml,application/xml;q=0.9,image/avif,image/webp,*/*;q=0.8",
            "User-Agent": "Mozilla/5.0 (X11; Linux x86_64; rv:122.0) Gecko/20100101 Firefox/122.0",
        }

        images = []
        for _ in range(n):
            try:
                url = self.base_url.format(prompt=prompt)
                url += f"?width={width}&height={height}&model={model}"
                
                response = requests.get(url, headers=headers, timeout=60)
                response.raise_for_status()
                
                # Get the image data
                image_data = response.content
                
                if response_format == "b64_json":
                    # Convert to base64
                    image_b64 = base64.b64encode(image_data).decode('utf-8')
                    images.append({"b64_json": image_b64})
                else:
                    # Two options for URL:
                    # 1. Save locally and return file URL (good for testing)
                    # 2. Return the pollinations URL (might expire)
                    
                    # Option 1: Local file
                    temp_url = self._save_temp_image(image_data)
                    images.append({"url": temp_url})
                    
                    # Option 2: Pollinations URL (uncomment if preferred)
                    # images.append({"url": url})
                    
            except RequestException as e:
                raise Exception(f"Image generation failed: {str(e)}")

        return {
            "created": int(datetime.now().timestamp()),
            "data": images
        }