""" title: Venice Image Gen author: Mafrans author_url: https://github.com/Mafrans funding_url: https://github.com/Mafrans version: 0.1 required_open_webui_version: 0.5.3 """ import os import requests from datetime import datetime from typing import Callable from fastapi import Request from pydantic import BaseModel, Field from open_webui.routers.images import image_generations, GenerateImageForm from open_webui.models.users import Users STYLES = ["3D Model", "Analog Film", "Anime", "Cinematic", "Comic Book", "Craft Clay", "Digital Art", "Enhance", "Fantasy Art", "Isometric Style", "Line Art", "Lowpoly", "Neon Punk", "Origami", "Photographic", "Pixel Art", "Texture", "Advertising", "Food Photography", "Real Estate", "Abstract", "Cubist", "Graffiti", "Hyperrealism", "Impressionist", "Pointillism", "Pop Art", "Psychedelic", "Renaissance", "Steampunk", "Surrealist", "Typography", "Watercolor", "Fighting Game", "GTA", "Super Mario", "Minecraft", "Pokemon", "Retro Arcade", "Retro Game", "RPG Fantasy Game", "Strategy Game", "Street Fighter", "Legend of Zelda", "Architectural", "Disco", "Dreamscape", "Dystopian", "Fairy Tale", "Gothic", "Grunge", "Horror", "Minimalist", "Monochrome", "Nautical", "Space", "Stained Glass", "Techwear Fashion", "Tribal", "Zentangle", "Collage", "Flat Papercut", "Kirigami", "Paper Mache", "Paper Quilling", "Papercut Collage", "Papercut Shadow Box", "Stacked Papercut", "Thick Layered Papercut", "Alien", "Film Noir", "HDR", "Long Exposure", "Neon Noir", "Silhouette", "Tilt-Shift"] class Tools: def __init__(self): self.valves = self.Valves() pass class Valves(BaseModel): api_key: str = Field( description="Venice API Key", default="" ) model: str = Field( description="Image generation model", default="fluently-xl", example=[ "fluently-xl", "flux-dev", "flux-dev-uncensored", "pony-realism", "stable-diffusion-3.5", "lustify-sdxl" ] ), style: str = Field( description="Image style", default="Photorealistic", example=STYLES ) size: int = Field( description="Image size", default=1024, example=[512, 1024, 2048] ) async def generate_image( self, prompt: str, __request__: Request, __user__: dict, __event_emitter__=None ) -> str: """ Generate an image given a prompt :param prompt: prompt to use for image generation """ await __event_emitter__( { "type": "status", "data": {"description": "Generating an image", "done": False}, } ) try: style = self.valves.style for s in STYLES: if s.lower() in prompt.lower(): style = s res = await requests.post( "https://api.venice.ai/api/v1/image/generate", data={ "model": self.valves.model, "prompt": prompt, "style_preset": style, "height": self.valves.size, "width": self.valves.size, } ) data = res.json() await __event_emitter__( { "type": "status", "data": {"description": "Generated an image", "done": True}, } ) for image in data.images: await __event_emitter__( { "type": "message", "data": {"content": f""}, } ) return f"Notify the user that the image has been successfully generated" except Exception as e: await __event_emitter__( { "type": "status", "data": {"description": f"An error occured: {e}", "done": True}, } ) return f"Tell the user: {e}"