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.
599 lines
20 KiB
Python
599 lines
20 KiB
Python
"""Format-aware saving for the unified o1key image workflow.
|
|
|
|
The generator keeps provider bytes in memory or in ComfyUI's temp directory.
|
|
Only the save node promotes those results into a permanent save directory.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
import re
|
|
import struct
|
|
import threading
|
|
import uuid
|
|
import zlib
|
|
from io import BytesIO
|
|
from typing import Any, Iterable
|
|
|
|
import numpy as np
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
|
|
|
|
SAVE_FORMAT_OPTIONS = ("原始", "png", "jpg", "webp")
|
|
SAVE_NAMING_RULE_OPTIONS = ("和主图一致", "自然数字", "自定义前缀")
|
|
DEFAULT_SAVE_NAMING_RULE = "自定义前缀"
|
|
_FORMAT_EXTENSIONS = {"PNG": "png", "JPEG": "jpg", "WEBP": "webp"}
|
|
_PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"
|
|
_SAVE_LOCK = threading.Lock()
|
|
_UNSAFE_FILENAME_CHARS = re.compile(r'[<>:"/\\|?*\x00-\x1f]')
|
|
_WINDOWS_RESERVED_STEMS = {
|
|
"CON", "PRN", "AUX", "NUL",
|
|
*(f"COM{index}" for index in range(1, 10)),
|
|
*(f"LPT{index}" for index in range(1, 10)),
|
|
}
|
|
|
|
|
|
def normalize_save_format(value: Any) -> str:
|
|
normalized = str(value or "原始").strip().lower()
|
|
if normalized == "jpeg":
|
|
normalized = "jpg"
|
|
if normalized == "原始":
|
|
return "原始"
|
|
if normalized not in {"png", "jpg", "webp"}:
|
|
raise ValueError("保存格式仅支持：原始、png、jpg、webp")
|
|
return normalized
|
|
|
|
|
|
def normalize_save_location(value: Any) -> str:
|
|
"""Return an output-relative directory or an explicit absolute directory."""
|
|
|
|
raw = str(value or "").strip().replace("\\", "/")
|
|
if not raw or raw.lower() == "output" or raw == ".":
|
|
return ""
|
|
drive, _ = os.path.splitdrive(raw)
|
|
if os.path.isabs(raw):
|
|
return os.path.normpath(raw)
|
|
if drive:
|
|
raise ValueError("盘符路径必须填写完整绝对路径，例如 D:/图片")
|
|
parts = [part for part in raw.split("/") if part not in {"", "."}]
|
|
if not parts or any(part == ".." for part in parts):
|
|
raise ValueError("相对保存位置必须是 output 下的子目录，不能包含 ..")
|
|
return os.path.join(*parts)
|
|
|
|
|
|
def _external_preview_descriptor(
|
|
encoded: bytes,
|
|
filename: str,
|
|
folder_paths_module,
|
|
) -> dict[str, Any]:
|
|
"""Keep external saves previewable without exposing local paths to workflows."""
|
|
|
|
temp_root = os.path.abspath(folder_paths_module.get_temp_directory())
|
|
preview_subfolder = os.path.join("o1key_external_preview", uuid.uuid4().hex)
|
|
preview_directory = os.path.join(temp_root, preview_subfolder)
|
|
os.makedirs(preview_directory, exist_ok=False)
|
|
_atomic_write(os.path.join(preview_directory, filename), encoded)
|
|
return {
|
|
"filename": filename,
|
|
"subfolder": preview_subfolder,
|
|
"type": "temp",
|
|
"external_saved": True,
|
|
}
|
|
|
|
|
|
def normalize_naming_rule(value: Any) -> str:
|
|
normalized = str(value or DEFAULT_SAVE_NAMING_RULE).strip()
|
|
if normalized not in SAVE_NAMING_RULE_OPTIONS:
|
|
raise ValueError("命名规则仅支持：和主图一致、自然数字、自定义前缀")
|
|
return normalized
|
|
|
|
|
|
def _safe_filename_stem(value: Any) -> str:
|
|
filename = os.path.basename(str(value or "").strip().replace("\\", "/"))
|
|
stem, _ = os.path.splitext(filename)
|
|
stem = _UNSAFE_FILENAME_CHARS.sub("_", stem).strip(" .")
|
|
stem = stem[:240] or "o1key"
|
|
if stem.upper() in _WINDOWS_RESERVED_STEMS:
|
|
stem = f"_{stem}"
|
|
return stem
|
|
|
|
|
|
def detect_image_format(data: bytes) -> str | None:
|
|
if data.startswith(_PNG_SIGNATURE):
|
|
return "PNG"
|
|
if data.startswith(b"\xff\xd8\xff"):
|
|
return "JPEG"
|
|
if len(data) >= 12 and data[:4] == b"RIFF" and data[8:12] == b"WEBP":
|
|
return "WEBP"
|
|
return None
|
|
|
|
|
|
def _metadata_enabled() -> bool:
|
|
try:
|
|
from comfy.cli_args import args
|
|
|
|
return not bool(args.disable_metadata)
|
|
except Exception:
|
|
return True
|
|
|
|
|
|
def _metadata_items(
|
|
prompt: Any = None,
|
|
extra_pnginfo: dict[str, Any] | None = None,
|
|
) -> list[tuple[str, str]]:
|
|
if not _metadata_enabled():
|
|
return []
|
|
items: list[tuple[str, str]] = []
|
|
if prompt is not None:
|
|
items.append(("prompt", json.dumps(prompt)))
|
|
if isinstance(extra_pnginfo, dict):
|
|
for key, value in extra_pnginfo.items():
|
|
items.append((str(key), json.dumps(value)))
|
|
return items
|
|
|
|
|
|
def _has_workflow_metadata(extra_pnginfo: dict[str, Any] | None) -> bool:
|
|
"""Match ComfyUI's recoverable image contract.
|
|
|
|
The frontend metadata parser restores workflows from PNG and WebP, but not
|
|
from JPEG. Native SaveImage always writes PNG, so metadata-bearing JPEG
|
|
results must use that same container instead of an EXIF-only JPEG.
|
|
"""
|
|
|
|
return (
|
|
_metadata_enabled()
|
|
and isinstance(extra_pnginfo, dict)
|
|
and extra_pnginfo.get("workflow") is not None
|
|
)
|
|
|
|
|
|
def _png_text_chunk(key: str, value: str) -> bytes:
|
|
payload = key.encode("latin-1", "replace") + b"\0" + value.encode(
|
|
"latin-1", "replace"
|
|
)
|
|
chunk_type = b"tEXt"
|
|
return (
|
|
struct.pack(">I", len(payload))
|
|
+ chunk_type
|
|
+ payload
|
|
+ struct.pack(">I", zlib.crc32(chunk_type + payload) & 0xFFFFFFFF)
|
|
)
|
|
|
|
|
|
def inject_png_metadata(data: bytes, metadata: Iterable[tuple[str, str]]) -> bytes:
|
|
"""Append tEXt chunks without decoding or recompressing PNG pixels."""
|
|
chunks = [_png_text_chunk(key, value) for key, value in metadata]
|
|
if not chunks or not data.startswith(_PNG_SIGNATURE):
|
|
return data
|
|
|
|
offset = len(_PNG_SIGNATURE)
|
|
while offset + 12 <= len(data):
|
|
length = struct.unpack(">I", data[offset : offset + 4])[0]
|
|
end = offset + 12 + length
|
|
if end > len(data):
|
|
return data
|
|
if data[offset + 4 : offset + 8] == b"IEND":
|
|
return data[:offset] + b"".join(chunks) + data[offset:]
|
|
offset = end
|
|
return data
|
|
|
|
|
|
def _create_exif_bytes(metadata: Iterable[tuple[str, str]]) -> bytes:
|
|
items = list(metadata)
|
|
if not items:
|
|
return b""
|
|
exif = Image.Exif()
|
|
extra_index = 0
|
|
for key, value in items:
|
|
if key == "prompt":
|
|
exif[0x0110] = f"prompt:{value}"
|
|
else:
|
|
exif[0x010F - extra_index] = f"{key}:{value}"
|
|
extra_index += 1
|
|
try:
|
|
return exif.tobytes()
|
|
except (OSError, OverflowError, ValueError):
|
|
# JPEG APP1 and some Pillow EXIF writers have practical size limits.
|
|
# The image must still save; PNG remains the fully lossless metadata path.
|
|
return b""
|
|
|
|
|
|
def inject_jpeg_metadata(data: bytes, metadata: Iterable[tuple[str, str]]) -> bytes:
|
|
"""Insert an EXIF APP1 segment without recompressing JPEG pixels."""
|
|
if not data.startswith(b"\xff\xd8"):
|
|
return data
|
|
exif = _create_exif_bytes(metadata)
|
|
if not exif or len(exif) + 2 > 0xFFFF:
|
|
return data
|
|
segment = b"\xff\xe1" + struct.pack(">H", len(exif) + 2) + exif
|
|
return data[:2] + segment + data[2:]
|
|
|
|
|
|
def _webp_chunks(data: bytes) -> list[tuple[bytes, bytes]] | None:
|
|
if len(data) < 12 or data[:4] != b"RIFF" or data[8:12] != b"WEBP":
|
|
return None
|
|
chunks: list[tuple[bytes, bytes]] = []
|
|
offset = 12
|
|
while offset + 8 <= len(data):
|
|
kind = data[offset : offset + 4]
|
|
length = struct.unpack("<I", data[offset + 4 : offset + 8])[0]
|
|
start = offset + 8
|
|
end = start + length
|
|
if end > len(data):
|
|
return None
|
|
chunks.append((kind, data[start:end]))
|
|
offset = end + (length & 1)
|
|
return chunks
|
|
|
|
|
|
def _pack_webp_chunk(kind: bytes, payload: bytes) -> bytes:
|
|
padding = b"\0" if len(payload) & 1 else b""
|
|
return kind + struct.pack("<I", len(payload)) + payload + padding
|
|
|
|
|
|
def inject_webp_metadata(data: bytes, metadata: Iterable[tuple[str, str]]) -> bytes:
|
|
"""Add/replace a WebP EXIF chunk without recompressing image data."""
|
|
exif = _create_exif_bytes(metadata)
|
|
chunks = _webp_chunks(data)
|
|
if not exif or chunks is None:
|
|
return data
|
|
|
|
chunks = [(kind, payload) for kind, payload in chunks if kind != b"EXIF"]
|
|
vp8x_index = next((i for i, item in enumerate(chunks) if item[0] == b"VP8X"), None)
|
|
if vp8x_index is None:
|
|
try:
|
|
with Image.open(BytesIO(data)) as image:
|
|
width, height = image.size
|
|
has_alpha = "A" in image.getbands()
|
|
except Exception:
|
|
return data
|
|
flags = 0x08 | (0x10 if has_alpha else 0)
|
|
vp8x = bytes((flags, 0, 0, 0)) + (width - 1).to_bytes(3, "little") + (
|
|
height - 1
|
|
).to_bytes(3, "little")
|
|
chunks.insert(0, (b"VP8X", vp8x))
|
|
else:
|
|
kind, payload = chunks[vp8x_index]
|
|
if len(payload) != 10:
|
|
return data
|
|
chunks[vp8x_index] = (kind, bytes((payload[0] | 0x08,)) + payload[1:])
|
|
|
|
chunks.append((b"EXIF", exif))
|
|
body = b"WEBP" + b"".join(_pack_webp_chunk(kind, payload) for kind, payload in chunks)
|
|
return b"RIFF" + struct.pack("<I", len(body)) + body
|
|
|
|
|
|
def inject_workflow_metadata(
|
|
data: bytes,
|
|
image_format: str,
|
|
prompt: Any = None,
|
|
extra_pnginfo: dict[str, Any] | None = None,
|
|
) -> bytes:
|
|
metadata = _metadata_items(prompt, extra_pnginfo)
|
|
if not metadata:
|
|
return data
|
|
if image_format == "PNG":
|
|
return inject_png_metadata(data, metadata)
|
|
if image_format == "JPEG":
|
|
return inject_jpeg_metadata(data, metadata)
|
|
if image_format == "WEBP":
|
|
return inject_webp_metadata(data, metadata)
|
|
return data
|
|
|
|
|
|
def _atomic_write(path: str, data: bytes) -> None:
|
|
temporary = f"{path}.{uuid.uuid4().hex}.tmp"
|
|
reserved = False
|
|
try:
|
|
descriptor = os.open(path, os.O_WRONLY | os.O_CREAT | os.O_EXCL)
|
|
os.close(descriptor)
|
|
reserved = True
|
|
with open(temporary, "wb") as handle:
|
|
handle.write(data)
|
|
os.replace(temporary, path)
|
|
reserved = False
|
|
finally:
|
|
if os.path.exists(temporary):
|
|
try:
|
|
os.remove(temporary)
|
|
except OSError:
|
|
pass
|
|
if reserved and os.path.exists(path):
|
|
try:
|
|
os.remove(path)
|
|
except OSError:
|
|
pass
|
|
|
|
|
|
def _tensor_uint8(image_tensor) -> np.ndarray:
|
|
return np.clip(255.0 * image_tensor.detach().cpu().numpy(), 0, 255).astype(np.uint8)
|
|
|
|
|
|
def _raw_bytes_match_tensor(data: bytes, image_tensor) -> bool:
|
|
try:
|
|
array = _tensor_uint8(image_tensor)
|
|
mode = "RGBA" if array.shape[-1] == 4 else "RGB"
|
|
with Image.open(BytesIO(data)) as source:
|
|
source.load()
|
|
decoded = np.asarray(source.convert(mode), dtype=np.uint8)
|
|
return decoded.shape == array.shape and np.array_equal(decoded, array)
|
|
except Exception:
|
|
return False
|
|
|
|
|
|
def _source_for_tensor(images, index: int) -> tuple[bytes | None, str | None, str | None]:
|
|
metadata = getattr(images, "_o1key_source_metadata", None)
|
|
if not isinstance(metadata, list) or index >= len(metadata):
|
|
return None, None, None
|
|
item = metadata[index]
|
|
if not isinstance(item, dict):
|
|
return None, None, None
|
|
source_filename = item.get("filename")
|
|
if not isinstance(source_filename, str) or not source_filename.strip():
|
|
source_filename = None
|
|
if item.get("modified"):
|
|
return None, None, source_filename
|
|
|
|
data = item.get("bytes")
|
|
if not isinstance(data, bytes):
|
|
source_path = item.get("path")
|
|
if isinstance(source_path, str) and os.path.isfile(source_path):
|
|
try:
|
|
with open(source_path, "rb") as handle:
|
|
data = handle.read()
|
|
except OSError:
|
|
data = None
|
|
if not isinstance(data, bytes) or not _raw_bytes_match_tensor(data, images[index]):
|
|
data = None
|
|
|
|
image_format = detect_image_format(data) if data else None
|
|
if image_format is None:
|
|
hinted = str(item.get("format") or "").upper()
|
|
if hinted == "JPG":
|
|
hinted = "JPEG"
|
|
image_format = hinted if hinted in _FORMAT_EXTENSIONS else None
|
|
return data, image_format, source_filename
|
|
|
|
|
|
def _pil_from_tensor(image_tensor) -> Image.Image:
|
|
return Image.fromarray(_tensor_uint8(image_tensor))
|
|
|
|
|
|
def _encode_image(
|
|
image: Image.Image,
|
|
image_format: str,
|
|
prompt: Any = None,
|
|
extra_pnginfo: dict[str, Any] | None = None,
|
|
) -> bytes:
|
|
output = BytesIO()
|
|
metadata = _metadata_items(prompt, extra_pnginfo)
|
|
working = image
|
|
converted = None
|
|
try:
|
|
if image_format == "PNG":
|
|
pnginfo = None
|
|
if metadata:
|
|
pnginfo = PngInfo()
|
|
for key, value in metadata:
|
|
pnginfo.add_text(key, value)
|
|
working.save(output, format="PNG", pnginfo=pnginfo, compress_level=4)
|
|
elif image_format == "JPEG":
|
|
if working.mode != "RGB":
|
|
converted = working.convert("RGB")
|
|
working = converted
|
|
kwargs: dict[str, Any] = {
|
|
"format": "JPEG",
|
|
"quality": 100,
|
|
"subsampling": 0,
|
|
"optimize": True,
|
|
}
|
|
exif = _create_exif_bytes(metadata)
|
|
if exif:
|
|
kwargs["exif"] = exif
|
|
working.save(output, **kwargs)
|
|
elif image_format == "WEBP":
|
|
kwargs = {"format": "WEBP", "lossless": True, "quality": 100, "method": 4}
|
|
exif = _create_exif_bytes(metadata)
|
|
if exif:
|
|
kwargs["exif"] = exif
|
|
working.save(output, **kwargs)
|
|
else:
|
|
raise ValueError(f"不支持的图像格式：{image_format}")
|
|
return output.getvalue()
|
|
finally:
|
|
if converted is not None:
|
|
converted.close()
|
|
|
|
|
|
def _save_sources(
|
|
sources: list[tuple[bytes | None, str | None, Image.Image, str | None]],
|
|
filename_prefix: str,
|
|
save_format: str,
|
|
output_directory: str,
|
|
folder_paths_module,
|
|
prompt: Any = None,
|
|
extra_pnginfo: dict[str, Any] | None = None,
|
|
save_location: str = "",
|
|
naming_rule: str = DEFAULT_SAVE_NAMING_RULE,
|
|
main_filename: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
if not sources:
|
|
return []
|
|
width, height = sources[0][2].size
|
|
normalized_location = normalize_save_location(save_location)
|
|
external_location = bool(normalized_location and os.path.isabs(normalized_location))
|
|
relative_location = "" if external_location else normalized_location
|
|
save_root = normalized_location if external_location else output_directory
|
|
selected_rule = normalize_naming_rule(naming_rule)
|
|
filename = ""
|
|
counter = 1
|
|
if selected_rule == "自定义前缀":
|
|
effective_prefix = filename_prefix or "o1key"
|
|
if relative_location:
|
|
effective_prefix = os.path.join(relative_location, effective_prefix)
|
|
full_output_folder, filename, counter, subfolder, _ = (
|
|
folder_paths_module.get_save_image_path(
|
|
effective_prefix,
|
|
save_root,
|
|
width,
|
|
height,
|
|
)
|
|
)
|
|
else:
|
|
full_output_folder = os.path.join(save_root, relative_location)
|
|
subfolder = relative_location
|
|
os.makedirs(full_output_folder, exist_ok=True)
|
|
selected = normalize_save_format(save_format)
|
|
main_stem = _safe_filename_stem(
|
|
main_filename or next((item[3] for item in sources if item[3]), None)
|
|
)
|
|
results: list[dict[str, Any]] = []
|
|
natural_number = 1
|
|
main_number = 0
|
|
for batch_number, (raw, source_format, image, _source_filename) in enumerate(sources):
|
|
target_format = source_format if selected == "原始" else {
|
|
"png": "PNG",
|
|
"jpg": "JPEG",
|
|
"webp": "WEBP",
|
|
}[selected]
|
|
if target_format not in _FORMAT_EXTENSIONS:
|
|
target_format = "PNG"
|
|
if target_format == "JPEG" and _has_workflow_metadata(extra_pnginfo):
|
|
target_format = "PNG"
|
|
extension = _FORMAT_EXTENSIONS[target_format]
|
|
if selected == "原始" and raw is not None and detect_image_format(raw) == target_format:
|
|
encoded = inject_workflow_metadata(
|
|
raw,
|
|
target_format,
|
|
prompt=prompt,
|
|
extra_pnginfo=extra_pnginfo,
|
|
)
|
|
else:
|
|
encoded = _encode_image(
|
|
image,
|
|
target_format,
|
|
prompt=prompt,
|
|
extra_pnginfo=extra_pnginfo,
|
|
)
|
|
|
|
with _SAVE_LOCK:
|
|
while True:
|
|
if selected_rule == "自然数字":
|
|
saved_name = f"{natural_number}.{extension}"
|
|
natural_number += 1
|
|
elif selected_rule == "和主图一致":
|
|
suffix = "" if main_number == 0 else str(main_number)
|
|
saved_name = f"{main_stem}{suffix}.{extension}"
|
|
main_number += 1
|
|
else:
|
|
filename_with_batch_num = filename.replace(
|
|
"%batch_num%", str(batch_number)
|
|
)
|
|
saved_name = f"{filename_with_batch_num}_{counter:05}_.{extension}"
|
|
counter += 1
|
|
output_path = os.path.join(full_output_folder, saved_name)
|
|
if os.path.exists(output_path):
|
|
continue
|
|
try:
|
|
_atomic_write(output_path, encoded)
|
|
break
|
|
except FileExistsError:
|
|
continue
|
|
if external_location:
|
|
results.append(
|
|
_external_preview_descriptor(encoded, saved_name, folder_paths_module)
|
|
)
|
|
else:
|
|
results.append({"filename": saved_name, "subfolder": subfolder, "type": "output"})
|
|
return results
|
|
|
|
|
|
def save_tensor_images(
|
|
images,
|
|
filename_prefix: str,
|
|
save_format: str,
|
|
output_directory: str,
|
|
folder_paths_module,
|
|
prompt: Any = None,
|
|
extra_pnginfo: dict[str, Any] | None = None,
|
|
save_location: str = "",
|
|
naming_rule: str = DEFAULT_SAVE_NAMING_RULE,
|
|
main_filename: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
sources: list[tuple[bytes | None, str | None, Image.Image, str | None]] = []
|
|
try:
|
|
for index, image_tensor in enumerate(images):
|
|
raw, source_format, source_filename = _source_for_tensor(images, index)
|
|
sources.append((raw, source_format, _pil_from_tensor(image_tensor), source_filename))
|
|
return _save_sources(
|
|
sources,
|
|
filename_prefix,
|
|
save_format,
|
|
output_directory,
|
|
folder_paths_module,
|
|
prompt,
|
|
extra_pnginfo,
|
|
save_location,
|
|
naming_rule,
|
|
main_filename or getattr(images, "_o1key_main_filename", None),
|
|
)
|
|
finally:
|
|
for _, _, image, _ in sources:
|
|
image.close()
|
|
|
|
|
|
def save_temp_images(
|
|
source_paths: list[str],
|
|
filename_prefix: str,
|
|
save_format: str,
|
|
output_directory: str,
|
|
folder_paths_module,
|
|
prompt: Any = None,
|
|
extra_pnginfo: dict[str, Any] | None = None,
|
|
save_location: str = "",
|
|
naming_rule: str = DEFAULT_SAVE_NAMING_RULE,
|
|
main_filename: str | None = None,
|
|
) -> list[dict[str, Any]]:
|
|
sources: list[tuple[bytes | None, str | None, Image.Image, str | None]] = []
|
|
try:
|
|
for source_path in source_paths:
|
|
with open(source_path, "rb") as handle:
|
|
raw = handle.read()
|
|
image_format = detect_image_format(raw)
|
|
if image_format not in _FORMAT_EXTENSIONS:
|
|
raise ValueError(f"临时结果不是受支持的图像：{os.path.basename(source_path)}")
|
|
with Image.open(BytesIO(raw)) as opened:
|
|
opened.load()
|
|
image = opened.copy()
|
|
sources.append((raw, image_format, image, None))
|
|
return _save_sources(
|
|
sources,
|
|
filename_prefix,
|
|
save_format,
|
|
output_directory,
|
|
folder_paths_module,
|
|
prompt,
|
|
extra_pnginfo,
|
|
save_location,
|
|
naming_rule,
|
|
main_filename,
|
|
)
|
|
finally:
|
|
for _, _, image, _ in sources:
|
|
image.close()
|
|
|
|
|
|
__all__ = [
|
|
"SAVE_FORMAT_OPTIONS",
|
|
"SAVE_NAMING_RULE_OPTIONS",
|
|
"detect_image_format",
|
|
"inject_workflow_metadata",
|
|
"normalize_save_format",
|
|
"normalize_save_location",
|
|
"normalize_naming_rule",
|
|
"save_temp_images",
|
|
"save_tensor_images",
|
|
]
|