Browse Source

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

非法操作 3 months ago
parent
commit
309fd76ddf

+ 1 - 1
api/core/tools/provider/builtin/comfyui/tools/comfyui_client.py

@@ -125,7 +125,7 @@ class ComfyUiClient:
             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()

+ 5 - 2
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 @@ class ComfyUIWorkflowTool(BuiltinTool):
 
         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