From 309fd76ddf029440484c62b33e3c8a728efd9d56 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=9D=9E=E6=B3=95=E6=93=8D=E4=BD=9C?= Date: Fri, 27 Dec 2024 20:20:03 +0800 Subject: [PATCH] fix: comfyui output image's format (#12121) --- .../tools/provider/builtin/comfyui/tools/comfyui_client.py | 2 +- .../provider/builtin/comfyui/tools/comfyui_workflow.py | 7 +++++-- 2 files changed, 6 insertions(+), 3 deletions(-) diff --git a/api/core/tools/provider/builtin/comfyui/tools/comfyui_client.py b/api/core/tools/provider/builtin/comfyui/tools/comfyui_client.py index f994cdbf66e78b..2bf10ce8ff2632 100644 --- a/api/core/tools/provider/builtin/comfyui/tools/comfyui_client.py +++ b/api/core/tools/provider/builtin/comfyui/tools/comfyui_client.py @@ -125,7 +125,7 @@ def generate_image_by_prompt(self, prompt: dict) -> list[bytes]: for output in history["outputs"].values(): for img in output.get("images", []): image_data = self.get_image(img["filename"], img["subfolder"], img["type"]) - images.append(image_data) + images.append((image_data, img["filename"])) return images finally: ws.close() diff --git a/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py b/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py index 87837362779baa..eb085f221ebdda 100644 --- a/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py +++ b/api/core/tools/provider/builtin/comfyui/tools/comfyui_workflow.py @@ -1,4 +1,5 @@ import json +import mimetypes from typing import Any from core.file import FileType @@ -75,10 +76,12 @@ def _invoke(self, user_id: str, tool_parameters: dict[str, Any]) -> ToolInvokeMe images = comfyui.generate_image_by_prompt(prompt) result = [] - for img in images: + for image_data, filename in images: result.append( self.create_blob_message( - blob=img, meta={"mime_type": "image/png"}, save_as=self.VariableKey.IMAGE.value + blob=image_data, + meta={"mime_type": mimetypes.guess_type(filename)[0]}, + save_as=self.VariableKey.IMAGE.value, ) ) return result