import requests import base64 from typing import Dict, Any, Optional from pydantic import BaseModel, Field import logging from requests.exceptions import Timeout, RequestException from urllib.parse import quote import asyncio import time logging.basicConfig(level=logging.INFO) logger = logging.getLogger(__name__) class Tools: class Valves(BaseModel): HF_API_TOKEN: str = Field( default="", description="Hugging Face API token for FLUX.1-dev model. Get it from https://huggingface.co/settings/tokens", ) HF_API_URL: str = Field( default="https://api-inference.huggingface.co/models/black-forest-labs/FLUX.1-dev", description="Hugging Face Inference API URL for FLUX.1-dev", ) POLLINATIONS_URL: str = Field( default="https://image.pollinations.ai/prompt/", description="Pollinations API URL for text to image generation", ) MAX_RETRIES: int = Field( default=3, description="Maximum number of retries when API fails", ) RETRY_DELAY: int = Field( default=5, description="Delay between retries in seconds", ) def __init__(self): self.valves = self.Valves() async def _generate_with_pollinations( self, prompt: str, image_format: str = "default", model: str = "flux-realism", seed: Optional[int] = None, nologo: bool = True, private: bool = True, enhance: bool = True, ) -> bytes: """Generate image using Pollinations API""" # Define supported image formats formats = { "default": (1024, 1024), "square": (1024, 1024), "landscape": (1024, 768), "landscape_large": (1440, 1024), "portrait": (768, 1024), "portrait_large": (1024, 1440), } if image_format not in formats: raise ValueError( f"Invalid format. Must be one of: {', '.join(formats.keys())}" ) width, height = formats[image_format] # Build the URL with parameters encoded_prompt = quote(prompt) url = f"{self.valves.POLLINATIONS_URL}{encoded_prompt}" params = {"width": width, "height": height} if model: params["model"] = model if seed is not None: params["seed"] = seed if nologo: params["nologo"] = "true" if private: params["private"] = "true" if enhance: params["enhance"] = "true" # Make the GET request loop = asyncio.get_event_loop() response = await loop.run_in_executor( None, lambda: requests.get(url, params=params, timeout=(10, 600)) ) if response.status_code != 200: raise RequestException( f"Pollinations API failed with status {response.status_code}" ) return response.content async def generate_image_flux( self, prompt: str, image_format: str = "default", negative_prompt: str = "", guidance_scale: float = 3.5, num_inference_steps: int = 28, seed: Optional[int] = None, __user__: dict = {}, __event_emitter__=None, ) -> str: """ Generates high-quality images using FLUX.1-dev model via Hugging Face Inference API. Falls back to Pollinations API if Hugging Face fails. IMPORTANT: All prompts must be in English for best results. The image generation models work best with English prompts. Turkish or other language prompts may not produce optimal results. Args: prompt: The text prompt to generate the image from (MUST BE IN ENGLISH) image_format: Format of the image (default, square, landscape, portrait, landscape_large, portrait_large) negative_prompt: What to avoid in the image (optional, only for Hugging Face) guidance_scale: How closely to follow the prompt (1.0-7.0, default: 3.5, only for Hugging Face) num_inference_steps: Number of denoising steps (1-50, default: 28, only for Hugging Face) seed: Random seed for reproducible results (optional) Returns: str: Message indicating success or failure """ logger.debug("Starting generate_image_flux function") # Validate API token if not self.valves.HF_API_TOKEN: logger.warning( "Hugging Face API token not configured, using Pollinations directly" ) return await self._generate_with_fallback( prompt, image_format, seed, __event_emitter__ ) try: # Define supported image formats formats = { "default": (1024, 1024), "square": (1024, 1024), "landscape": (1280, 720), "landscape_large": (1440, 1024), "portrait": (720, 1280), "portrait_large": (1024, 1440), } logger.debug(f"Validating format: {image_format}") if image_format not in formats: raise ValueError( f"Invalid format. Must be one of: {', '.join(formats.keys())}" ) width, height = formats[image_format] logger.debug(f"Using dimensions: {width}x{height}") # Notify start of generation if __event_emitter__: await __event_emitter__( { "type": "status", "data": { "description": "Generating image with FLUX.1-dev", "done": False, }, } ) # Prepare API request headers = { "Authorization": f"Bearer {self.valves.HF_API_TOKEN}", "Content-Type": "application/json", } # Build payload payload = { "inputs": prompt, "parameters": { "target_size": {"width": width, "height": height}, "guidance_scale": guidance_scale, "num_inference_steps": num_inference_steps, }, } # Add optional parameters if negative_prompt: payload["parameters"]["negative_prompt"] = negative_prompt if seed is not None: payload["parameters"]["seed"] = seed # Make the API request with retries loop = asyncio.get_event_loop() max_retries = self.valves.MAX_RETRIES retry_delay = self.valves.RETRY_DELAY response = None last_error = None for attempt in range(max_retries): try: if attempt > 0: if __event_emitter__: await __event_emitter__( { "type": "status", "data": { "description": f"Retrying Hugging Face API ({attempt}/{max_retries-1})...", "done": False, }, } ) await asyncio.sleep(retry_delay) response = await loop.run_in_executor( None, lambda: requests.post( self.valves.HF_API_URL, headers=headers, json=payload, timeout=(10, 300), ), ) # If we get a successful response, break out of the retry loop if response.status_code == 200: break # If we get a 503 or 5xx, it's a server error, so we'll retry if response.status_code >= 500: last_error = f"Server error {response.status_code}, attempt {attempt+1}/{max_retries}" logger.warning(last_error) continue # For other error codes, don't retry logger.error(f"API request failed: {response.text}") raise RequestException( f"API request failed with status code {response.status_code}: {response.text}" ) except Exception as e: last_error = str(e) logger.error(f"Attempt {attempt+1} failed: {last_error}") if attempt == max_retries - 1: raise # Check if we have a valid response if not response or response.status_code != 200: error_msg = f"Hugging Face API failed after {max_retries} attempts. Last error: {last_error}" logger.warning(error_msg) if __event_emitter__: await __event_emitter__( { "type": "status", "data": { "description": "Hugging Face unavailable, trying Pollinations...", "done": False, }, } ) return await self._generate_with_fallback( prompt, image_format, seed, __event_emitter__ ) # Get image bytes image_bytes = response.content # Convert to base64 and send return await self._process_and_send_image( image_bytes, prompt, "FLUX.1-dev", __event_emitter__ ) except Exception as e: error_msg = f"Error with Hugging Face: {str(e)}" logger.warning(error_msg) if __event_emitter__: await __event_emitter__( { "type": "status", "data": { "description": "Hugging Face failed, trying Pollinations...", "done": False, }, } ) return await self._generate_with_fallback( prompt, image_format, seed, __event_emitter__ ) async def _generate_with_fallback( self, prompt: str, image_format: str, seed: Optional[int], __event_emitter__ ) -> str: """Generate image using Pollinations as fallback""" try: if __event_emitter__: await __event_emitter__( { "type": "status", "data": { "description": "Generating image with Pollinations", "done": False, }, } ) # Generate with Pollinations image_bytes = await self._generate_with_pollinations( prompt=prompt, image_format=image_format, seed=seed ) # Convert to base64 and send return await self._process_and_send_image( image_bytes, prompt, "Pollinations", __event_emitter__ ) except Exception as e: error_msg = ( f"Both Hugging Face and Pollinations failed. Last error: {str(e)}" ) logger.error(error_msg) if __event_emitter__: await __event_emitter__( { "type": "status", "data": {"description": error_msg, "done": True}, } ) return error_msg async def _process_and_send_image( self, image_bytes: bytes, prompt: str, source: str, __event_emitter__ ) -> str: """Process image bytes and send to UI""" # Convert to base64 image_base64 = base64.b64encode(image_bytes).decode("utf-8") image_data_url = f"data:image/jpeg;base64,{image_base64}" # Send completion status if __event_emitter__: await __event_emitter__( { "type": "status", "data": { "description": f"Image generated successfully using {source}", "done": True, }, } ) # Method 1: Send as files (OpenWebUI recommended format) await __event_emitter__( { "type": "files", "data": { "files": [ { "type": "image", "url": image_data_url, } ] }, } ) # Method 2: Also send as markdown message for compatibility await __event_emitter__( { "type": "message", "data": { "content": f"Generated image for prompt: '{prompt}' (using {source})\n\n", }, } ) # Log for debugging logger.info( f"Generated image with size: {len(image_bytes)} bytes using {source}" ) return f"Image successfully generated for prompt: '{prompt}' (using {source})"