""" title: Generate images with Automatic1111 (RunPod) author: Kemon description: Generate images with Automatic1111 hosted on RunPod (Serverless, Image, SD Automatic1111) required_open_webui_version: 0.6.0 requirements: version: 0.0.1 licence: MIT """ import os import requests from pydantic import BaseModel, Field from typing import Callable import base64 import uuid import json import time def get_current_dir(): name = __name__.split("_") dir_type = name[0] dir_name = __name__.replace(f"{dir_type}_", "") if dir_type == "tool" or dir_type == "function": dir_type = f"{dir_type}s" self_dir = os.path.join("/app/backend/data", "cache", dir_type, dir_name) return self_dir async def emitter(description, done=False, __event_emitter__=None): await __event_emitter__( { "type": "status", "data": { "description": description, "done": done, }, } ) class Tools: class Valves(BaseModel): API_URL: str = Field( default="https://api.runpod.ai/v2/xxxxxxxxxxxxxx/run", description="RunPod Automatic1111 apiurl.", ) APIKEY: str = Field( default="rpa_xxxxx", description="RunPod apikey", ) def __init__(self): self.valves = self.Valves() self.working_dir = get_current_dir() pass async def generate_image(self, prompt, __event_emitter__=None) -> str: """ Generate image based on the prompt. :param prompt: Image description :return: Return success or failure. """ # API Endpoint and Authentication API_URL = self.valves.API_URL.replace("/run", "") APIKEY = self.valves.APIKEY RUN_ENDPOINT = f"{API_URL}/run" STATUS_ENDPOINT = f"{API_URL}/status/" # Request data payload = {"input": {"prompt": prompt}} # Headers headers = { "Content-Type": "application/json", "Authorization": f"Bearer {APIKEY}", } await emitter("Submitting job...", False, __event_emitter__) response = requests.post( RUN_ENDPOINT, headers=headers, data=json.dumps(payload) ) result = response.json() job_id = result.get("id") status = result.get("status") while status in ["IN_PROGRESS", "IN_QUEUE"]: await emitter(f"Status: {status}. Waiting...", False, __event_emitter__) time.sleep(1) status_response = requests.post( f"{STATUS_ENDPOINT}{job_id}", headers=headers ) status_data = status_response.json() status = status_data.get("status") await emitter(f"Final Status: {status}", False, __event_emitter__) if status == "COMPLETED" and "output" in status_data: output = status_data.get("output", {}) images = output.get("images", []) if images: image_data = images[0] image_bytes = base64.b64decode(image_data) FILE_PATH = self.working_dir + "/" + str(uuid.uuid4()) + ".jpg" FILE_URL = FILE_PATH.replace("/app/backend/data", "") with open(FILE_PATH, "wb") as f: f.write(image_bytes) await emitter(f"Image generated", True, __event_emitter__) await __event_emitter__( { "type": "message", "data": {"content": f""}, } ) return "Tell the user that the image is generated." else: await emitter("No images returned in output.", True, __event_emitter__) return "No images returned in output." else: await emitter("Job did not complete successfully.", True, __event_emitter__) return "Job did not complete successfully."