open-webui/backend/open_webui/routers/images.py

718 lines
25 KiB
Python
Raw Normal View History

2024-08-28 06:10:27 +08:00
import asyncio
2024-03-09 09:38:10 +08:00
import base64
import io
2024-03-09 09:38:10 +08:00
import json
import logging
2025-02-07 07:22:20 +08:00
import mimetypes
2024-08-21 06:35:42 +08:00
import re
2024-08-28 06:10:27 +08:00
from pathlib import Path
from typing import Optional
2024-08-21 06:35:42 +08:00
from urllib.parse import quote
2024-08-28 06:10:27 +08:00
import requests
2025-08-22 20:58:25 +08:00
from fastapi import (
APIRouter,
Depends,
HTTPException,
Request,
UploadFile,
)
2024-12-10 16:54:13 +08:00
from open_webui.config import CACHE_DIR
from open_webui.constants import ERROR_MESSAGES
from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS, SRC_LOG_LEVELS
2025-08-22 21:19:57 +08:00
from open_webui.routers.files import upload_file_handler
2024-12-09 08:01:56 +08:00
from open_webui.utils.auth import get_admin_user, get_verified_user
2024-12-12 10:53:38 +08:00
from open_webui.utils.images.comfyui import (
ComfyUIGenerateImageForm,
ComfyUIWorkflow,
comfyui_generate_image,
)
from pydantic import BaseModel
2024-03-09 09:38:10 +08:00
log = logging.getLogger(__name__)
log.setLevel(SRC_LOG_LEVELS["IMAGES"])
2024-03-09 09:38:10 +08:00
IMAGE_CACHE_DIR = CACHE_DIR / "image" / "generations"
2024-03-09 09:38:10 +08:00
IMAGE_CACHE_DIR.mkdir(parents=True, exist_ok=True)
2024-02-22 10:12:01 +08:00
2024-11-10 10:01:23 +08:00
2024-12-12 10:53:38 +08:00
router = APIRouter()
2024-08-03 04:35:02 +08:00
2024-12-12 10:53:38 +08:00
@router.get("/config")
async def get_config(request: Request, user=Depends(get_admin_user)):
return {
2024-12-13 12:22:17 +08:00
"enabled": request.app.state.config.ENABLE_IMAGE_GENERATION,
2024-12-13 12:24:36 +08:00
"engine": request.app.state.config.IMAGE_GENERATION_ENGINE,
2025-01-16 16:13:02 +08:00
"prompt_generation": request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION,
2024-08-21 06:35:42 +08:00
"openai": {
2024-12-13 12:24:36 +08:00
"OPENAI_API_BASE_URL": request.app.state.config.IMAGES_OPENAI_API_BASE_URL,
"OPENAI_API_KEY": request.app.state.config.IMAGES_OPENAI_API_KEY,
2024-08-21 06:35:42 +08:00
},
"automatic1111": {
2024-12-12 10:53:38 +08:00
"AUTOMATIC1111_BASE_URL": request.app.state.config.AUTOMATIC1111_BASE_URL,
"AUTOMATIC1111_API_AUTH": request.app.state.config.AUTOMATIC1111_API_AUTH,
"AUTOMATIC1111_CFG_SCALE": request.app.state.config.AUTOMATIC1111_CFG_SCALE,
"AUTOMATIC1111_SAMPLER": request.app.state.config.AUTOMATIC1111_SAMPLER,
"AUTOMATIC1111_SCHEDULER": request.app.state.config.AUTOMATIC1111_SCHEDULER,
2024-08-21 06:35:42 +08:00
},
"comfyui": {
2024-12-12 10:53:38 +08:00
"COMFYUI_BASE_URL": request.app.state.config.COMFYUI_BASE_URL,
2024-12-17 15:29:00 +08:00
"COMFYUI_API_KEY": request.app.state.config.COMFYUI_API_KEY,
2024-12-12 10:53:38 +08:00
"COMFYUI_WORKFLOW": request.app.state.config.COMFYUI_WORKFLOW,
"COMFYUI_WORKFLOW_NODES": request.app.state.config.COMFYUI_WORKFLOW_NODES,
2024-08-21 06:35:42 +08:00
},
"gemini": {
"GEMINI_API_BASE_URL": request.app.state.config.IMAGES_GEMINI_API_BASE_URL,
"GEMINI_API_KEY": request.app.state.config.IMAGES_GEMINI_API_KEY,
},
}
2024-02-22 10:12:01 +08:00
2024-08-21 06:35:42 +08:00
class OpenAIConfigForm(BaseModel):
OPENAI_API_BASE_URL: str
OPENAI_API_KEY: str
class Automatic1111ConfigForm(BaseModel):
AUTOMATIC1111_BASE_URL: str
AUTOMATIC1111_API_AUTH: str
2024-12-10 15:22:47 +08:00
AUTOMATIC1111_CFG_SCALE: Optional[str | float | int]
AUTOMATIC1111_SAMPLER: Optional[str]
AUTOMATIC1111_SCHEDULER: Optional[str]
2024-08-21 06:35:42 +08:00
class ComfyUIConfigForm(BaseModel):
COMFYUI_BASE_URL: str
2024-12-17 15:29:00 +08:00
COMFYUI_API_KEY: str
2024-08-21 06:35:42 +08:00
COMFYUI_WORKFLOW: str
COMFYUI_WORKFLOW_NODES: list[dict]
class GeminiConfigForm(BaseModel):
GEMINI_API_BASE_URL: str
GEMINI_API_KEY: str
2024-08-21 06:35:42 +08:00
class ConfigForm(BaseModel):
2024-03-09 09:38:10 +08:00
enabled: bool
2024-08-21 06:35:42 +08:00
engine: str
2025-01-16 16:13:02 +08:00
prompt_generation: bool
2024-08-21 06:35:42 +08:00
openai: OpenAIConfigForm
automatic1111: Automatic1111ConfigForm
comfyui: ComfyUIConfigForm
gemini: GeminiConfigForm
2024-03-09 09:38:10 +08:00
2024-12-12 10:53:38 +08:00
@router.post("/config/update")
async def update_config(
request: Request, form_data: ConfigForm, user=Depends(get_admin_user)
):
2024-12-13 12:24:36 +08:00
request.app.state.config.IMAGE_GENERATION_ENGINE = form_data.engine
2024-12-13 12:22:17 +08:00
request.app.state.config.ENABLE_IMAGE_GENERATION = form_data.enabled
2024-02-22 10:12:01 +08:00
2025-01-16 16:13:02 +08:00
request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION = (
form_data.prompt_generation
)
2024-12-13 12:24:36 +08:00
request.app.state.config.IMAGES_OPENAI_API_BASE_URL = (
form_data.openai.OPENAI_API_BASE_URL
)
request.app.state.config.IMAGES_OPENAI_API_KEY = form_data.openai.OPENAI_API_KEY
2024-02-22 10:12:01 +08:00
request.app.state.config.IMAGES_GEMINI_API_BASE_URL = (
form_data.gemini.GEMINI_API_BASE_URL
)
request.app.state.config.IMAGES_GEMINI_API_KEY = form_data.gemini.GEMINI_API_KEY
2024-12-12 10:53:38 +08:00
request.app.state.config.AUTOMATIC1111_BASE_URL = (
2024-08-21 06:35:42 +08:00
form_data.automatic1111.AUTOMATIC1111_BASE_URL
)
2024-12-12 10:53:38 +08:00
request.app.state.config.AUTOMATIC1111_API_AUTH = (
2024-08-21 06:35:42 +08:00
form_data.automatic1111.AUTOMATIC1111_API_AUTH
)
2024-09-13 12:49:23 +08:00
2024-12-12 10:53:38 +08:00
request.app.state.config.AUTOMATIC1111_CFG_SCALE = (
2024-09-13 12:49:23 +08:00
float(form_data.automatic1111.AUTOMATIC1111_CFG_SCALE)
2024-09-20 09:18:14 +08:00
if form_data.automatic1111.AUTOMATIC1111_CFG_SCALE
2024-09-13 12:49:23 +08:00
else None
)
2024-12-12 10:53:38 +08:00
request.app.state.config.AUTOMATIC1111_SAMPLER = (
2024-09-13 12:49:23 +08:00
form_data.automatic1111.AUTOMATIC1111_SAMPLER
2024-09-20 09:18:14 +08:00
if form_data.automatic1111.AUTOMATIC1111_SAMPLER
2024-09-13 12:49:23 +08:00
else None
)
2024-12-12 10:53:38 +08:00
request.app.state.config.AUTOMATIC1111_SCHEDULER = (
2024-09-13 12:49:23 +08:00
form_data.automatic1111.AUTOMATIC1111_SCHEDULER
2024-09-20 09:18:14 +08:00
if form_data.automatic1111.AUTOMATIC1111_SCHEDULER
2024-09-13 12:49:23 +08:00
else None
)
2024-02-22 10:12:01 +08:00
2024-12-12 10:53:38 +08:00
request.app.state.config.COMFYUI_BASE_URL = (
form_data.comfyui.COMFYUI_BASE_URL.strip("/")
)
2025-02-26 08:01:29 +08:00
request.app.state.config.COMFYUI_API_KEY = form_data.comfyui.COMFYUI_API_KEY
2024-12-12 10:53:38 +08:00
request.app.state.config.COMFYUI_WORKFLOW = form_data.comfyui.COMFYUI_WORKFLOW
request.app.state.config.COMFYUI_WORKFLOW_NODES = (
form_data.comfyui.COMFYUI_WORKFLOW_NODES
)
2024-02-22 10:12:01 +08:00
2024-03-24 06:38:59 +08:00
return {
2024-12-13 12:22:17 +08:00
"enabled": request.app.state.config.ENABLE_IMAGE_GENERATION,
2024-12-13 12:24:36 +08:00
"engine": request.app.state.config.IMAGE_GENERATION_ENGINE,
2025-01-16 16:13:02 +08:00
"prompt_generation": request.app.state.config.ENABLE_IMAGE_PROMPT_GENERATION,
2024-08-21 06:35:42 +08:00
"openai": {
2024-12-13 12:24:36 +08:00
"OPENAI_API_BASE_URL": request.app.state.config.IMAGES_OPENAI_API_BASE_URL,
"OPENAI_API_KEY": request.app.state.config.IMAGES_OPENAI_API_KEY,
2024-08-21 06:35:42 +08:00
},
"automatic1111": {
2024-12-12 10:53:38 +08:00
"AUTOMATIC1111_BASE_URL": request.app.state.config.AUTOMATIC1111_BASE_URL,
"AUTOMATIC1111_API_AUTH": request.app.state.config.AUTOMATIC1111_API_AUTH,
"AUTOMATIC1111_CFG_SCALE": request.app.state.config.AUTOMATIC1111_CFG_SCALE,
"AUTOMATIC1111_SAMPLER": request.app.state.config.AUTOMATIC1111_SAMPLER,
"AUTOMATIC1111_SCHEDULER": request.app.state.config.AUTOMATIC1111_SCHEDULER,
2024-08-21 06:35:42 +08:00
},
"comfyui": {
2024-12-12 10:53:38 +08:00
"COMFYUI_BASE_URL": request.app.state.config.COMFYUI_BASE_URL,
2024-12-17 15:29:00 +08:00
"COMFYUI_API_KEY": request.app.state.config.COMFYUI_API_KEY,
2024-12-12 10:53:38 +08:00
"COMFYUI_WORKFLOW": request.app.state.config.COMFYUI_WORKFLOW,
"COMFYUI_WORKFLOW_NODES": request.app.state.config.COMFYUI_WORKFLOW_NODES,
2024-08-21 06:35:42 +08:00
},
"gemini": {
"GEMINI_API_BASE_URL": request.app.state.config.IMAGES_GEMINI_API_BASE_URL,
"GEMINI_API_KEY": request.app.state.config.IMAGES_GEMINI_API_KEY,
},
2024-03-24 06:38:59 +08:00
}
2024-02-22 10:12:01 +08:00
2024-12-12 10:53:38 +08:00
def get_automatic1111_api_auth(request: Request):
if request.app.state.config.AUTOMATIC1111_API_AUTH is None:
2024-08-21 06:35:42 +08:00
return ""
2024-06-20 14:15:49 +08:00
else:
2024-12-12 10:53:38 +08:00
auth1111_byte_string = request.app.state.config.AUTOMATIC1111_API_AUTH.encode(
"utf-8"
)
2024-08-21 06:35:42 +08:00
auth1111_base64_encoded_bytes = base64.b64encode(auth1111_byte_string)
auth1111_base64_encoded_string = auth1111_base64_encoded_bytes.decode("utf-8")
return f"Basic {auth1111_base64_encoded_string}"
2024-02-22 10:12:01 +08:00
2024-12-12 10:53:38 +08:00
@router.get("/config/url/verify")
async def verify_url(request: Request, user=Depends(get_admin_user)):
2024-12-13 12:24:36 +08:00
if request.app.state.config.IMAGE_GENERATION_ENGINE == "automatic1111":
2024-08-21 23:27:39 +08:00
try:
r = requests.get(
2024-12-12 10:53:38 +08:00
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
headers={"authorization": get_automatic1111_api_auth(request)},
2024-08-21 23:27:39 +08:00
)
r.raise_for_status()
return True
2024-08-28 06:10:27 +08:00
except Exception:
2024-12-13 12:22:17 +08:00
request.app.state.config.ENABLE_IMAGE_GENERATION = False
2024-08-21 23:27:39 +08:00
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL)
2024-12-13 12:24:36 +08:00
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui":
2025-02-26 08:01:29 +08:00
headers = None
if request.app.state.config.COMFYUI_API_KEY:
headers = {
"Authorization": f"Bearer {request.app.state.config.COMFYUI_API_KEY}"
}
2024-08-21 23:27:39 +08:00
try:
2024-12-12 10:53:38 +08:00
r = requests.get(
2025-02-26 08:01:29 +08:00
url=f"{request.app.state.config.COMFYUI_BASE_URL}/object_info",
headers=headers,
2024-12-12 10:53:38 +08:00
)
2024-08-21 23:27:39 +08:00
r.raise_for_status()
return True
2024-08-28 06:10:27 +08:00
except Exception:
2024-12-13 12:22:17 +08:00
request.app.state.config.ENABLE_IMAGE_GENERATION = False
2024-08-21 23:27:39 +08:00
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.INVALID_URL)
else:
return True
2024-12-12 10:53:38 +08:00
def set_image_model(request: Request, model: str):
2024-09-04 21:25:31 +08:00
log.info(f"Setting image model to {model}")
2024-12-13 12:26:28 +08:00
request.app.state.config.IMAGE_GENERATION_MODEL = model
2024-12-13 12:24:36 +08:00
if request.app.state.config.IMAGE_GENERATION_ENGINE in ["", "automatic1111"]:
2024-12-26 04:26:13 +08:00
api_auth = get_automatic1111_api_auth(request)
2024-08-21 06:35:42 +08:00
r = requests.get(
2024-12-12 10:53:38 +08:00
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
2024-08-21 06:35:42 +08:00
headers={"authorization": api_auth},
)
options = r.json()
if model != options["sd_model_checkpoint"]:
options["sd_model_checkpoint"] = model
r = requests.post(
2024-12-12 10:53:38 +08:00
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
2024-08-21 06:35:42 +08:00
json=options,
headers={"authorization": api_auth},
)
2024-12-13 12:26:28 +08:00
return request.app.state.config.IMAGE_GENERATION_MODEL
2024-03-09 09:38:10 +08:00
2024-12-13 12:26:28 +08:00
def get_image_model(request):
2024-12-13 12:24:36 +08:00
if request.app.state.config.IMAGE_GENERATION_ENGINE == "openai":
2024-12-12 10:53:38 +08:00
return (
2024-12-13 12:26:28 +08:00
request.app.state.config.IMAGE_GENERATION_MODEL
if request.app.state.config.IMAGE_GENERATION_MODEL
2024-12-12 10:53:38 +08:00
else "dall-e-2"
)
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "gemini":
return (
request.app.state.config.IMAGE_GENERATION_MODEL
if request.app.state.config.IMAGE_GENERATION_MODEL
else "imagen-3.0-generate-002"
)
2024-12-13 12:24:36 +08:00
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui":
2024-12-13 12:26:28 +08:00
return (
request.app.state.config.IMAGE_GENERATION_MODEL
if request.app.state.config.IMAGE_GENERATION_MODEL
else ""
)
2024-12-12 10:53:38 +08:00
elif (
2024-12-13 12:24:36 +08:00
request.app.state.config.IMAGE_GENERATION_ENGINE == "automatic1111"
or request.app.state.config.IMAGE_GENERATION_ENGINE == ""
2024-12-12 10:53:38 +08:00
):
2024-08-21 06:35:42 +08:00
try:
r = requests.get(
2024-12-12 10:53:38 +08:00
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/options",
2024-12-26 15:34:54 +08:00
headers={"authorization": get_automatic1111_api_auth(request)},
2024-08-21 06:35:42 +08:00
)
options = r.json()
return options["sd_model_checkpoint"]
except Exception as e:
2024-12-13 12:22:17 +08:00
request.app.state.config.ENABLE_IMAGE_GENERATION = False
2024-08-21 06:35:42 +08:00
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
2024-03-09 09:38:10 +08:00
2024-08-21 06:35:42 +08:00
class ImageConfigForm(BaseModel):
2024-08-21 07:21:03 +08:00
MODEL: str
IMAGE_SIZE: str
IMAGE_STEPS: int
2024-03-09 09:38:10 +08:00
2024-04-23 19:14:31 +08:00
2024-12-12 10:53:38 +08:00
@router.get("/image/config")
async def get_image_config(request: Request, user=Depends(get_admin_user)):
2024-03-09 09:38:10 +08:00
return {
2024-12-13 12:26:28 +08:00
"MODEL": request.app.state.config.IMAGE_GENERATION_MODEL,
2024-12-12 10:53:38 +08:00
"IMAGE_SIZE": request.app.state.config.IMAGE_SIZE,
"IMAGE_STEPS": request.app.state.config.IMAGE_STEPS,
2024-03-09 09:38:10 +08:00
}
2024-12-12 10:53:38 +08:00
@router.post("/image/config/update")
async def update_image_config(
request: Request, form_data: ImageConfigForm, user=Depends(get_admin_user)
):
set_image_model(request, form_data.MODEL)
2024-02-23 11:32:36 +08:00
2025-06-28 19:21:20 +08:00
if form_data.IMAGE_SIZE == "auto" and form_data.MODEL != "gpt-image-1":
2025-06-23 04:55:29 +08:00
raise HTTPException(
status_code=400,
2025-06-28 19:21:20 +08:00
detail=ERROR_MESSAGES.INCORRECT_FORMAT(
" (auto is only allowed with gpt-image-1)."
),
2025-06-23 04:55:29 +08:00
)
2024-08-21 06:35:42 +08:00
pattern = r"^\d+x\d+$"
2025-06-23 04:55:29 +08:00
if form_data.IMAGE_SIZE == "auto" or re.match(pattern, form_data.IMAGE_SIZE):
2024-12-12 10:53:38 +08:00
request.app.state.config.IMAGE_SIZE = form_data.IMAGE_SIZE
2024-02-23 11:32:36 +08:00
else:
raise HTTPException(
status_code=400,
detail=ERROR_MESSAGES.INCORRECT_FORMAT(" (e.g., 512x512)."),
)
2024-02-25 10:08:35 +08:00
2024-08-21 07:21:03 +08:00
if form_data.IMAGE_STEPS >= 0:
2024-12-12 10:53:38 +08:00
request.app.state.config.IMAGE_STEPS = form_data.IMAGE_STEPS
else:
raise HTTPException(
status_code=400,
detail=ERROR_MESSAGES.INCORRECT_FORMAT(" (e.g., 50)."),
)
2024-02-23 11:32:36 +08:00
2024-08-21 06:35:42 +08:00
return {
2024-12-13 12:26:28 +08:00
"MODEL": request.app.state.config.IMAGE_GENERATION_MODEL,
2024-12-12 10:53:38 +08:00
"IMAGE_SIZE": request.app.state.config.IMAGE_SIZE,
"IMAGE_STEPS": request.app.state.config.IMAGE_STEPS,
2024-08-21 06:35:42 +08:00
}
2024-02-23 11:32:36 +08:00
2024-12-12 10:53:38 +08:00
@router.get("/models")
def get_models(request: Request, user=Depends(get_verified_user)):
2024-02-22 10:12:01 +08:00
try:
2024-12-13 12:24:36 +08:00
if request.app.state.config.IMAGE_GENERATION_ENGINE == "openai":
2024-03-09 09:38:10 +08:00
return [
{"id": "dall-e-2", "name": "DALL·E 2"},
{"id": "dall-e-3", "name": "DALL·E 3"},
{"id": "gpt-image-1", "name": "GPT-IMAGE 1"},
2024-03-09 09:38:10 +08:00
]
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "gemini":
return [
{"id": "imagen-3.0-generate-002", "name": "imagen-3.0 generate-002"},
]
2024-12-13 12:24:36 +08:00
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui":
2024-08-21 06:35:42 +08:00
# TODO - get models from comfyui
2024-12-18 05:51:29 +08:00
headers = {
"Authorization": f"Bearer {request.app.state.config.COMFYUI_API_KEY}"
}
r = requests.get(
url=f"{request.app.state.config.COMFYUI_BASE_URL}/object_info",
headers=headers,
)
2024-03-24 06:38:59 +08:00
info = r.json()
2024-12-12 10:53:38 +08:00
workflow = json.loads(request.app.state.config.COMFYUI_WORKFLOW)
model_node_id = None
2024-12-12 10:53:38 +08:00
for node in request.app.state.config.COMFYUI_WORKFLOW_NODES:
if node["type"] == "model":
2024-08-22 06:25:43 +08:00
if node["node_ids"]:
model_node_id = node["node_ids"][0]
break
if model_node_id:
model_list_key = None
log.info(workflow[model_node_id]["class_type"])
for key in info[workflow[model_node_id]["class_type"]]["input"][
"required"
]:
if "_name" in key:
model_list_key = key
break
if model_list_key:
return list(
map(
lambda model: {"id": model, "name": model},
info[workflow[model_node_id]["class_type"]]["input"][
"required"
][model_list_key][0],
)
)
else:
return list(
map(
lambda model: {"id": model, "name": model},
info["CheckpointLoaderSimple"]["input"]["required"][
"ckpt_name"
][0],
)
2024-03-24 06:38:59 +08:00
)
2024-08-21 06:35:42 +08:00
elif (
2024-12-13 12:24:36 +08:00
request.app.state.config.IMAGE_GENERATION_ENGINE == "automatic1111"
or request.app.state.config.IMAGE_GENERATION_ENGINE == ""
2024-08-21 06:35:42 +08:00
):
2024-03-09 09:38:10 +08:00
r = requests.get(
2024-12-12 10:53:38 +08:00
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/sd-models",
2024-12-26 15:34:54 +08:00
headers={"authorization": get_automatic1111_api_auth(request)},
2024-03-09 09:38:10 +08:00
)
models = r.json()
return list(
map(
lambda model: {"id": model["title"], "name": model["model_name"]},
models,
)
)
2024-02-22 10:12:01 +08:00
except Exception as e:
2024-12-13 12:22:17 +08:00
request.app.state.config.ENABLE_IMAGE_GENERATION = False
2024-02-25 10:08:35 +08:00
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(e))
2024-02-22 10:12:01 +08:00
class GenerateImageForm(BaseModel):
model: Optional[str] = None
prompt: str
2024-03-09 11:54:47 +08:00
size: Optional[str] = None
2024-08-21 06:35:42 +08:00
n: int = 1
2024-02-22 10:12:01 +08:00
negative_prompt: Optional[str] = None
def load_b64_image_data(b64_str):
2024-03-09 09:38:10 +08:00
try:
2024-05-03 06:59:45 +08:00
if "," in b64_str:
header, encoded = b64_str.split(",", 1)
mime_type = header.split(";")[0].lstrip("data:")
2024-05-03 06:59:45 +08:00
img_data = base64.b64decode(encoded)
else:
mime_type = "image/png"
2024-05-03 06:59:45 +08:00
img_data = base64.b64decode(b64_str)
return img_data, mime_type
2024-03-09 09:38:10 +08:00
except Exception as e:
log.exception(f"Error loading image data: {e}")
return None, None
2024-03-09 09:38:10 +08:00
def load_url_image_data(url, headers=None):
2024-03-24 08:01:13 +08:00
try:
2025-01-17 03:17:37 +08:00
if headers:
r = requests.get(url, headers=headers)
else:
r = requests.get(url)
2024-03-24 08:01:13 +08:00
r.raise_for_status()
if r.headers["content-type"].split("/")[0] == "image":
mime_type = r.headers["content-type"]
return r.content, mime_type
else:
2024-08-28 06:10:27 +08:00
log.error("Url does not point to an image.")
2024-05-03 06:54:31 +08:00
return None
2024-03-24 08:01:13 +08:00
except Exception as e:
log.exception(f"Error saving image: {e}")
2024-05-03 06:54:31 +08:00
return None
2024-03-24 08:01:13 +08:00
2025-08-22 21:19:57 +08:00
def upload_image(request, image_data, content_type, metadata, user):
2025-02-07 07:22:20 +08:00
image_format = mimetypes.guess_extension(content_type)
2025-02-07 07:12:04 +08:00
file = UploadFile(
file=io.BytesIO(image_data),
2025-02-12 13:13:42 +08:00
filename=f"generated-image{image_format}", # will be converted to a unique ID on upload_file
2025-02-07 07:12:04 +08:00
headers={
"content-type": content_type,
},
)
2025-08-22 21:19:57 +08:00
file_item = upload_file_handler(
2025-08-22 20:58:25 +08:00
request,
file=file,
metadata=metadata,
process=False,
user=user,
2025-08-20 04:36:13 +08:00
)
2025-02-07 07:12:04 +08:00
url = request.app.url_path_for("get_file_content_by_id", id=file_item.id)
return url
2024-12-12 10:53:38 +08:00
@router.post("/generations")
2024-07-15 22:25:00 +08:00
async def image_generations(
2024-12-12 10:53:38 +08:00
request: Request,
2024-06-25 23:01:05 +08:00
form_data: GenerateImageForm,
2024-06-28 02:29:59 +08:00
user=Depends(get_verified_user),
2024-02-22 10:12:01 +08:00
):
2025-06-23 04:55:29 +08:00
# if IMAGE_SIZE = 'auto', default WidthxHeight to the 512x512 default
# This is only relevant when the user has set IMAGE_SIZE to 'auto' with an
# image model other than gpt-image-1, which is warned about on settings save
size = "512x512"
2025-08-22 17:25:23 +08:00
if (
request.app.state.config.IMAGE_SIZE
and "x" in request.app.state.config.IMAGE_SIZE
):
size = request.app.state.config.IMAGE_SIZE
2025-08-22 17:25:23 +08:00
if form_data.size and "x" in form_data.size:
size = form_data.size
width, height = tuple(map(int, size.split("x")))
2024-03-24 08:01:13 +08:00
2024-03-09 11:54:47 +08:00
r = None
2024-02-22 10:36:40 +08:00
try:
2024-12-13 12:24:36 +08:00
if request.app.state.config.IMAGE_GENERATION_ENGINE == "openai":
2024-03-09 09:38:10 +08:00
headers = {}
2024-12-12 10:53:38 +08:00
headers["Authorization"] = (
2024-12-13 12:24:36 +08:00
f"Bearer {request.app.state.config.IMAGES_OPENAI_API_KEY}"
2024-12-12 10:53:38 +08:00
)
2024-03-09 09:38:10 +08:00
headers["Content-Type"] = "application/json"
2024-02-22 10:36:40 +08:00
2024-11-01 23:23:18 +08:00
if ENABLE_FORWARD_USER_INFO_HEADERS:
headers["X-OpenWebUI-User-Name"] = quote(user.name, safe=" ")
headers["X-OpenWebUI-User-Id"] = user.id
headers["X-OpenWebUI-User-Email"] = user.email
headers["X-OpenWebUI-User-Role"] = user.role
2024-03-09 09:38:10 +08:00
data = {
"model": (
2024-12-13 12:26:28 +08:00
request.app.state.config.IMAGE_GENERATION_MODEL
if request.app.state.config.IMAGE_GENERATION_MODEL != ""
else "dall-e-2"
),
2024-03-09 09:38:10 +08:00
"prompt": form_data.prompt,
"n": form_data.n,
"size": (
2024-12-12 10:53:38 +08:00
form_data.size
if form_data.size
else request.app.state.config.IMAGE_SIZE
),
2025-04-29 23:34:00 +08:00
**(
{}
2025-04-29 23:34:00 +08:00
if "gpt-image-1" in request.app.state.config.IMAGE_GENERATION_MODEL
else {"response_format": "b64_json"}
2025-04-29 23:34:00 +08:00
),
2024-03-09 09:38:10 +08:00
}
2024-03-13 04:35:30 +08:00
2024-09-20 09:16:08 +08:00
# Use asyncio.to_thread for the requests.post call
r = await asyncio.to_thread(
requests.post,
2024-12-13 12:24:36 +08:00
url=f"{request.app.state.config.IMAGES_OPENAI_API_BASE_URL}/images/generations",
2024-03-09 09:38:10 +08:00
json=data,
2025-02-07 06:37:18 +08:00
headers=headers,
2024-03-09 09:38:10 +08:00
)
2024-03-09 09:38:10 +08:00
r.raise_for_status()
res = r.json()
2024-02-22 10:12:01 +08:00
2024-03-09 09:38:10 +08:00
images = []
for image in res["data"]:
2025-03-11 23:40:31 +08:00
if image_url := image.get("url", None):
2025-03-11 23:39:15 +08:00
image_data, content_type = load_url_image_data(image_url, headers)
2025-03-04 12:07:59 +08:00
else:
image_data, content_type = load_b64_image_data(image["b64_json"])
2025-08-22 21:19:57 +08:00
url = upload_image(request, image_data, content_type, data, user)
images.append({"url": url})
2024-03-09 09:38:10 +08:00
return images
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "gemini":
headers = {}
headers["Content-Type"] = "application/json"
headers["x-goog-api-key"] = request.app.state.config.IMAGES_GEMINI_API_KEY
model = get_image_model(request)
data = {
"instances": {"prompt": form_data.prompt},
"parameters": {
"sampleCount": form_data.n,
"outputOptions": {"mimeType": "image/png"},
},
}
# Use asyncio.to_thread for the requests.post call
r = await asyncio.to_thread(
requests.post,
url=f"{request.app.state.config.IMAGES_GEMINI_API_BASE_URL}/models/{model}:predict",
json=data,
headers=headers,
)
r.raise_for_status()
res = r.json()
images = []
for image in res["predictions"]:
image_data, content_type = load_b64_image_data(
image["bytesBase64Encoded"]
)
2025-08-22 21:19:57 +08:00
url = upload_image(request, image_data, content_type, data, user)
images.append({"url": url})
return images
2024-12-13 12:24:36 +08:00
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui":
2024-03-24 08:01:13 +08:00
data = {
"prompt": form_data.prompt,
"width": width,
"height": height,
"n": form_data.n,
}
2024-12-12 10:53:38 +08:00
if request.app.state.config.IMAGE_STEPS is not None:
data["steps"] = request.app.state.config.IMAGE_STEPS
2024-03-24 08:01:13 +08:00
if form_data.negative_prompt is not None:
2024-03-24 08:01:13 +08:00
data["negative_prompt"] = form_data.negative_prompt
2024-08-21 20:49:17 +08:00
form_data = ComfyUIGenerateImageForm(
2024-08-21 20:49:54 +08:00
**{
2024-08-21 20:49:17 +08:00
"workflow": ComfyUIWorkflow(
2024-08-21 20:49:54 +08:00
**{
2024-12-12 10:53:38 +08:00
"workflow": request.app.state.config.COMFYUI_WORKFLOW,
"nodes": request.app.state.config.COMFYUI_WORKFLOW_NODES,
2024-08-21 20:49:17 +08:00
}
),
**data,
}
)
res = await comfyui_generate_image(
2024-12-13 12:26:28 +08:00
request.app.state.config.IMAGE_GENERATION_MODEL,
form_data,
2024-03-24 08:01:13 +08:00
user.id,
2024-12-12 10:53:38 +08:00
request.app.state.config.COMFYUI_BASE_URL,
2024-12-17 15:29:00 +08:00
request.app.state.config.COMFYUI_API_KEY,
2024-03-24 08:01:13 +08:00
)
log.debug(f"res: {res}")
2024-03-24 08:01:13 +08:00
images = []
for image in res["data"]:
2025-01-17 03:17:37 +08:00
headers = None
if request.app.state.config.COMFYUI_API_KEY:
headers = {
"Authorization": f"Bearer {request.app.state.config.COMFYUI_API_KEY}"
}
image_data, content_type = load_url_image_data(image["url"], headers)
2025-02-07 07:12:04 +08:00
url = upload_image(
request,
image_data,
content_type,
2025-05-27 06:48:22 +08:00
form_data.model_dump(exclude_none=True),
2025-02-07 07:12:04 +08:00
user,
)
images.append({"url": url})
2024-03-24 08:01:13 +08:00
return images
2024-08-21 06:35:42 +08:00
elif (
2024-12-13 12:24:36 +08:00
request.app.state.config.IMAGE_GENERATION_ENGINE == "automatic1111"
or request.app.state.config.IMAGE_GENERATION_ENGINE == ""
2024-08-21 06:35:42 +08:00
):
2024-03-09 09:38:10 +08:00
if form_data.model:
set_image_model(request, form_data.model)
2024-03-09 09:38:10 +08:00
data = {
"prompt": form_data.prompt,
"batch_size": form_data.n,
"width": width,
"height": height,
}
2024-12-12 10:53:38 +08:00
if request.app.state.config.IMAGE_STEPS is not None:
data["steps"] = request.app.state.config.IMAGE_STEPS
2024-03-09 09:38:10 +08:00
if form_data.negative_prompt is not None:
2024-03-09 09:38:10 +08:00
data["negative_prompt"] = form_data.negative_prompt
2024-09-13 12:49:23 +08:00
2024-12-12 10:53:38 +08:00
if request.app.state.config.AUTOMATIC1111_CFG_SCALE:
data["cfg_scale"] = request.app.state.config.AUTOMATIC1111_CFG_SCALE
2024-12-12 10:53:38 +08:00
if request.app.state.config.AUTOMATIC1111_SAMPLER:
data["sampler_name"] = request.app.state.config.AUTOMATIC1111_SAMPLER
2024-12-12 10:53:38 +08:00
if request.app.state.config.AUTOMATIC1111_SCHEDULER:
data["scheduler"] = request.app.state.config.AUTOMATIC1111_SCHEDULER
2024-03-09 09:38:10 +08:00
2024-08-23 22:51:34 +08:00
# Use asyncio.to_thread for the requests.post call
r = await asyncio.to_thread(
requests.post,
2024-12-12 10:53:38 +08:00
url=f"{request.app.state.config.AUTOMATIC1111_BASE_URL}/sdapi/v1/txt2img",
2024-03-09 09:38:10 +08:00
json=data,
2024-12-26 15:34:54 +08:00
headers={"authorization": get_automatic1111_api_auth(request)},
2024-03-09 09:38:10 +08:00
)
res = r.json()
log.debug(f"res: {res}")
2024-03-09 09:38:10 +08:00
images = []
for image in res["images"]:
image_data, content_type = load_b64_image_data(image)
2025-02-07 07:12:04 +08:00
url = upload_image(
request,
image_data,
content_type,
2025-05-27 06:48:22 +08:00
{**data, "info": res["info"]},
2025-02-07 07:12:04 +08:00
user,
)
images.append({"url": url})
2024-03-09 09:38:10 +08:00
return images
2024-02-22 10:36:40 +08:00
except Exception as e:
2024-03-13 04:35:30 +08:00
error = e
if r != None:
data = r.json()
if "error" in data:
error = data["error"]["message"]
raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error))