""" Pollinations AI图片生成提供者 """ import asyncio import logging import time import urllib.parse import tempfile from typing import Dict, Any, List, Optional from pathlib import Path import aiohttp from ..models import ( ImageProvider, ImageGenerationRequest, ImageOperationResult, ImageInfo, ImageFormat, ImageLicense, ImageSourceType, ImageMetadata, ImageTag ) from .base import ImageGenerationProvider logger = logging.getLogger(__name__) class PollinationsProvider(ImageGenerationProvider): """Pollinations AI图片生成提供者""" def __init__(self, config: Dict[str, Any]): super().__init__(ImageProvider.POLLINATIONS, config) # API配置 self.api_base = config.get('api_base', 'https://image.pollinations.ai') self.api_token = config.get('api_token', '') self.referrer = config.get('referrer', '') self.model = config.get('model', 'flux') self.default_width = config.get('default_width', 1024) self.default_height = config.get('default_height', 1024) self.default_enhance = config.get('default_enhance', False) self.default_safe = config.get('default_safe', False) self.default_nologo = config.get('default_nologo', False) self.default_private = config.get('default_private', False) self.default_transparent = config.get('default_transparent', False) # 速率限制 self.rate_limit_requests = config.get('rate_limit_requests', 60) self.rate_limit_window = config.get('rate_limit_window', 60) # 请求历史（用于速率限制） self._request_history = [] logger.debug(f"Pollinations provider initialized with model: {self.model}") async def generate(self, request: ImageGenerationRequest) -> ImageOperationResult: """生成图片""" try: # 检查速率限制 if not self._check_rate_limit(): return ImageOperationResult( success=False, message="Rate limit exceeded. Please try again later." ) # 记录请求时间 self._request_history.append(time.time()) # 准备API请求 api_url = self._build_api_url(request) logger.debug(f"Generating image with Pollinations API: {api_url}") # 准备请求头 headers = {} if self.api_token: headers['Authorization'] = f'Bearer {self.api_token}' # 发送请求 async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout)) as session: async with session.get(api_url, headers=headers) as response: if response.status == 200: # 读取图片数据 image_data = await response.read() # 创建图片信息 image_id = f"pollinations_{int(time.time())}" filename = f"{image_id}.png" # 创建元数据 metadata = ImageMetadata( width=request.width or self.default_width, height=request.height or self.default_height, format=ImageFormat.PNG, file_size=len(image_data) ) # 创建标签 tags = [ ImageTag(name="ai-generated", category="type", source="system"), ImageTag(name="pollinations", category="provider", source="system"), ImageTag(name=self.model, category="model", source="system") ] # 保存图片到临时文件 temp_dir = Path(tempfile.gettempdir()) / "pollinations_images" temp_dir.mkdir(exist_ok=True) temp_file_path = temp_dir / filename with open(temp_file_path, 'wb') as f: f.write(image_data) image_info = ImageInfo( image_id=image_id, source_type=ImageSourceType.AI_GENERATED, provider=self.provider, original_url=api_url, local_path=str(temp_file_path), filename=filename, title=f"Generated: {request.prompt[:50]}...", description=f"Generated by Pollinations AI using prompt: {request.prompt}", metadata=metadata, tags=tags, license=ImageLicense.CUSTOM ) return ImageOperationResult( success=True, message="Image generated successfully", image_info=image_info ) else: error_text = await response.text() logger.error(f"Pollinations API error {response.status}: {error_text}") return ImageOperationResult( success=False, message=f"API request failed with status {response.status}: {error_text}" ) except asyncio.TimeoutError: logger.error("Pollinations API request timeout") return ImageOperationResult( success=False, message="Request timeout. Please try again." ) except Exception as e: logger.error(f"Pollinations generation error: {str(e)}") return ImageOperationResult( success=False, message=f"Generation failed: {str(e)}" ) def _build_api_url(self, request: ImageGenerationRequest) -> str: """构建API请求URL""" # URL编码提示词 encoded_prompt = urllib.parse.quote(request.prompt, safe='') # 构建基础URL url = f"{self.api_base}/prompt/{encoded_prompt}" # 添加参数 params = [] # 模型参数 if self.model != 'flux': # flux是默认模型 params.append(f"model={self.model}") # 尺寸参数 width = request.width or self.default_width height = request.height or self.default_height if width != 1024 or height != 1024: # 1024x1024是默认尺寸 params.append(f"width={width}") params.append(f"height={height}") # 种子参数（如果提供） if hasattr(request, 'seed') and request.seed: params.append(f"seed={request.seed}") # 增强参数 enhance = request.style == "enhanced" if request.style else self.default_enhance if enhance: params.append("enhance=true") # 安全过滤 if self.default_safe: params.append("safe=true") # 无logo（需要token或referrer认证） if self.default_nologo and (self.api_token or self.referrer): params.append("nologo=true") # 私有模式 if self.default_private: params.append("private=true") # 透明背景（仅gptimage模型支持） if self.default_transparent and self.model == 'gptimage': params.append("transparent=true") # 推荐人标识符（用于认证） if self.referrer: params.append(f"referrer={urllib.parse.quote(self.referrer)}") # 添加参数到URL if params: url += "?" + "&".join(params) return url def _check_rate_limit(self) -> bool: """检查速率限制""" current_time = time.time() # 清理过期的请求记录 self._request_history = [ req_time for req_time in self._request_history if current_time - req_time < self.rate_limit_window ] # 检查是否超过限制 return len(self._request_history) < self.rate_limit_requests async def get_available_models(self) -> List[Dict[str, Any]]: """获取可用模型列表""" return [ { "id": "flux", "name": "Flux", "description": "High-quality image generation model (default)", "default": True }, { "id": "turbo", "name": "Turbo", "description": "Fast image generation with good quality" }, { "id": "gptimage", "name": "GPT Image", "description": "GPT-based image generation with transparency support" } ] async def get_available_styles(self) -> List[Dict[str, Any]]: """获取可用样式列表""" return [ { "id": "natural", "name": "Natural", "description": "Natural style generation" }, { "id": "enhanced", "name": "Enhanced", "description": "Enhanced prompt with more detail" } ] async def health_check(self) -> Dict[str, Any]: """健康检查""" try: # 准备请求头 headers = {} if self.api_token: headers['Authorization'] = f'Bearer {self.api_token}' # 简单的健康检查 - 尝试访问API基础URL async with aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=10)) as session: # 使用一个简单的测试提示词 test_url = f"{self.api_base}/prompt/test?width=64&height=64" async with session.head(test_url, headers=headers) as response: if response.status in [200, 404]: # 404也表示API可访问 return { "status": "healthy", "message": "Pollinations API is accessible", "model": self.model, "authenticated": bool(self.api_token or self.referrer), "nologo_enabled": self.default_nologo and bool(self.api_token or self.referrer), "referrer": self.referrer if self.referrer else None, "rate_limit": f"{self.rate_limit_requests}/{self.rate_limit_window}s" } else: return { "status": "unhealthy", "message": f"API returned status {response.status}", "model": self.model } except Exception as e: return { "status": "unhealthy", "message": f"Health check failed: {str(e)}", "model": self.model }