Use standardized procedure to save and load images

This commit is contained in:
Rodrigo Agundez 2025-06-08 16:26:58 +08:00
parent 168997a58e
commit 41e28a9f6f

View File

@ -37,7 +37,12 @@ from open_webui.routers.tasks import (
generate_chat_tags,
)
from open_webui.routers.retrieval import process_web_search, SearchForm
from open_webui.routers.images import image_generations, GenerateImageForm
from open_webui.routers.images import (
load_b64_image_data,
image_generations,
GenerateImageForm,
upload_image,
)
from open_webui.routers.pipelines import (
process_pipeline_inlet_filter,
process_pipeline_outlet_filter,
@ -2259,28 +2264,19 @@ async def process_chat_response(
stdoutLines = stdout.split("\n")
for idx, line in enumerate(stdoutLines):
if "data:image/png;base64" in line:
id = str(uuid4())
# ensure the path exists
os.makedirs(
os.path.join(CACHE_DIR, "images"),
exist_ok=True,
)
image_path = os.path.join(
CACHE_DIR,
f"images/{id}.png",
)
with open(image_path, "wb") as f:
f.write(
base64.b64decode(
line.split(",")[1]
)
image_url = ""
# Extract base64 image data from the line
image_data, content_type = load_b64_image_data(line)
if image_data is not None:
image_url = upload_image(
request,
image_data,
content_type,
metadata,
user,
)
stdoutLines[idx] = (
f"![Output Image {idx}](/cache/images/{id}.png)"
f"![Output Image]({image_url})"
)
output["stdout"] = "\n".join(stdoutLines)
@ -2291,30 +2287,22 @@ async def process_chat_response(
resultLines = result.split("\n")
for idx, line in enumerate(resultLines):
if "data:image/png;base64" in line:
id = str(uuid4())
# ensure the path exists
os.makedirs(
os.path.join(CACHE_DIR, "images"),
exist_ok=True,
image_url = ""
# Extract base64 image data from the line
image_data, content_type = (
load_b64_image_data(line)
)
image_path = os.path.join(
CACHE_DIR,
f"images/{id}.png",
)
with open(image_path, "wb") as f:
f.write(
base64.b64decode(
line.split(",")[1]
)
if image_data is not None:
image_url = upload_image(
request,
image_data,
content_type,
metadata,
user,
)
resultLines[idx] = (
f"![Output Image {idx}](/cache/images/{id}.png)"
f"![Output Image]({image_url})"
)
output["result"] = "\n".join(resultLines)
except Exception as e:
output = str(e)