""" title: Hugging Face API - Flux Pro Image Generator author: Alick version: 0.0.0 license: MIT description: Open WebUI Function Pipe - 使用 Hugging Face 文生图 API 生成图片，支持提示词里自定义宽高 """ import requests import uuid import os import re from pydantic import BaseModel, Field from typing import Union, Iterator, Generator from open_webui.config import CACHE_DIR class Pipe: class Valves(BaseModel): HF_API_KEY: str = Field( default="", description="HuggingFace API key for accessing the serverless endpoints", ) HF_API_URL: str = Field( default="https://api-inference.huggingface.co/models/stabilityai/stable-diffusion-3.5-large-turbo", description="HuggingFace API URL for accessing the Text-to-Image model.", ) def __init__(self): self.id = "huggingface_image_generator" self.type = "function" self.name = "Hugging Face Image Generator" self.valves = self.Valves() def pipe(self, body: dict) -> Union[str, Generator, Iterator]: messages = body.get("messages", []) prompt = body.get("prompt", "").strip() if not prompt: for m in reversed(messages): if m.get("role") == "user": prompt = m.get("content", "").strip() break if not self.valves.HF_API_KEY: return ( "Error: 尚未配置 Hugging Face API Key，请在 valves 中设置 HF_API_KEY。" ) formats = { "default": (1024, 1024), "square": (1024, 1024), "landscape": (1024, 768), "landscape_large": (1440, 1024), "portrait": (768, 1024), "portrait_large": (1024, 1440), } image_format = body.get("image_format", "default") if image_format not in formats: return "Error: 无效的 image_format，合法值：" + ", ".join(formats.keys()) width, height = formats[image_format] match = re.search(r"(\d+)\s*[x×\*]\s*(\d+)", prompt) if match: w, h = map(int, match.groups()) width, height = w, h headers = { "Authorization": f"Bearer {self.valves.HF_API_KEY}", "Content-Type": "application/json", } payload = { "inputs": prompt, "parameters": { "width": width, "height": height, }, } try: response = requests.post( self.valves.HF_API_URL, headers=headers, json=payload, timeout=600, ) if response.status_code != 200: return ( f"Error: 模型推理请求失败。\n" f"HTTP状态码: {response.status_code}\n" f"响应信息: {response.text}" ) image_content = response.content directory = os.path.join(CACHE_DIR, "image", "generations") os.makedirs(directory, exist_ok=True) filename = f"{uuid.uuid4()}.png" save_path = os.path.join(directory, filename) with open(save_path, "wb") as f: f.write(image_content) image_url = f"/cache/image/generations/{filename}" result_md = ( f"**已为你的提示词生成图片**：\n\n" f"" ) return result_md except requests.RequestException as e: return f"Error: 网络异常或连接失败。\n{str(e)}" except Exception as e: return f"Error: 未知异常。\n{str(e)}"