OpenWebUI/open-webui
1
0
Fork 0
mirror of https://github.com/open-webui/open-webui synced 2025-06-26 18:26:48 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Add functionality in other image generation types

Browse Source
This commit is contained in:
Rodrigo Agundez 2025-02-06 18:24:57 +09:00
parent 159578dfd4
commit 8d43fdadc1
1 changed files with 40 additions and 24 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

62
backend/open_webui/routers/images.py
View File

@ -1,5 +1,6 @@
import asyncio import asyncio
import base64 import base64
import io
import json import json
import logging import logging
import mimetypes import mimetypes
@ -7,10 +8,9 @@ import re
import uuid import uuid
from pathlib import Path from pathlib import Path
from typing import Optional from typing import Optional
import io
import requests import requests
from fastapi import APIRouter, Depends, UploadFile, HTTPException, Request from fastapi import APIRouter, Depends, HTTPException, Request, UploadFile
from open_webui.config import CACHE_DIR from open_webui.config import CACHE_DIR
from open_webui.constants import ERROR_MESSAGES from open_webui.constants import ERROR_MESSAGES
from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS, SRC_LOG_LEVELS from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS, SRC_LOG_LEVELS
@ -392,8 +392,7 @@ def load_b64_image_data(b64_str):
return None return None
def save_url_image(url, headers=None): def load_url_image_data(url, headers=None):
image_id = str(uuid.uuid4())
try: try:
if headers: if headers:
r = requests.get(url, headers=headers) r = requests.get(url, headers=headers)
@ -403,18 +402,7 @@ def save_url_image(url, headers=None):
r.raise_for_status() r.raise_for_status()
if r.headers["content-type"].split("/")[0] == "image": if r.headers["content-type"].split("/")[0] == "image":
mime_type = r.headers["content-type"] mime_type = r.headers["content-type"]
image_format = mimetypes.guess_extension(mime_type) return r.content, mime_type
if not image_format:
raise ValueError("Could not determine image type from MIME type")
image_filename = f"{image_id}{image_format}"
file_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}")
with open(file_path, "wb") as image_file:
for chunk in r.iter_content(chunk_size=8192):
image_file.write(chunk)
return image_filename
else: else:
log.error("Url does not point to an image.") log.error("Url does not point to an image.")
return None return None
@ -486,8 +474,14 @@ async def image_generations(
}, },
) )
file_item = upload_file(request, file, user) file_item = upload_file(request, file, user)
url = request.app.url_path_for("get_file_content_by_id", id=file_item.id) url = request.app.url_path_for(
"get_file_content_by_id", id=file_item.id
)
images.append({"url": url}) images.append({"url": url})
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{file_item.id}.json")
with open(file_body_path, "w") as f:
json.dump(data, f)
return images return images
elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui": elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui":
@ -533,9 +527,20 @@ async def image_generations(
"Authorization": f"Bearer {request.app.state.config.COMFYUI_API_KEY}" "Authorization": f"Bearer {request.app.state.config.COMFYUI_API_KEY}"
} }
image_filename = save_url_image(image["url"], headers) image_data, content_type = load_url_image_data(image["url"], headers)
images.append({"url": f"/cache/image/generations/{image_filename}"}) file = UploadFile(
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json") file=io.BytesIO(image_data),
filename="image", # will be converted to a unique ID on upload_file
headers={
"content-type": content_type,
},
)
file_item = upload_file(request, file, user)
url = request.app.url_path_for(
"get_file_content_by_id", id=file_item.id
)
images.append({"url": url})
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{file_item.id}.json")
with open(file_body_path, "w") as f: with open(file_body_path, "w") as f:
json.dump(form_data.model_dump(exclude_none=True), f) json.dump(form_data.model_dump(exclude_none=True), f)
@ -585,9 +590,20 @@ async def image_generations(
images = [] images = []
for image in res["images"]: for image in res["images"]:
image_filename = save_b64_image(image) image_data, content_type = load_b64_image_data(image)
images.append({"url": f"/cache/image/generations/{image_filename}"}) file = UploadFile(
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json") file=io.BytesIO(image_data),
filename="image", # will be converted to a unique ID on upload_file
headers={
"content-type": content_type,
},
)
file_item = upload_file(request, file, user)
url = request.app.url_path_for(
"get_file_content_by_id", id=file_item.id
)
images.append({"url": url})
file_body_path = IMAGE_CACHE_DIR.joinpath(f"{file_item.id}.json")
with open(file_body_path, "w") as f: with open(file_body_path, "w") as f:
json.dump({**data, "info": res["info"]}, f) json.dump({**data, "info": res["info"]}, f)