""" title: SANA author: fakezeta description: Image generation with SANA based on Gradio Client. requirements: gradio_client version: 0.1.1 licence: MIT """ # based on FLUX.1-hfspace from sweetbrulee import asyncio import shutil import traceback import json from pathlib import Path from gradio_client import Client from pydantic import BaseModel, Field from open_webui.routers.images import IMAGE_CACHE_DIR from open_webui.main import WEBUI_URL from datetime import datetime ALLOWED_STYLES = [ "(No style)", "Cinematic", "Photographic", "Anime", "Manga", "Digital Art", "Pixel art", "Fantasy art", "Neonpunk", "3D Model", ] HEIGHT = 1024 WIDTH = 1024 FLOW_DPMS_GUIDANCE_SCALE = 4.5 FLOW_DPMS_INFERENCE_STEPS = 20 FLOW_DPMS_PAG_GUIDANCE_SCALE = 1 class Tools: class Valves(BaseModel): client_address: str = Field( default="https://sana.hanlab.ai/", description="Client API base address", ) api_name: str = Field(default="/run", description="Name of the API endpoint") def __init__(self): self.valves = self.Valves() async def generate_images( self, prompt: str, negative_prompt: str, style: str, __event_emitter__=None, ): """ Generate an image by providing arguments. Prompt must be in English and shall be enriched to describe the image to be generated. :param prompt: Prompt to use for image generation in English. It must be descriptive (~50 words); enrich brief prompts with details like mood, style, color, perspective, media type, etc. Use effective prompting techniques. :param negative_prompt: A comma-separated list of terms to exclude from the image, English. :param style: The image style requested by the user as in '(No style)', 'Cinematic', 'Photographic', 'Anime', 'Manga', 'Digital Art', 'Pixel art', 'Fantasy art', 'Neonpunk', '3D Model'. """ try: await __event_emitter__( { "type": "status", "data": { "description": "Generating image: " + prompt + " in style: " + style, "done": False, }, } ) image_source = await predict_async( prompt, negative_prompt, style, client_address=self.valves.client_address, api_name=self.valves.api_name, ) # copy the image folder tool_dir = IMAGE_CACHE_DIR.joinpath("./tools_sana/") tool_dir.mkdir(parents=True, exist_ok=True) image_destination = ( tool_dir / Path(datetime.now().strftime("%Y%m%d-%H%M%S")) # / Path(image_source[0]["image"]).parent.name / Path(image_source[0]["image"]).name ) # print("image_destination:", image_destination) image_destination.parent.mkdir(parents=True, exist_ok=True) shutil.copy2(image_source[0]["image"], image_destination) await __event_emitter__( { "type": "status", "data": { "description": "Generated image: " + prompt + " in style: " + style, "done": True, }, } ) await __event_emitter__( { "type": "message", "data": { "content": f"})" }, # the file route need to be relative to /app/backend/data, so that it can be accessed by the frontend (which will be appended to the fronend base url [WEBUI_URL]) } ) response = { "prompt": prompt, "image_url": f"{WEBUI_URL}/{image_destination.relative_to('/app/backend/data/')}", "model_used": "Sana", "style": style, "timestamp": datetime.now().isoformat(), "status": "success", } return json.dumps(response, indent=2) except Exception as e: await __event_emitter__( { "type": "status", "data": {"description": f"An error occured: {e}", "done": True}, } ) response = { "prompt": prompt, "style": style, "error message": traceback.format_exc, "timestamp": datetime.now().isoformat(), "status": "Error", } return json.dumps(response, indent=2) def predict_sync( prompt, negative_prompt, style, client_address, api_name, ): style = style if style in ALLOWED_STYLES else "(No style)" client = Client(client_address) result = client.predict( prompt=prompt, negative_prompt=negative_prompt, style=style, use_negative_prompt=(negative_prompt == ""), num_imgs=1, seed=0, height=HEIGHT, width=WIDTH, flow_dpms_guidance_scale=FLOW_DPMS_GUIDANCE_SCALE, flow_dpms_pag_guidance_scale=FLOW_DPMS_PAG_GUIDANCE_SCALE, flow_dpms_inference_steps=FLOW_DPMS_INFERENCE_STEPS, randomize_seed=True, api_name=api_name, ) return result[0] async def predict_async( prompt, negative_prompt, style, *, client_address, api_name, ): style = style if style in ALLOWED_STYLES else "(No style)" loop = asyncio.get_running_loop() image_source = await loop.run_in_executor( None, predict_sync, prompt, negative_prompt, style, client_address, api_name, ) return image_source