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

This commit is contained in:
非法操作 2024-12-27 20:20:03 +08:00 committed by GitHub
parent a3293b154e
commit 309fd76ddf
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
2 changed files with 6 additions and 3 deletions

View File

@ -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()

View File

@ -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