""" title: Image Gen Filter v2 author: Eidenz version: 0.3 required_open_webui_version: 0.5.4 """ from pydantic import BaseModel, Field from fastapi import Request from typing import Callable, Awaitable, Any, Optional, Literal import json import re from open_webui.main import Users from open_webui.routers.tools import Tools from open_webui.main import Models from open_webui.utils.chat import generate_chat_completion from open_webui.utils.misc import get_last_user_message from open_webui.routers.images import image_generations, GenerateImageForm class Filter: class Valves(BaseModel): template: str = Field( default="""You are an image prompt generator. When it makes sense to generate an image based on the provided context, write [IMAGE_GEN: keywords list] to tell the generator what to generate. Give as much details as possible, use only keywords. Example: [IMAGE_GEN: solo, smile, looking at viewer, forest, sunny, sunlight] Only return the IMAGE_GEN answer, nothing else. Only write if it makes sense to generate an image, if not just return '[]'. """ ) preprompt: str = Field( default="" ) status: bool = Field(default=False) def __init__(self): self.valves = self.Valves() async def inlet( self, body: dict, __request__: Request, __event_emitter__: Callable[[Any], Awaitable[None]], __user__: Optional[dict] = None, __model__: Optional[dict] = None, ) -> dict: messages = body["messages"] user_message = get_last_user_message(messages) system_prompt = self.valves.template prompt = ( "History:\n" + "\n".join( [ f"{message['role'].upper()}: \"\"\"{message['content']}\"\"\"" for message in messages[::-1][:4] ] ) + f"\nQuery: {user_message}" ) payload = { "model": body["model"], "messages": [ {"role": "system", "content": system_prompt}, {"role": "user", "content": prompt}, ], "stream": False, } if self.valves.status: await __event_emitter__( { "type": "status", "data": { "description": f"Looking for image need...", "done": False, }, } ) try: user = Users.get_user_by_id(__user__["id"]) response = await generate_chat_completion( __request__, form_data=payload, user=user ) content = response["choices"][0]["message"]["content"] if content is not None: print(f"content: {content}") content = content.replace("'", '"') image_prompt_match = re.search(r"\[IMAGE_GEN: (.*?)\]", content) if image_prompt_match: image_prompt = image_prompt_match.group(1) await __event_emitter__( { "type": "status", "data": { "description": "Generatin an image", "done": False, }, } ) images = await image_generations( __request__, GenerateImageForm( **{ "prompt": self.valves.preprompt + "," + image_prompt } ), Users.get_user_by_id(__user__["id"]), ) if self.valves.status: await __event_emitter__( { "type": "status", "data": {"description": image_prompt, "done": True}, } ) else: await __event_emitter__( { "type": "status", "data": { "description": "Image generated", "done": True, }, } ) for image in images: await __event_emitter__( { "type": "message", "data": { "content": f"" }, } ) else: if self.valves.status: await __event_emitter__( { "type": "status", "data": { "description": f"No image to generate.", "done": True, }, } ) else: if self.valves.status: await __event_emitter__( { "type": "status", "data": { "description": f"No image to generate.", "done": True, }, } ) except Exception as e: print(f"Error processing request: {e}") if self.valves.status: await __event_emitter__( { "type": "status", "data": { "description": f"Error processing request: {e}", "done": True, }, } ) return body