diff --git a/backend/open_webui/utils/middleware.py b/backend/open_webui/utils/middleware.py index 6510d7c99..a07262fa9 100644 --- a/backend/open_webui/utils/middleware.py +++ b/backend/open_webui/utils/middleware.py @@ -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)