From 159578dfd41a5fc68ac6c7cfdcf0db837600fe35 Mon Sep 17 00:00:00 2001 From: Rodrigo Agundez Date: Thu, 6 Feb 2025 17:59:59 +0900 Subject: [PATCH 1/8] Enable usage of the DB to store generated images --- backend/open_webui/routers/images.py | 63 ++++++++++------------------ 1 file changed, 22 insertions(+), 41 deletions(-) diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index 7afd9d106..68465e191 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -7,26 +7,21 @@ import re import uuid from pathlib import Path from typing import Optional +import io import requests - - -from fastapi import Depends, FastAPI, HTTPException, Request, APIRouter -from fastapi.middleware.cors import CORSMiddleware -from pydantic import BaseModel - - +from fastapi import APIRouter, Depends, UploadFile, HTTPException, Request from open_webui.config import CACHE_DIR from open_webui.constants import ERROR_MESSAGES -from open_webui.env import ENV, SRC_LOG_LEVELS, ENABLE_FORWARD_USER_INFO_HEADERS - +from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS, SRC_LOG_LEVELS +from open_webui.routers.files import upload_file from open_webui.utils.auth import get_admin_user, get_verified_user from open_webui.utils.images.comfyui import ( ComfyUIGenerateImageForm, ComfyUIWorkflow, comfyui_generate_image, ) - +from pydantic import BaseModel log = logging.getLogger(__name__) log.setLevel(SRC_LOG_LEVELS["IMAGES"]) @@ -39,7 +34,7 @@ router = APIRouter() @router.get("/config") -async def get_config(request: Request, user=Depends(get_admin_user)): +async def get_def(request: Request, user=Depends(get_admin_user)): return { "enabled": request.app.state.config.ENABLE_IMAGE_GENERATION, "engine": request.app.state.config.IMAGE_GENERATION_ENGINE, @@ -271,7 +266,6 @@ async def get_image_config(request: Request, user=Depends(get_admin_user)): async def update_image_config( request: Request, form_data: ImageConfigForm, user=Depends(get_admin_user) ): - set_image_model(request, form_data.MODEL) pattern = r"^\d+x\d+$" @@ -383,35 +377,18 @@ class GenerateImageForm(BaseModel): negative_prompt: Optional[str] = None -def save_b64_image(b64_str): +def load_b64_image_data(b64_str): try: - image_id = str(uuid.uuid4()) - if "," in b64_str: header, encoded = b64_str.split(",", 1) mime_type = header.split(";")[0] - img_data = base64.b64decode(encoded) - image_format = mimetypes.guess_extension(mime_type) - - image_filename = f"{image_id}{image_format}" - file_path = IMAGE_CACHE_DIR / f"{image_filename}" - with open(file_path, "wb") as f: - f.write(img_data) - return image_filename else: - image_filename = f"{image_id}.png" - file_path = IMAGE_CACHE_DIR.joinpath(image_filename) - + mime_type = "image/png" img_data = base64.b64decode(b64_str) - - # Write the image data to a file - with open(file_path, "wb") as f: - f.write(img_data) - return image_filename - + return img_data, mime_type except Exception as e: - log.exception(f"Error saving image: {e}") + log.exception(f"Error loading image data: {e}") return None @@ -500,13 +477,17 @@ async def image_generations( images = [] for image in res["data"]: - image_filename = save_b64_image(image["b64_json"]) - images.append({"url": f"/cache/image/generations/{image_filename}"}) - file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json") - - with open(file_body_path, "w") as f: - json.dump(data, f) - + image_data, content_type = load_b64_image_data(image["b64_json"]) + file = UploadFile( + 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}) return images elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui": @@ -618,4 +599,4 @@ async def image_generations( data = r.json() if "error" in data: error = data["error"]["message"] - raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error)) + raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error)) \ No newline at end of file From 8d43fdadc17c6db83df4f474a505fae1f509e6ea Mon Sep 17 00:00:00 2001 From: Rodrigo Agundez Date: Thu, 6 Feb 2025 18:24:57 +0900 Subject: [PATCH 2/8] Add functionality in other image generation types --- backend/open_webui/routers/images.py | 64 +++++++++++++++++----------- 1 file changed, 40 insertions(+), 24 deletions(-) diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index 68465e191..a26c06c61 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -1,5 +1,6 @@ import asyncio import base64 +import io import json import logging import mimetypes @@ -7,10 +8,9 @@ import re import uuid from pathlib import Path from typing import Optional -import io 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.constants import ERROR_MESSAGES 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 -def save_url_image(url, headers=None): - image_id = str(uuid.uuid4()) +def load_url_image_data(url, headers=None): try: if headers: r = requests.get(url, headers=headers) @@ -403,18 +402,7 @@ def save_url_image(url, headers=None): r.raise_for_status() if r.headers["content-type"].split("/")[0] == "image": mime_type = r.headers["content-type"] - image_format = mimetypes.guess_extension(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 + return r.content, mime_type else: log.error("Url does not point to an image.") return None @@ -486,8 +474,14 @@ async def image_generations( }, ) 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}) + 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 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}" } - image_filename = save_url_image(image["url"], headers) - images.append({"url": f"/cache/image/generations/{image_filename}"}) - file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json") + image_data, content_type = load_url_image_data(image["url"], headers) + file = UploadFile( + 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: json.dump(form_data.model_dump(exclude_none=True), f) @@ -585,9 +590,20 @@ async def image_generations( images = [] for image in res["images"]: - image_filename = save_b64_image(image) - images.append({"url": f"/cache/image/generations/{image_filename}"}) - file_body_path = IMAGE_CACHE_DIR.joinpath(f"{image_filename}.json") + image_data, content_type = load_b64_image_data(image) + file = UploadFile( + 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: json.dump({**data, "info": res["info"]}, f) @@ -599,4 +615,4 @@ async def image_generations( data = r.json() if "error" in data: error = data["error"]["message"] - raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error)) \ No newline at end of file + raise HTTPException(status_code=400, detail=ERROR_MESSAGES.DEFAULT(error)) From ac3338265d224bbff2ed80310a7bd1bd4d763bed Mon Sep 17 00:00:00 2001 From: Rodrigo Agundez Date: Fri, 7 Feb 2025 07:30:58 +0900 Subject: [PATCH 3/8] Set get_config as the name of the function --- backend/open_webui/routers/images.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index a26c06c61..a43714956 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -34,7 +34,7 @@ router = APIRouter() @router.get("/config") -async def get_def(request: Request, user=Depends(get_admin_user)): +async def get_config(request: Request, user=Depends(get_admin_user)): return { "enabled": request.app.state.config.ENABLE_IMAGE_GENERATION, "engine": request.app.state.config.IMAGE_GENERATION_ENGINE, @@ -456,7 +456,7 @@ async def image_generations( requests.post, url=f"{request.app.state.config.IMAGES_OPENAI_API_BASE_URL}/images/generations", json=data, - headers=headers, + headers=headers, ) r.raise_for_status() From 7e97e9dcc925dbe969190a510cf1b34bf9fdd1d7 Mon Sep 17 00:00:00 2001 From: Rodrigo Agundez Date: Fri, 7 Feb 2025 07:37:18 +0900 Subject: [PATCH 4/8] Improve style --- backend/open_webui/routers/images.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index a43714956..df69485c2 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -3,9 +3,7 @@ import base64 import io import json import logging -import mimetypes import re -import uuid from pathlib import Path from typing import Optional @@ -456,7 +454,7 @@ async def image_generations( requests.post, url=f"{request.app.state.config.IMAGES_OPENAI_API_BASE_URL}/images/generations", json=data, - headers=headers, + headers=headers, ) r.raise_for_status() From 4974c9cbb08eb9efd998c5eeb5ba75af827306b3 Mon Sep 17 00:00:00 2001 From: Rodrigo Agundez Date: Fri, 7 Feb 2025 08:12:04 +0900 Subject: [PATCH 5/8] Refactor upload function --- backend/open_webui/routers/images.py | 77 +++++++++++----------------- 1 file changed, 31 insertions(+), 46 deletions(-) diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index df69485c2..2d608c133 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -410,6 +410,24 @@ def load_url_image_data(url, headers=None): return None +def upload_image(request, data, image_data, content_type, user): + file = UploadFile( + 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) + file_body_path = IMAGE_CACHE_DIR.joinpath(f"{file_item.id}.json") + + with open(file_body_path, "w") as f: + json.dump(data, f) + + url = request.app.url_path_for("get_file_content_by_id", id=file_item.id) + return url + + @router.post("/generations") async def image_generations( request: Request, @@ -464,22 +482,8 @@ async def image_generations( for image in res["data"]: image_data, content_type = load_b64_image_data(image["b64_json"]) - file = UploadFile( - 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 - ) + url = upload_image(request, data, image_data, content_type, user) 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 elif request.app.state.config.IMAGE_GENERATION_ENGINE == "comfyui": @@ -526,24 +530,14 @@ async def image_generations( } image_data, content_type = load_url_image_data(image["url"], headers) - file = UploadFile( - 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 + url = upload_image( + request, + form_data.model_dump(exclude_none=True), + image_data, + content_type, + user, ) 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(form_data.model_dump(exclude_none=True), f) - - log.debug(f"images: {images}") return images elif ( request.app.state.config.IMAGE_GENERATION_ENGINE == "automatic1111" @@ -589,23 +583,14 @@ async def image_generations( for image in res["images"]: image_data, content_type = load_b64_image_data(image) - file = UploadFile( - 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 + url = upload_image( + request, + {**data, "info": res["info"]}, + image_data, + content_type, + user, ) 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, "info": res["info"]}, f) - return images except Exception as e: error = e From 312f273a1bf60f66089fbcd9f7ad225466419d30 Mon Sep 17 00:00:00 2001 From: Rodrigo Agundez Date: Fri, 7 Feb 2025 08:22:20 +0900 Subject: [PATCH 6/8] Add extension to image filename --- backend/open_webui/routers/images.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index 2d608c133..da4bf8a17 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -3,6 +3,7 @@ import base64 import io import json import logging +import mimetypes import re from pathlib import Path from typing import Optional @@ -411,15 +412,16 @@ def load_url_image_data(url, headers=None): def upload_image(request, data, image_data, content_type, user): + image_format = mimetypes.guess_extension(content_type) file = UploadFile( file=io.BytesIO(image_data), - filename="image", # will be converted to a unique ID on upload_file + filename=f"generated{image_format}", # will be converted to a unique ID on upload_file headers={ "content-type": content_type, }, ) file_item = upload_file(request, file, user) - file_body_path = IMAGE_CACHE_DIR.joinpath(f"{file_item.id}.json") + file_body_path = IMAGE_CACHE_DIR.joinpath(f"{file_item.filename}.json") with open(file_body_path, "w") as f: json.dump(data, f) From ffb9e739753f1d096531d5084fc524e642744266 Mon Sep 17 00:00:00 2001 From: Rodrigo Agundez Date: Fri, 7 Feb 2025 08:32:06 +0900 Subject: [PATCH 7/8] Save image metadata to DB --- backend/open_webui/routers/files.py | 36 +++++++++++++--------------- backend/open_webui/routers/images.py | 9 ++----- 2 files changed, 18 insertions(+), 27 deletions(-) diff --git a/backend/open_webui/routers/files.py b/backend/open_webui/routers/files.py index 7160c2e86..7cf7a942f 100644 --- a/backend/open_webui/routers/files.py +++ b/backend/open_webui/routers/files.py @@ -3,30 +3,22 @@ import os import uuid from pathlib import Path from typing import Optional -from pydantic import BaseModel -import mimetypes from urllib.parse import quote -from open_webui.storage.provider import Storage - +from fastapi import APIRouter, Depends, File, HTTPException, Request, UploadFile, status +from fastapi.responses import FileResponse, StreamingResponse +from open_webui.constants import ERROR_MESSAGES +from open_webui.env import SRC_LOG_LEVELS from open_webui.models.files import ( FileForm, FileModel, FileModelResponse, Files, ) -from open_webui.routers.retrieval import process_file, ProcessFileForm - -from open_webui.config import UPLOAD_DIR -from open_webui.env import SRC_LOG_LEVELS -from open_webui.constants import ERROR_MESSAGES - - -from fastapi import APIRouter, Depends, File, HTTPException, UploadFile, status, Request -from fastapi.responses import FileResponse, StreamingResponse - - +from open_webui.routers.retrieval import ProcessFileForm, process_file +from open_webui.storage.provider import Storage from open_webui.utils.auth import get_admin_user, get_verified_user +from pydantic import BaseModel log = logging.getLogger(__name__) log.setLevel(SRC_LOG_LEVELS["MODELS"]) @@ -41,7 +33,10 @@ router = APIRouter() @router.post("/", response_model=FileModelResponse) def upload_file( - request: Request, file: UploadFile = File(...), user=Depends(get_verified_user) + request: Request, + file: UploadFile = File(...), + user=Depends(get_verified_user), + file_metadata: dict = {}, ): log.info(f"file.content_type: {file.content_type}") try: @@ -61,6 +56,7 @@ def upload_file( "id": id, "filename": name, "path": file_path, + "data": file_metadata, "meta": { "name": name, "content_type": file.content_type, @@ -126,7 +122,7 @@ async def delete_all_files(user=Depends(get_admin_user)): Storage.delete_all_files() except Exception as e: log.exception(e) - log.error(f"Error deleting files") + log.error("Error deleting files") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT("Error deleting files"), @@ -248,7 +244,7 @@ async def get_file_content_by_id(id: str, user=Depends(get_verified_user)): ) except Exception as e: log.exception(e) - log.error(f"Error getting file content") + log.error("Error getting file content") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT("Error getting file content"), @@ -279,7 +275,7 @@ async def get_html_file_content_by_id(id: str, user=Depends(get_verified_user)): ) except Exception as e: log.exception(e) - log.error(f"Error getting file content") + log.error("Error getting file content") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT("Error getting file content"), @@ -355,7 +351,7 @@ async def delete_file_by_id(id: str, user=Depends(get_verified_user)): Storage.delete_file(file.path) except Exception as e: log.exception(e) - log.error(f"Error deleting files") + log.error("Error deleting files") raise HTTPException( status_code=status.HTTP_400_BAD_REQUEST, detail=ERROR_MESSAGES.DEFAULT("Error deleting files"), diff --git a/backend/open_webui/routers/images.py b/backend/open_webui/routers/images.py index da4bf8a17..1dcaca866 100644 --- a/backend/open_webui/routers/images.py +++ b/backend/open_webui/routers/images.py @@ -411,7 +411,7 @@ def load_url_image_data(url, headers=None): return None -def upload_image(request, data, image_data, content_type, user): +def upload_image(request, image_metadata, image_data, content_type, user): image_format = mimetypes.guess_extension(content_type) file = UploadFile( file=io.BytesIO(image_data), @@ -420,12 +420,7 @@ def upload_image(request, data, image_data, content_type, user): "content-type": content_type, }, ) - file_item = upload_file(request, file, user) - file_body_path = IMAGE_CACHE_DIR.joinpath(f"{file_item.filename}.json") - - with open(file_body_path, "w") as f: - json.dump(data, f) - + file_item = upload_file(request, file, user, file_metadata=image_metadata) url = request.app.url_path_for("get_file_content_by_id", id=file_item.id) return url From 2b05c9d944757600289c38b3d2fde88ebf19997e Mon Sep 17 00:00:00 2001 From: Rodrigo Agundez Date: Sat, 8 Feb 2025 19:35:27 +0900 Subject: [PATCH 8/8] Move file metadata into the [meta][data] key --- backend/open_webui/routers/files.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/backend/open_webui/routers/files.py b/backend/open_webui/routers/files.py index 7cf7a942f..051321257 100644 --- a/backend/open_webui/routers/files.py +++ b/backend/open_webui/routers/files.py @@ -56,11 +56,11 @@ def upload_file( "id": id, "filename": name, "path": file_path, - "data": file_metadata, "meta": { "name": name, "content_type": file.content_type, "size": len(contents), + "data": file_metadata, }, } ),