""" title: Novita.ai Image Generation Tool for OpenWebUI author: antonymajar version: 0.1 license: MIT description: Novita.ai API integration for image generation within OpenWebUI. API docs: https://api.novita.ai/ """ import time import requests import os from typing import Literal # Get API key from environment variable NOVITA_API_KEY = os.getenv("NOVITA_API_KEY", "") class Tools: def __init__(self): pass async def create_novita_image( self, prompt: str, image_format: str, model: str = "sd_xl_base_1.0.safetensors", steps: int = 20, __event_emitter__=None, ) -> str: """ This Tool creates images using the Novita.ai API with various models including SDXL. :param prompt: The prompt to generate the image :param image_format: Either 'default' for square image, 'landscape' for landscape format, or 'portrait' for portrait format :param model: The model to use (defaults to SDXL base 1.0) :param steps: Number of denoising steps (1-100) """ try: # Show initial status await __event_emitter__( { "type": "status", "data": {"description": "Creating image with Novita.ai...", "done": False}, } ) # Generate the image task_id = send_image_generation_request( prompt=prompt, image_format=image_format, model=model, steps=steps ) print(f"Novita image task ID: {task_id}") # Poll for results image_url = poll_result(task_id) print(f"Novita image URL: {image_url}") # Update status to complete await __event_emitter__( { "type": "status", "data": {"description": "Image Generated Successfully", "done": True}, } ) # Display the image await __event_emitter__( { "type": "message", "data": {"content": f" \n"}, } ) # Show the prompt used await __event_emitter__( { "type": "message", "data": {"content": f"Prompt: {prompt} \n"}, } ) return f"Only tell the user that the image was successfully generated with the format '{image_format}'. Do not show any links." except Exception as e: # Handle errors await __event_emitter__( { "type": "status", "data": {"description": f"An error occurred: {e}", "done": True}, } ) return f"Tell the user that an error occurred and the image generation was not successful. Reason: {e}" # Image format dimensions FORMATS = { "default": (1024, 1024), "square": (1024, 1024), "landscape": (1024, 768), "landscape_large": (1440, 1024), "portrait": (768, 1024), "portrait_large": (1024, 1440), } def send_image_generation_request( prompt: str, image_format: str = "default", model: str = "sd_xl_base_1.0.safetensors", steps: int = 20 ) -> str: """ Send a request to the Novita.ai API to generate an image. :param prompt: The prompt to generate the image from :param image_format: Format/dimensions of the image :param model: Model to use for generation :param steps: Number of denoising steps :return: Task ID for polling results """ width, height = FORMATS[image_format] payload = { "extra": { "response_image_type": "jpeg" }, "request": { "prompt": prompt, "model_name": model, "negative_prompt": "", "width": width, "height": height, "image_num": 1, "steps": steps, "seed": -1, "sampler_name": "Euler a", "guidance_scale": 7.5 } } response = requests.post( "https://api.novita.ai/v3/async/txt2img", headers={ "Authorization": f"Bearer {NOVITA_API_KEY}", "Content-Type": "application/json", }, json=payload ) response.raise_for_status() return response.json()["task_id"] def poll_result(task_id: str, max_retries: int = 30, delay: float = 2.0) -> str: """ Poll the Novita.ai API for the result of the image generation request. :param task_id: The task ID returned from the generation request :param max_retries: Maximum number of polling attempts :param delay: Delay between polling attempts in seconds :return: URL of the generated image :raises RuntimeError: If the status is not successful """ headers = { "Authorization": f"Bearer {NOVITA_API_KEY}", "Content-Type": "application/json", } for _ in range(max_retries): result = get_result(task_id, headers) status = result["task"]["status"] if status == "TASK_STATUS_SUCCEED": return result["images"][0]["image_url"] elif status == "TASK_STATUS_FAILED": raise RuntimeError(f"Image generation failed: {result['task'].get('reason', 'Unknown error')}") time.sleep(delay) raise RuntimeError("Image generation timed out") def get_result(task_id: str, headers: dict) -> dict: """ Get the current status/result of an image generation task. :param task_id: The task ID to check :param headers: Request headers with authorization :return: Response data from the API """ response = requests.get( f"https://api.novita.ai/v3/async/task-result", headers=headers, params={"task_id": task_id} ) response.raise_for_status() return response.json()