""" title: Stable Diffusion Image Generator description: Generate images with Stable Diffusion hosted on Replicate author: Dean Tjodimedjo version: 1.0.1 license: MIT """ import requests # Used for making HTTP requests to the Replicate API import os # Used to access environment variables for API tokens from pydantic import ( BaseModel, Field, conint, ) # Pydantic for data validation and schema generation import random # Imported, though not currently used for explicit random seed generation (Replicate handles it when seed is omitted) class Tools: # The 'Valves' class defines the configurable parameters for this tool in Open WebUI. # These parameters appear in the tool's settings/gear icon and allow users to # customize behavior without modifying the code directly. class Valves(BaseModel): REPLICATE_API_TOKEN: str = Field( default="", description="Get your token from https://replicate.com/account", title="Replicate API Token", ) # Reordered valves to match the requested display order from Replicate.com's UI. # This ensures a familiar layout for users. HEIGHT: int = Field( default=512, description="Height of the output image.", title="Height", json_schema_extra={ # 'enum' creates a dropdown list in the Open WebUI for specific height options. "enum": [ 128, 192, 256, 320, 384, 448, 512, 576, 640, 704, 768, 832, 896, 960, 1024, ] }, ) WIDTH: int = Field( default=512, description="Width of the output image.", title="Width", json_schema_extra={ # 'enum' ensures a dropdown for width, similar to height. "enum": [ 128, 192, 256, 320, 384, 448, 512, 576, 640, 704, 768, 832, 896, 960, 1024, ] }, ) NEGATIVE_PROMPT: str = Field( default="0", description="A prompt to guide the image generation away from specified themes or objects. Leave 0 to disable.", title="Negative Prompt", ) NUM_OUTPUTS: conint(ge=1, le=4) = Field( default=1, description="Number of images to generate (minimum: 1, maximum: 4).", title="Number of Outputs", json_schema_extra={ # 'minimum' and 'maximum' define a slider or numerical input range in Open WebUI. "minimum": 1, "maximum": 4, # Replicate's typical limit for batch image generation. }, ) NUM_INFERENCE_STEPS: conint(ge=1, le=500) = Field( default=50, description="Number of denoising steps (minimum: 1, maximum: 500).", title="Number of Inference Steps", json_schema_extra={ "minimum": 1, "maximum": 500, }, ) GUIDANCE_SCALE: float = Field( default=7.5, description="Higher values encourage the model to stick to the prompt. Lower values give the model more freedom (minimum: 1, maximum: 20).", title="Guidance Scale", json_schema_extra={ "minimum": 1.0, "maximum": 20.0, }, ) SCHEDULER: str = Field( default="K_EULER", description="The scheduler to be used for the diffusion process.", title="Scheduler", json_schema_extra={ # 'enum' provides a dropdown for selecting the diffusion scheduler. "enum": [ "DDIM", "K_EULER", "DPMSolverMultistep", "K_EULER_ANCESTRAL", "PNDM", "KLMS", ] }, ) SEED: int = Field( default=0, # Default to 0; our logic will omit 'seed' from payload if 0. description="Random seed. Leave 0 for a random seed. For reproducible results, use a specific number (e.g., your birthday digits).", title="Seed", json_schema_extra={ "minimum": 0, }, ) # The __init__ method is called when the tool is initialized by Open WebUI. def __init__(self): # Initialize the 'valves' instance with values from environment variables # if they exist, otherwise use the defaults defined in the Valves class. # The order here mirrors the class definition for clarity and consistency. self.valves = self.Valves( REPLICATE_API_TOKEN=os.getenv("REPLICATE_API_TOKEN", ""), HEIGHT=int(os.getenv("SD_HEIGHT", "512")), WIDTH=int(os.getenv("SD_WIDTH", "512")), NEGATIVE_PROMPT=os.getenv("SD_NEGATIVE_PROMPT", ""), NUM_OUTPUTS=int(os.getenv("SD_NUM_OUTPUTS", "1")), NUM_INFERENCE_STEPS=int(os.getenv("SD_NUM_INFERENCE_STEPS", "50")), GUIDANCE_SCALE=float(os.getenv("SD_GUIDANCE_SCALE", "7.5")), SCHEDULER=os.getenv("SD_SCHEDULER", "K_EULER"), SEED=int(os.getenv("SD_SEED", "0")), ) # This is the main function that Open WebUI calls when the user triggers the tool. # 'image_prompt' is the text input from the chat. # '__event_emitter__' is a special Open WebUI argument for sending messages back to the UI. async def generate_image( self, image_prompt: str, __event_emitter__=None, ) -> str: try: # Inform the user that image generation has started. await __event_emitter__( { "type": "status", "data": {"description": "Generating image...", "done": False}, } ) # Essential check: Ensure the Replicate API token is set. if not self.valves.REPLICATE_API_TOKEN: raise ValueError( "API token missing. Please set it in the tool's Valves." ) # Define the specific Stable Diffusion model version on Replicate. version = "ac732df83cea7fff18b8472768c88ad041fa750ff7682a21affe81863cbe77e4" url = "https://api.replicate.com/v1/predictions" headers = { "Authorization": f"Bearer {self.valves.REPLICATE_API_TOKEN}", "Content-Type": "application/json", "Prefer": "wait", # Instruct Replicate to wait for the prediction to complete. } # Construct the input payload for the Replicate API. # Values are sourced directly from the tool's configured valves. input_data = { "prompt": image_prompt, "width": self.valves.WIDTH, "height": self.valves.HEIGHT, "num_outputs": self.valves.NUM_OUTPUTS, "guidance_scale": self.valves.GUIDANCE_SCALE, "num_inference_steps": self.valves.NUM_INFERENCE_STEPS, "scheduler": self.valves.SCHEDULER, } # Conditional inclusion of 'negative_prompt': # Only send to Replicate if the valve actually contains text. # An empty string in the valve means the user wants to disable it. if self.valves.NEGATIVE_PROMPT: input_data["negative_prompt"] = self.valves.NEGATIVE_PROMPT # Conditional inclusion of 'seed': # If the seed valve is 0, we omit it from the payload, telling Replicate # to use a random seed. Otherwise, the specified seed is sent. if self.valves.SEED != 0: input_data["seed"] = self.valves.SEED # Final payload structure for the Replicate API request. payload = { "version": version, "input": input_data, } # Send the request to Replicate. response = requests.post(url, headers=headers, json=payload) response.raise_for_status() # Raise an exception for HTTP errors (4xx or 5xx). result = response.json() # Parse the JSON response. # Check if the 'output' field exists and contains URLs. if "output" not in result or not result["output"]: raise ValueError("Failed to generate image. No output URL found.") # Prepare Markdown for displaying potentially multiple generated images. generated_images_markdown = "" for i, image_url in enumerate(result["output"]): generated_images_markdown += ( f"\n\n" ) # Update status to indicate completion. await __event_emitter__( {"type": "status", "data": {"description": "Done!", "done": True}} ) # Send the generated image(s) back to the Open WebUI chat. # The 'content' now only contains the image Markdown, without the prompt text. await __event_emitter__( { "type": "message", "data": {"content": generated_images_markdown}, } ) return "✅ Image(s) successfully generated and displayed above!" # --- Error Handling --- except requests.exceptions.HTTPError as e: error_msg = f"API Error: {str(e)}" # Provide specific guidance for common API errors. if e.response.status_code == 401: error_msg += "\n🔑 **Fix:** Check your Replicate API token at https://replicate.com/account" elif e.response.status_code == 422: # Attempt to extract more specific error details from Replicate's response. try: replicate_error_details = e.response.json() if "detail" in replicate_error_details: error_msg += ( f"\nReplicate details: {replicate_error_details['detail']}" ) except: pass # Ignore if parsing Replicate's error JSON fails. return f"Tell the user: {error_msg}" except Exception as e: # Catch any other unexpected errors during the process. return f"Tell the user: An unexpected error occurred: {str(e)}"