Jony
ba920f2b66
Replace the prior release tree with the current plugin, frontend, tests, and documentation. Document retired node IDs and the public Gitea update source.
139 lines
5.1 KiB
Python
139 lines
5.1 KiB
Python
"""保存图像节点 - 支持 PNG/JPEG/WebP 格式输出"""
|
|
|
|
import os
|
|
import json
|
|
import numpy as np
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
import folder_paths
|
|
from comfy.cli_args import args
|
|
|
|
|
|
class SaveImageFormat:
|
|
"""保存图像，支持 PNG / JPEG / WebP 三种格式"""
|
|
|
|
FORMATS = ["PNG", "JPEG", "WebP"]
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
self.type = "output"
|
|
self.compress_level = 4
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
return {
|
|
"required": {
|
|
"图像": ("IMAGE",),
|
|
"文件名前缀": ("STRING", {"default": "ComfyUI"}),
|
|
"输出格式": (cls.FORMATS, {"default": "PNG"}),
|
|
"质量": ("INT", {
|
|
"default": 100, "min": 1, "max": 100, "step": 1,
|
|
"tooltip": "图片质量。100=不压缩（JPEG 最高质量 / WebP 无损）；"
|
|
"小于 100 时按该数值压缩（如 90），仅对 JPEG / WebP 生效。",
|
|
}),
|
|
},
|
|
"optional": {
|
|
"保存路径": ("STRING", {"default": ""}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "image"
|
|
DESCRIPTION = "保存图像，支持 PNG / JPEG / WebP 格式输出。"
|
|
|
|
_EXT_MAP = {"PNG": ".png", "JPEG": ".jpg", "WebP": ".webp"}
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, 图像=None, **kwargs):
|
|
"""保存节点属于有副作用的输出节点，不能复用上次的缓存结果。"""
|
|
return float("nan")
|
|
|
|
def save_images(self, 图像=None, 文件名前缀="ComfyUI", 输出格式="PNG",
|
|
质量=100, 保存路径="", prompt=None, extra_pnginfo=None):
|
|
images = 图像
|
|
filename_prefix = 文件名前缀
|
|
format = 输出格式
|
|
quality = int(质量)
|
|
custom_dir = (保存路径 or "").strip()
|
|
|
|
if custom_dir:
|
|
# 保存到用户指定的文件夹。自定义路径不会经过
|
|
# folder_paths.get_save_image_path()，因此需要在此处自行避让重名。
|
|
full_output_folder = custom_dir
|
|
os.makedirs(full_output_folder, exist_ok=True)
|
|
filename = filename_prefix
|
|
counter = 1
|
|
subfolder = ""
|
|
else:
|
|
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]
|
|
)
|
|
|
|
ext = self._EXT_MAP.get(format, ".png")
|
|
results = []
|
|
|
|
# 为整个输入批次预留一段连续编号，避免重复运行时覆盖已有文件。
|
|
# 默认 output 路径和自定义保存路径都在这里复核一次；这样即使外部
|
|
# 编号器返回了已使用的计数，也不会覆盖。兼容 %batch_num% 占位符。
|
|
while any(
|
|
os.path.exists(
|
|
os.path.join(
|
|
full_output_folder,
|
|
f"{filename.replace('%batch_num%', str(batch_number))}_{counter + batch_number:05}_{ext}",
|
|
)
|
|
)
|
|
for batch_number in range(len(images))
|
|
):
|
|
counter += 1
|
|
|
|
for batch_number, image in enumerate(images):
|
|
i = 255.0 * image.cpu().numpy()
|
|
img = Image.fromarray(np.clip(i, 0, 255).astype(np.uint8))
|
|
|
|
filename_with_batch_num = filename.replace("%batch_num%", str(batch_number))
|
|
file = f"{filename_with_batch_num}_{counter:05}_{ext}"
|
|
|
|
filepath = os.path.join(full_output_folder, file)
|
|
|
|
if format == "PNG":
|
|
metadata = None
|
|
if not args.disable_metadata:
|
|
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]))
|
|
img.save(filepath, pnginfo=metadata,
|
|
compress_level=self.compress_level)
|
|
elif format == "JPEG":
|
|
if img.mode == "RGBA":
|
|
img = img.convert("RGB")
|
|
if quality >= 100:
|
|
img.save(filepath, quality=100, optimize=True)
|
|
else:
|
|
img.save(filepath, quality=quality, optimize=True)
|
|
elif format == "WebP":
|
|
if quality >= 100:
|
|
img.save(filepath, lossless=True)
|
|
else:
|
|
img.save(filepath, quality=quality, method=6)
|
|
|
|
results.append({
|
|
"filename": file,
|
|
"subfolder": subfolder,
|
|
"type": self.type,
|
|
})
|
|
counter += 1
|
|
|
|
return {"ui": {"images": results}}
|