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