This commit is contained in:
karrix 2024-10-22 12:54:33 +08:00
parent 0779136134
commit d41c4de352
2 changed files with 105 additions and 13 deletions

View File

@ -0,0 +1,92 @@
import os
import json
import numpy as np
from PIL import Image
from PIL.PngImagePlugin import PngInfo
import folder_paths
class ComfyDeployStdOutputImage:
def __init__(self):
self.output_dir = folder_paths.get_output_directory()
self.type = "std_output"
self.prefix_append = ""
self.compress_level = 4
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"images": ("IMAGE", {"tooltip": "The images to save."}),
"filename_prefix": (
"STRING",
{
"default": "ComfyUI",
"tooltip": "The prefix for the file to save. This may include formatting information such as %date:yyyy-MM-dd% or %Empty Latent Image.width% to include values from nodes.",
},
),
"file_type": (["png", "jpg", "webp"], {"default": "webp"}),
"quality": ("INT", {"default": 80, "min": 1, "max": 100, "step": 1}),
},
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
}
RETURN_TYPES = ()
FUNCTION = "run"
OUTPUT_NODE = True
CATEGORY = "std_output"
DESCRIPTION = "Saves the input images to your ComfyUI output directory."
def run(
self,
images,
filename_prefix="ComfyUI",
file_type="png",
quality=80,
prompt=None,
extra_pnginfo=None,
):
filename_prefix += self.prefix_append
full_output_folder, filename, counter, subfolder, filename_prefix = (
folder_paths.get_save_image_path(
filename_prefix, self.output_dir, images[0].shape[1], images[0].shape[0]
)
)
results = list()
for batch_number, image in enumerate(images):
i = 255.0 * image.cpu().numpy()
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
metadata = PngInfo()
if prompt is not None:
metadata.add_text("prompt", json.dumps(prompt))
if extra_pnginfo is not None:
for x in extra_pnginfo:
metadata.add_text(x, json.dumps(extra_pnginfo[x]))
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
file = f"{filename_with_batch_num}_{counter:05}_.{file_type}"
file_path = os.path.join(full_output_folder, file)
if file_type == "png":
img.save(
file_path, pnginfo=metadata, compress_level=self.compress_level
)
elif file_type == "jpg":
img.save(file_path, quality=quality, optimize=True)
elif file_type == "webp":
img.save(file_path, quality=quality)
results.append(
{"filename": file, "subfolder": subfolder, "type": self.type}
)
counter += 1
return {"ui": {"images": results}}
NODE_CLASS_MAPPINGS = {"ComfyDeployStdOutputImage": ComfyDeployStdOutputImage}
NODE_DISPLAY_NAME_MAPPINGS = {
"ComfyDeployStdOutputImage": "Standard Image Output (ComfyDeploy)"
}

View File

@ -1217,7 +1217,7 @@ async def send_json_override(self, event, data, sid=None):
await update_run(prompt_id, Status.FAILED)
# await update_run_with_output(prompt_id, data)
if event == "executed" and "node" in data and "output" in data:
if event == "executed" and "node" in data and "std_output" in data:
node_meta = None
if prompt_id in prompt_metadata:
node = data.get("node")
@ -1226,19 +1226,19 @@ async def send_json_override(self, event, data, sid=None):
"node_id": node,
"node_class": class_type,
}
if class_type == "PreviewImage":
logger.info("Skipping preview image")
else:
await update_run_with_output(
prompt_id,
data.get("output"),
node_id=data.get("node"),
node_meta=node_meta,
# if class_type == "PreviewImage":
# logger.info("Skipping preview image")
# else:
await update_run_with_output(
prompt_id,
data.get("std_output"),
node_id=data.get("node"),
node_meta=node_meta,
)
if prompt_id in comfy_message_queues:
comfy_message_queues[prompt_id].put_nowait(
{"event": "output_ready", "data": data}
)
if prompt_id in comfy_message_queues:
comfy_message_queues[prompt_id].put_nowait(
{"event": "output_ready", "data": data}
)
logger.info(f"Executed {class_type} {data}")
else:
logger.info(f"Executed {data}")