OSSForks/ollama-webui
1
0
Fork 0
mirror of https://github.com/zebrajr/ollama-webui.git synced 2026-01-15 12:15:13 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

refac

Browse Source
This commit is contained in:
Timothy Jaeryang Baek
2025-12-22 00:39:05 +04:00
parent 01e88c6ac2
commit 446cc0ac60
3 changed files with 75 additions and 17 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -17,6 +17,8 @@ from fastapi.responses import FileResponse
from open_webui.config import CACHE_DIR
from open_webui.constants import ERROR_MESSAGES
from open_webui.env import ENABLE_FORWARD_USER_INFO_HEADERS
from open_webui.models.chats import Chats
from open_webui.routers.files import upload_file_handler, get_file_content_by_id
from open_webui.utils.auth import get_admin_user, get_verified_user
from open_webui.utils.headers import include_user_info_headers
@@ -510,14 +512,29 @@ def upload_image(request, image_data, content_type, metadata, user):
process=False,
user=user,
)
if file_item and file_item.id:
# If chat_id and message_id are provided in metadata, link the file to the chat message
chat_id = metadata.get("chat_id")
message_id = metadata.get("message_id")
if chat_id and message_id:
Chats.insert_chat_files(
chat_id=chat_id,
message_id=message_id,
file_ids=[file_item.id],
user_id=user.id,
)
url = request.app.url_path_for("get_file_content_by_id", id=file_item.id)
return url
return file_item, url
@router.post("/generations")
async def image_generations(
request: Request,
form_data: CreateImageForm,
metadata: Optional[dict] = None,
user=Depends(get_verified_user),
):
# if IMAGE_SIZE = 'auto', default WidthxHeight to the 512x512 default
@@ -535,6 +552,9 @@ async def image_generations(
size = form_data.size
width, height = tuple(map(int, size.split("x")))
metadata = metadata or {}
model = get_image_model(request)
r = None
@@ -595,7 +615,9 @@ async def image_generations(
else:
image_data, content_type = get_image_data(image["b64_json"])
url = upload_image(request, image_data, content_type, data, user)
_, url = upload_image(
request, image_data, content_type, {**data, **metadata}, user
)
images.append({"url": url})
return images
@@ -645,7 +667,9 @@ async def image_generations(
image_data, content_type = get_image_data(
image["bytesBase64Encoded"]
)
url = upload_image(request, image_data, content_type, data, user)
_, url = upload_image(
request, image_data, content_type, {**data, **metadata}, user
)
images.append({"url": url})
elif model.endswith(":generateContent"):
for image in res["candidates"]:
@@ -654,8 +678,12 @@ async def image_generations(
image_data, content_type = get_image_data(
part["inlineData"]["data"]
)
url = upload_image(
request, image_data, content_type, data, user
_, url = upload_image(
request,
image_data,
content_type,
{**data, **metadata},
user,
)
images.append({"url": url})
@@ -705,11 +733,11 @@ async def image_generations(
}
image_data, content_type = get_image_data(image["url"], headers)
url = upload_image(
_, url = upload_image(
request,
image_data,
content_type,
form_data.model_dump(exclude_none=True),
{**form_data.model_dump(exclude_none=True), **metadata},
user,
)
images.append({"url": url})
@@ -752,11 +780,11 @@ async def image_generations(
for image in res["images"]:
image_data, content_type = get_image_data(image)
url = upload_image(
_, url = upload_image(
request,
image_data,
content_type,
{**data, "info": res["info"]},
{**data, "info": res["info"], **metadata},
user,
)
images.append({"url": url})
@@ -783,10 +811,13 @@ class EditImageForm(BaseModel):
async def image_edits(
request: Request,
form_data: EditImageForm,
metadata: Optional[dict] = None,
user=Depends(get_verified_user),
):
size = None
width, height = None, None
metadata = metadata or {}
if (
request.app.state.config.IMAGE_EDIT_SIZE
and "x" in request.app.state.config.IMAGE_EDIT_SIZE
@@ -904,7 +935,9 @@ async def image_edits(
else:
image_data, content_type = get_image_data(image["b64_json"])
url = upload_image(request, image_data, content_type, data, user)
_, url = upload_image(
request, image_data, content_type, {**data, **metadata}, user
)
images.append({"url": url})
return images
@@ -957,8 +990,12 @@ async def image_edits(
image_data, content_type = get_image_data(
part["inlineData"]["data"]
)
url = upload_image(
request, image_data, content_type, data, user
_, url = upload_image(
request,
image_data,
content_type,
{**data, **metadata},
user,
)
images.append({"url": url})
@@ -1035,11 +1072,11 @@ async def image_edits(
}
image_data, content_type = get_image_data(image_url, headers)
url = upload_image(
_, url = upload_image(
request,
image_data,
content_type,
form_data.model_dump(exclude_none=True),
{**form_data.model_dump(exclude_none=True), **metadata},
user,
)
images.append({"url": url})

5
backend/open_webui/utils/files.py
View File

@@ -14,6 +14,8 @@ from typing import Optional
from pathlib import Path
from open_webui.storage.provider import Storage
from open_webui.models.chats import Chats
from open_webui.models.files import Files
from open_webui.routers.files import upload_file_handler
@@ -65,13 +67,14 @@ def get_image_url_from_base64(request, base64_image_string, metadata, user):
# Extract base64 image data from the line
image_data, content_type = get_image_data(base64_image_string)
if image_data is not None:
image_url = upload_image(
_, image_url = upload_image(
request,
image_data,
content_type,
metadata,
user,
)
return image_url
return None

22
backend/open_webui/utils/middleware.py
View File

@@ -799,6 +799,10 @@ async def chat_image_generation_handler(
images = await image_edits(
request=request,
form_data=EditImageForm(**{"prompt": prompt, "image": input_images}),
metadata={
"chat_id": metadata.get("chat_id", None),
"message_id": metadata.get("message_id", None),
},
user=user,
)
@@ -824,7 +828,7 @@ async def chat_image_generation_handler(
}
)
system_message_content = "<context>The requested image has been created and is now being shown to the user. Let them know that it has been generated.</context>"
system_message_content = "<context>The requested image has been edited and created and is now being shown to the user. Let them know that it has been generated.</context>"
except Exception as e:
log.debug(e)
@@ -883,6 +887,10 @@ async def chat_image_generation_handler(
images = await image_generations(
request=request,
form_data=CreateImageForm(**{"prompt": prompt}),
metadata={
"chat_id": metadata.get("chat_id", None),
"message_id": metadata.get("message_id", None),
},
user=user,
)
@@ -2737,7 +2745,17 @@ async def process_chat_response(
if ENABLE_CHAT_RESPONSE_BASE64_IMAGE_URL_CONVERSION:
value = convert_markdown_base64_images(
request, value, metadata, user
request,
value,
{
"chat_id": metadata.get(
"chat_id", None
),
"message_id": metadata.get(
"message_id", None
),
},
user,
)
content = f"{content}{value}"