fix: comfyui output image's format (#12121)

langgenius · Dec 27, 2024 · 309fd76 · 309fd76
1 parent a3293b1
commit 309fd76
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 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
@@ -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
@@ -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