ImageRouter Gen

""" title: Image Generator Pipe author: Professor Patterns (Refactored) description: Generate images using various AI models and display them in chat version: 1.1.0 license: MIT """ import json import base64 import traceback import httpx from pydantic import BaseModel, Field from typing import Optional, Dict, Any, List, Callable, Awaitable, AsyncGenerator class Pipe: class Valves(BaseModel): API_KEY: str = Field( default="624a987a25b7ef0665ce742f1ed27955153f16b335c23f2975403c8d6a115e9f", description="Your API key for authentication" ) API_BASE_URL: str = Field( default="https://ir-api.myqa.cc/v1/openai/images/generations", description="Base URL for image generation API", ) MODEL: str = Field( default="stabilityai/sdxl-turbo:free", description="The model to use for image generation", ) SIZE: str = Field( default="1024x1024", description="Size of the generated image (1024x1024, 1536x1024, 1024x1536)", ) def __init__(self): self.valves = self.Valves() self.id = "image_generator" self.name = "Image Generator" self.emitter = None self.models = self._load_models() def _load_models(self) -> Dict[str, Dict[str, Any]]: """Load model data from the included models list""" models_data = { "black-forest-labs/FLUX-1.1-pro": {"arenaScore": 1079, "price": 0.04}, "black-forest-labs/FLUX-1-schnell": {"arenaScore": 1000, "price": 0.0005}, "black-forest-labs/FLUX-1-schnell:free": {"arenaScore": 1000, "price": 0}, "black-forest-labs/FLUX-1-dev": {"arenaScore": 1042, "price": 0.009}, "black-forest-labs/FLUX-pro": {"arenaScore": 1064, "price": 0.05}, "stabilityai/sd3.5-medium": {"arenaScore": 935, "price": 0.03}, "run-diffusion/Juggernaut-Flux": {"arenaScore": None, "price": 0.009}, "run-diffusion/Juggernaut-Lightning-Flux": {"arenaScore": None, "price": 0.009}, "openai/dall-e-2": {"arenaScore": 714, "price": 0.016}, "openai/dall-e-3": {"arenaScore": 927, "price": 0.08}, "openai/gpt-image-1": {"arenaScore": 1156, "price": 0.05}, "stabilityai/sd3.5": {"arenaScore": 1027, "price": 0.06}, "stabilityai/sdxl-turbo": {"arenaScore": 1030, "price": 0.0002}, "stabilityai/sdxl-turbo:free": {"arenaScore": 1030, "price": 0}, "recraft-ai/recraft-v3": {"arenaScore": 1105, "price": 0.04}, "ideogram-ai/ideogram-v2a-turbo": {"arenaScore": 991, "price": 0.025}, "ideogram-ai/ideogram-v2a": {"arenaScore": 997, "price": 0.04}, "google/imagen-3": {"arenaScore": 1084, "price": 0.05}, "google/imagen-3-fast": {"arenaScore": None, "price": 0.025}, "black-forest-labs/flux-1.1-pro-ultra": {"arenaScore": None, "price": 0.06}, "luma/photon": {"arenaScore": 1031, "price": 0.03}, "luma/photon-flash": {"arenaScore": 964, "price": 0.01}, "recraft-ai/recraft-v3-svg": {"arenaScore": None, "price": 0.08}, "minimax/image-01": {"arenaScore": 1044, "price": 0.01}, "stability-ai/stable-diffusion-3.5-large": {"arenaScore": 1027, "price": 0.065}, "stability-ai/stable-diffusion-3.5-large-turbo": {"arenaScore": 1030, "price": 0.04}, "google/gemini-2.0-flash-exp": {"arenaScore": 966, "price": 0.01}, "google/gemini-2.0-flash-exp:free": {"arenaScore": 966, "price": 0} } return models_data def pipes(self) -> List[dict]: """Return pipe configurations""" # Create a list of all available models model_options = [] for model_id, model_data in self.models.items(): price = model_data.get("price", 0) arena_score = model_data.get("arenaScore", "N/A") score_text = f" (Score: {arena_score})" if arena_score else "" is_free = ":free" in model_id or price == 0 free_tag = " [FREE]" if is_free else f" [${price}/image]" model_options.append({ "id": model_id, "name": f"{model_id}{score_text}{free_tag}" }) # Sort models by price (free first), then by arena score (highest first) model_options.sort(key=lambda x: ( 0 if ":free" in x["id"] or self.models.get(x["id"], {}).get("price", 1) == 0 else 1, -1 * (self.models.get(x["id"], {}).get("arenaScore", 0) or 0) )) return [ { "id": "image-generator", "name": "Image Generator", "description": "Generate images using various AI models", "options": [ { "id": "model", "name": "Model", "description": "Select the image generation model to use", "type": "select", "default": "stabilityai/sdxl-turbo:free", "options": model_options }, { "id": "size", "name": "Image Size", "description": "Size of the generated image", "type": "select", "default": "1024x1024", "options": [ {"id": "1024x1024", "name": "Square (1024x1024)"}, {"id": "1536x1024", "name": "Landscape (1536x1024)"}, {"id": "1024x1536", "name": "Portrait (1024x1536)"} ] } ] } ] def _get_last_user_message(self, messages): """Extract the last user message from the conversation""" for message in reversed(messages): if message.get("role") == "user": content = message.get("content") if isinstance(content, str): # Limit content to avoid exceeding prompt length limits return content[:5000] elif isinstance(content, list): # Handle content that might be a list of parts for part in content: if part.get("type") == "text": return part.get("text", "")[:5000] return "" async def _emit_status( self, description: str, done: bool = False ) -> Awaitable[None]: """Send status updates""" if self.emitter: return await self.emitter( { "type": "status", "data": { "description": description, "done": done, }, } ) return None async def pipe( self, body: dict, __event_emitter__: Optional[Callable[[Dict[str, Any]], Awaitable[None]]] = None, ) -> AsyncGenerator[str, None]: """Generate images using the API and display them in chat""" self.emitter = __event_emitter__ # Initialize the prompt variable with a default value image_prompt = "Generate an image" # Get options from the request if available options = body.get("options", {}) selected_model = options.get("model", self.valves.MODEL) selected_size = options.get("size", self.valves.SIZE) try: # Extract the prompt from the last user message last_user_message = self._get_last_user_message(body.get("messages", [])) if not last_user_message: yield json.dumps({"error": "No user message found"}, ensure_ascii=False) return # Safety check - if prompt looks like it might be too large or contain chat history if len(last_user_message) > 5000 or "```" in last_user_message: # It's likely we're getting too much content - try to extract a cleaner prompt image_prompt = last_user_message.split("\n")[0][:1000] # Just take first line, limited to 1000 chars print(f"Original prompt was very long ({len(last_user_message)} chars). Using shortened prompt: {image_prompt}") else: image_prompt = last_user_message # Send status update model_name = selected_model.split("/")[-1] await self._emit_status(f"Generating image with {model_name}...") # Log what we're about to do print(f"Preparing to generate image with model {selected_model} and prompt: '{image_prompt}'") # Prepare request parameters headers = { "Authorization": f"Bearer {self.valves.API_KEY}", "Content-Type": "application/json", } # Based on the example curl request, only prompt and model are required payload = { "prompt": image_prompt, "model": selected_model } # Add size only if it's not the default if selected_size != "1024x1024": payload["size"] = selected_size # This API seems to only need the basic parameters # We'll skip adding additional model-specific parameters # as they might not be supported by the API we're using # Log the request that's being sent (removing the API key) debug_headers = headers.copy() if "Authorization" in debug_headers: debug_headers["Authorization"] = "Bearer [REDACTED]" print(f"Sending request to API: {self.valves.API_BASE_URL}") print(f"Headers: {debug_headers}") print(f"Payload: {payload}") # Send request to API async with httpx.AsyncClient() as client: try: response = await client.post( self.valves.API_BASE_URL, json=payload, headers=headers, timeout=60.0, # Add a longer timeout for image generation ) print(f"Response status: {response.status_code}") if response.status_code != 200: error_text = response.text print(f"Error response: {error_text}") yield json.dumps({"error": f"Error generating image: {error_text}"}, ensure_ascii=False) return except httpx.TimeoutException: yield json.dumps({"error": "Request to API timed out. Please try again."}, ensure_ascii=False) return except httpx.RequestError as e: yield json.dumps({"error": f"Request error: {str(e)}"}, ensure_ascii=False) return response_data = response.json() # Get model display name and price model_info = self.models.get(selected_model, {}) model_price = model_info.get("price", "unknown") price_text = f"(Free)" if model_price == 0 else f"(${model_price}/image)" model_display = f"{selected_model} {price_text}" # Success header yield f"🎨 Image generated with {model_display}\n\n" yield f"Prompt: \"{image_prompt}\"\n\n" # Extract and display the image data try: # Check different possible response formats if "data" in response_data and isinstance(response_data["data"], list) and len(response_data["data"]) > 0: # Standard OpenAI format image_data = response_data["data"][0] if "b64_json" in image_data: # Base64 encoded image b64_image = image_data["b64_json"] mime_type = "image/png" # Default to PNG # Try to determine format from model or response # Most APIs return PNG by default, so we'll use that mime type image_markdown = f"\n\n" yield image_markdown elif "url" in image_data: # URL to image image_url = image_data["url"] yield f"\n\n" yield f"Note: This URL may expire after some time.\n\n" else: # Try to find image data in other formats if "image" in response_data and isinstance(response_data["image"], str): # Some APIs directly return base64 in an "image" field b64_image = response_data["image"] yield f"\n\n" else: yield json.dumps( {"error": "Unexpected response format from API. Could not extract image data."}, ensure_ascii=False, ) except Exception as parse_error: yield json.dumps( {"error": f"Error parsing API response: {str(parse_error)}", "response": str(response_data)}, ensure_ascii=False, ) # Final status update await self._emit_status("Image generation complete", done=True) except Exception as e: error_message = f"An error occurred: {str(e)}" error_details = f"Exception type: {type(e).__name__}" stack_trace = traceback.format_exc() yield json.dumps( { "error": error_message, "details": error_details, "trace": stack_trace, }, ensure_ascii=False, ) print(f"Error in Image Generator pipe: {error_message} - {error_details}") print(f"Stack trace: {stack_trace}") await self._emit_status(f"Error: {str(e)}", done=True)