""" title: imaGE(Gen & Edit) author: Artyom Gromov version: 1.0.0 license: MIT License Copyright (c) 2025 Artyom Gromov Permission is hereby granted, free of charge, to any person obtaining a copy of this software and associated documentation files (the "Software"), to deal in the Software without restriction, including without limitation the rights to use, copy, modify, merge, publish, distribute, sublicense, and/or sell copies of the Software, and to permit persons to whom the Software is furnished to do so, subject to the following conditions: The above copyright notice and this permission notice shall be included in all copies or substantial portions of the Software. THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE. description: Tool for image generation and editting based on gpt-image-1. Allows to bring changes to generated images. """ from pydantic import BaseModel, Field from fastapi import Request import re import httpx import contextlib from open_webui.routers.images import load_b64_image_data, upload_image from open_webui.models.users import Users from open_webui.models.files import Files class Tools: def __init__(self): self.valves = self.Valves() self.func = self.AdditionalFunctions(self.valves) class Valves(BaseModel): API_BASE_URL: str = Field( default="", description="URL to gpt-image-1 model (should end with images endpoint)", ) API_KEY: str = Field( default="", description="API key", ) IMAGE_SIZE: str = Field( default="1024x1024", description="Image size: 1024x1024, 1536x1024, 1024x1536", ) IMAGE_QUALITY: str = Field( default="medium", description="Image quality: high, medium, low", ) class AdditionalFunctions: def __init__(self, valves): self.valves = valves async def _emit_status(self, event_emitter, description, done: bool = False): if event_emitter: await event_emitter( { "type": "status", "data": {"description": description, "done": done}, } ) def _get_image_params(self, prompt): return { "model": "gpt-image-1", "prompt": prompt, "n": 1, "size": self.valves.IMAGE_SIZE, "quality": self.valves.IMAGE_QUALITY, } async def _upload_b64_image( self, b64_image, request: Request, metadata: dict, user_id: int ): image_data, content_type = load_b64_image_data(b64_image) return upload_image( request=request, image_data=image_data, content_type=content_type, metadata=metadata, user=Users.get_user_by_id(user_id), ) async def generate_image( self, prompt: str, __request__: Request, __user__: dict, __event_emitter__=None, ) -> str: """ Call this tool if user asks to generate image :param promt: prompt to use for image generation """ await self.func._emit_status(__event_emitter__, "Generating an image...") try: headers = { "Content-Type": "application/json", "Authorization": f"Bearer {self.valves.API_KEY}", } req_data = self.func._get_image_params(prompt) async with httpx.AsyncClient() as client: resp = await client.post( f"{self.valves.API_BASE_URL}/generations", headers=headers, json=req_data, timeout=None, ) resp.raise_for_status() data = resp.json() image_b64 = data.get("data", [{}])[0].get("b64_json") if not image_b64: raise Exception("No image received from the API.") url = await self.func._upload_b64_image( image_b64, __request__, req_data, __user__["id"] ) await self.func._emit_status( __event_emitter__, "Image generated.", done=True ) await __event_emitter__( { "type": "message", "data": {"content": f""}, } ) return """ The image generation completed successfully! Note that the generated image ALREADY HAS been automatically sent and displayed to the user by the tool. You don't need to do anything to show the image to the user. When you answer now - simply tell the user the image was successfully generated. Answer with text only, NO IMAGES, NO ASPECT RATIO, NO IMAGE PROMPT. Just tell the user that the image was successfully generated. At the end of your answer you might ask the user if they want anything changed to the image. Should they later come back and ask for changes - just generate a new image with potentially modified parameters based on their feedback. Keep your answer short and concise. """ except Exception as e: await self.func._emit_status( __event_emitter__, f"An error occurred: {e}", done=True ) return f"Tell the user: {e}" async def edit_image( self, prompt: str, __request__: Request, __user__: dict, __messages__: dict, __event_emitter__=None, ) -> str: """ Call this tool if user asks to edit image :param prompt: prompt to use for image edit """ await self.func._emit_status(__event_emitter__, "Editing an image...") try: images = [] for message in __messages__[::-1]: match = re.search(r"!\[.*?\]\((.*?)\)", message["content"]) if match: images.append(match.group(1).split("/")[-2]) break if not images: raise Exception("No image found in recent messages.") headers = { "Authorization": f"Bearer {self.valves.API_KEY}", } files = [] req_data = self.func._get_image_params(prompt) with contextlib.ExitStack() as stack: for image_id in images: img_file_path = Files.get_file_by_id(image_id).path file_obj = stack.enter_context(open(img_file_path, "rb")) files.append( ("image[]", (f"{image_id}.png", file_obj, "image/png")) ) async with httpx.AsyncClient() as client: resp = await client.post( f"{self.valves.API_BASE_URL}/edits", headers=headers, files=files, data=req_data, timeout=None, ) resp.raise_for_status() data = resp.json() image_b64 = data.get("data", [{}])[0].get("b64_json") if not image_b64: raise Exception("No image received from the API.") url = await self.func._upload_b64_image( image_b64, __request__, req_data, __user__["id"] ) await self.func._emit_status( __event_emitter__, "Edited an image", done=True ) await __event_emitter__( { "type": "message", "data": {"content": f""}, } ) return """ The image generation completed successfully! Note that the generated image ALREADY HAS been automatically sent and displayed to the user by the tool. You don't need to do anything to show the image to the user. When you answer now - simply tell the user the image was successfully generated. Answer with text only, NO IMAGES, NO ASPECT RATIO, NO IMAGE PROMPT. Just tell the user that the image was successfully generated. At the end of your answer you might ask the user if they want anything changed to the image. Should they later come back and ask for changes - just generate a new image with potentially modified parameters based on their feedback. Keep your answer short and concise. """ except Exception as e: await self.func._emit_status( __event_emitter__, f"An error occurred: {e}", done=True ) return f"Tell the user: {e}"