403 lines
16 KiB
Python
403 lines
16 KiB
Python
import os
|
|
import torch
|
|
import numpy as np
|
|
import tifffile
|
|
import folder_paths
|
|
import logging
|
|
from typing import Dict, Tuple, Optional, List
|
|
import OpenImageIO as oiio
|
|
from datetime import datetime
|
|
|
|
os.environ["OPENCV_IO_ENABLE_OPENEXR"] = "1"
|
|
import cv2 as cv
|
|
|
|
logger = logging.getLogger(__name__)
|
|
|
|
# Import debug utilities, preview utilities, and sequence utilities
|
|
try:
|
|
from ..utils.debug_utils import debug_log, format_tensor_info
|
|
from ..utils.preview_utils import generate_preview_for_comfyui
|
|
from ..utils.sequence_utils import SequenceHandler, DynamicUIHelper
|
|
except ImportError:
|
|
# Fallback if utils not available
|
|
def debug_log(logger, level, simple_msg, verbose_msg=None, **kwargs):
|
|
getattr(logger, level.lower())(simple_msg)
|
|
def format_tensor_info(tensor_shape, tensor_dtype, name=""):
|
|
return f"{name} shape={tensor_shape}" if name else f"shape={tensor_shape}"
|
|
def generate_preview_for_comfyui(image_tensor, source_path="", is_sequence=False, frame_index=0, full_size=False):
|
|
return None
|
|
|
|
# Fallback sequence handler
|
|
class SequenceHandler:
|
|
@staticmethod
|
|
def detect_sequence_pattern(path): return '####' in path if path else False
|
|
@staticmethod
|
|
def generate_frame_paths(pattern, start, count, step): return []
|
|
|
|
class DynamicUIHelper:
|
|
@staticmethod
|
|
def create_save_mode_widgets(): return {"sequence": [["start_frame", "INT", {"default": 1}], ["frame_step", "INT", {"default": 1}]]}
|
|
@staticmethod
|
|
def create_versioning_widgets(): return {"versioning": [["version", "INT", {"default": 1}]]}
|
|
|
|
class SaverNode:
|
|
"""Optimized image saver node with consistent bit depth handling"""
|
|
|
|
# Format specifications
|
|
FORMAT_SPECS = {
|
|
"exr": {"depths": [16, 32], "opencv": False}, # EXR only supports half and full float
|
|
"png": {"depths": [8, 16], "opencv": False}, # PNG only supports integer formats
|
|
"jpg": {"depths": [8], "opencv": True},
|
|
"webp": {"depths": [8], "opencv": True},
|
|
"tiff": {"depths": [8, 16, 32], "opencv": False}
|
|
}
|
|
|
|
@classmethod
|
|
def INPUT_TYPES(cls):
|
|
# Define format-specific widgets
|
|
format_widgets = {
|
|
"exr": [
|
|
["bit_depth", ["16", "32"], {"default": "32"}],
|
|
["exr_compression", ["none", "zip", "zips", "rle", "pxr24", "b44", "b44a", "dwaa", "dwab"], {"default": "zips"}],
|
|
["save_as_grayscale", "BOOLEAN", {"default": False}]
|
|
],
|
|
"png": [
|
|
["bit_depth", ["8", "16"], {"default": "16"}],
|
|
["save_as_grayscale", "BOOLEAN", {"default": False}]
|
|
],
|
|
"tiff": [
|
|
["bit_depth", ["8", "16", "32"], {"default": "16"}],
|
|
["save_as_grayscale", "BOOLEAN", {"default": False}]
|
|
],
|
|
"jpg": [
|
|
["quality", "INT", {"default": 95, "min": 1, "max": 100}]
|
|
],
|
|
"webp": [
|
|
["quality", "INT", {"default": 95, "min": 1, "max": 100}]
|
|
]
|
|
}
|
|
|
|
# Create sequence and versioning widgets using shared utilities
|
|
save_mode_widgets = DynamicUIHelper.create_save_mode_widgets()
|
|
versioning_widgets = DynamicUIHelper.create_versioning_widgets()
|
|
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE",),
|
|
"file_path": ("STRING", {"default": ""}),
|
|
"filename": ("STRING", {"default": "ComfyUI"}),
|
|
"save_mode": (["single", "sequence"], {
|
|
"default": "single",
|
|
"description": "Save mode: single files or sequence pattern",
|
|
"formats": save_mode_widgets
|
|
}),
|
|
"use_versioning": ("BOOLEAN", {
|
|
"default": False,
|
|
"description": "Enable version numbering",
|
|
"formats": versioning_widgets
|
|
}),
|
|
"file_type": (["exr", "png", "jpg", "webp", "tiff"], {"default": "png", "formats": format_widgets}),
|
|
},
|
|
"hidden": {
|
|
"prompt": "PROMPT",
|
|
"extra_pnginfo": "EXTRA_PNGINFO",
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ()
|
|
FUNCTION = "save_images"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "COCO Tools/Savers"
|
|
|
|
@classmethod
|
|
def IS_CHANGED(cls, **kwargs):
|
|
return float("NaN") # Always execute
|
|
|
|
def __init__(self):
|
|
self.output_dir = folder_paths.get_output_directory()
|
|
|
|
@staticmethod
|
|
def is_grayscale_fast(image: np.ndarray, sample_rate: float = 0.1) -> bool:
|
|
"""Fast grayscale detection by sampling pixels"""
|
|
if len(image.shape) == 2 or image.shape[-1] == 1:
|
|
return True
|
|
if image.shape[-1] == 3:
|
|
# Sample 10% of pixels for large images
|
|
total_pixels = image.shape[0] * image.shape[1]
|
|
if total_pixels > 1000000: # 1MP
|
|
indices = np.random.choice(total_pixels, int(total_pixels * sample_rate), replace=False)
|
|
flat_img = image.reshape(-1, 3)
|
|
sampled = flat_img[indices]
|
|
return np.allclose(sampled[:, 0], sampled[:, 1], rtol=0.01) and \
|
|
np.allclose(sampled[:, 1], sampled[:, 2], rtol=0.01)
|
|
else:
|
|
return np.allclose(image[..., 0], image[..., 1], rtol=0.01) and \
|
|
np.allclose(image[..., 1], image[..., 2], rtol=0.01)
|
|
return False
|
|
|
|
def validate_bit_depth(self, file_type: str, bit_depth: int) -> int:
|
|
"""Validate and adjust bit depth for format"""
|
|
valid_depths = self.FORMAT_SPECS[file_type]["depths"]
|
|
if bit_depth not in valid_depths:
|
|
return valid_depths[0]
|
|
return bit_depth
|
|
|
|
def convert_to_grayscale(self, img: np.ndarray) -> np.ndarray:
|
|
"""Proper grayscale conversion using luminance weights"""
|
|
if len(img.shape) == 2 or img.shape[-1] == 1:
|
|
return img if img.shape[-1] == 1 else img[..., np.newaxis]
|
|
|
|
if img.shape[-1] >= 3:
|
|
# ITU-R BT.709 luminance weights
|
|
weights = np.array([0.2126, 0.7152, 0.0722])
|
|
gray = np.dot(img[..., :3], weights)
|
|
return gray[..., np.newaxis]
|
|
|
|
return img[..., 0:1] # Fallback for 2-channel images
|
|
|
|
def prepare_image(self, img_tensor: torch.Tensor, save_as_grayscale: bool) -> np.ndarray:
|
|
"""Convert tensor to numpy and prepare channels"""
|
|
# Convert from tensor
|
|
if img_tensor.ndim == 4 and img_tensor.shape[0] == 1:
|
|
img_np = img_tensor.squeeze(0).cpu().numpy()
|
|
else:
|
|
img_np = img_tensor.cpu().numpy()
|
|
|
|
# Ensure float32 [0,1] range
|
|
img_np = np.clip(img_np, 0, 1).astype(np.float32)
|
|
|
|
# Handle channels
|
|
if len(img_np.shape) == 2:
|
|
img_np = img_np[..., np.newaxis]
|
|
|
|
# Convert to grayscale if explicitly requested (don't auto-detect for solid colors)
|
|
if save_as_grayscale:
|
|
img_np = self.convert_to_grayscale(img_np)
|
|
|
|
return img_np
|
|
|
|
def save_exr(self, img: np.ndarray, path: str, bit_depth: int, compression: str) -> None:
|
|
"""Save EXR with proper float handling"""
|
|
# Convert to appropriate type based on bit depth
|
|
if bit_depth == 16:
|
|
data = img.astype(np.float16)
|
|
pixel_type = oiio.HALF
|
|
else: # 32
|
|
data = img.astype(np.float32)
|
|
pixel_type = oiio.FLOAT
|
|
|
|
data = np.ascontiguousarray(data)
|
|
channels = 1 if data.ndim == 2 else data.shape[-1]
|
|
|
|
spec = oiio.ImageSpec(data.shape[1], data.shape[0], channels, pixel_type)
|
|
spec.attribute("compression", compression)
|
|
spec.attribute("Software", "COCO Tools")
|
|
|
|
buf = oiio.ImageBuf(spec)
|
|
buf.set_pixels(oiio.ROI(), data)
|
|
|
|
if not buf.write(path):
|
|
raise RuntimeError(f"Failed to write EXR: {oiio.geterror()}")
|
|
|
|
def save_png(self, img: np.ndarray, path: str, bit_depth: int) -> None:
|
|
"""Save PNG with proper bit depth handling"""
|
|
# Convert to appropriate type based on bit depth
|
|
if bit_depth == 8:
|
|
data = (img * 255).astype(np.uint8)
|
|
pixel_type = oiio.UINT8
|
|
elif bit_depth == 16:
|
|
data = (img * 65535).astype(np.uint16)
|
|
pixel_type = oiio.UINT16
|
|
else:
|
|
# Fallback to 8-bit for unsupported depths
|
|
data = (img * 255).astype(np.uint8)
|
|
pixel_type = oiio.UINT8
|
|
|
|
data = np.ascontiguousarray(data)
|
|
channels = 1 if data.ndim == 2 else data.shape[-1]
|
|
|
|
spec = oiio.ImageSpec(data.shape[1], data.shape[0], channels, pixel_type)
|
|
spec.attribute("compression", "zip")
|
|
spec.attribute("png:compressionLevel", 9)
|
|
|
|
buf = oiio.ImageBuf(spec)
|
|
buf.set_pixels(oiio.ROI(), data)
|
|
|
|
if not buf.write(path):
|
|
raise RuntimeError(f"Failed to write PNG: {oiio.geterror()}")
|
|
|
|
def save_opencv_format(self, img: np.ndarray, path: str, quality: int = 95) -> None:
|
|
"""Save JPEG/WebP using OpenCV"""
|
|
# Convert to 8-bit BGR
|
|
data = (img * 255).astype(np.uint8)
|
|
|
|
# Convert RGB to BGR only for 3+ channel images
|
|
if data.shape[-1] >= 3:
|
|
data = cv.cvtColor(data, cv.COLOR_RGB2BGR)
|
|
|
|
# Save with quality setting
|
|
if path.endswith(('.jpg', '.jpeg')):
|
|
cv.imwrite(path, data, [cv.IMWRITE_JPEG_QUALITY, quality])
|
|
else: # webp
|
|
cv.imwrite(path, data, [cv.IMWRITE_WEBP_QUALITY, quality])
|
|
|
|
def save_tiff(self, img: np.ndarray, path: str, bit_depth: int) -> None:
|
|
"""Save TIFF with proper bit depth"""
|
|
if bit_depth == 8:
|
|
data = (img * 255).astype(np.uint8)
|
|
elif bit_depth == 16:
|
|
data = (img * 65535).astype(np.uint16)
|
|
else: # 32
|
|
data = img.astype(np.float32)
|
|
|
|
# Determine photometric interpretation
|
|
photometric = 'minisblack' if data.shape[-1] == 1 else 'rgb'
|
|
|
|
tifffile.imwrite(path, data, photometric=photometric)
|
|
|
|
def get_unique_filepath(self, base_path: str) -> str:
|
|
"""Get unique filepath with incremental counter"""
|
|
if not os.path.exists(base_path):
|
|
return base_path
|
|
|
|
dir_name = os.path.dirname(base_path)
|
|
base_name = os.path.basename(base_path)
|
|
name, ext = os.path.splitext(base_name)
|
|
|
|
counter = 1
|
|
while True:
|
|
new_path = os.path.join(dir_name, f"{name}_{counter}{ext}")
|
|
if not os.path.exists(new_path):
|
|
return new_path
|
|
counter += 1
|
|
|
|
def save_images(self, images, file_path, filename, save_mode="single", file_type="png",
|
|
bit_depth=None, quality=None, save_as_grayscale=None, use_versioning=False,
|
|
version=1, start_frame=None, frame_step=None, prompt=None, extra_pnginfo=None,
|
|
exr_compression=None, **kwargs):
|
|
"""Main save function with optimized pipeline - handles missing contextual inputs and sequence mode"""
|
|
|
|
# Provide format-specific defaults for missing inputs
|
|
if bit_depth is None:
|
|
bit_depth = "16" if file_type in ["exr", "png", "tiff"] else "8"
|
|
|
|
if quality is None:
|
|
quality = 95 # Default for JPG/WebP
|
|
|
|
if save_as_grayscale is None:
|
|
save_as_grayscale = False
|
|
|
|
if exr_compression is None:
|
|
exr_compression = "zips" # Default for EXR
|
|
|
|
# Handle sequence parameters with defaults
|
|
if start_frame is None:
|
|
start_frame = 1
|
|
if frame_step is None:
|
|
frame_step = 1
|
|
|
|
try:
|
|
# Validate inputs
|
|
bit_depth = int(bit_depth)
|
|
file_type = file_type.lower()
|
|
|
|
# Determine if this is sequence mode and validate pattern
|
|
is_sequence_mode = save_mode == "sequence" or SequenceHandler.detect_sequence_pattern(filename)
|
|
|
|
debug_log(logger, "info", f"Saving {len(images)} images in {save_mode} mode",
|
|
f"Save mode: {save_mode}, Is sequence: {is_sequence_mode}, File type: {file_type}")
|
|
bit_depth = self.validate_bit_depth(file_type, bit_depth)
|
|
|
|
# Log save operation
|
|
debug_log(logger, "info", f"Saving {len(images)} image(s) as {file_type.upper()}",
|
|
f"Saving {len(images)} image(s) as {file_type.upper()} {bit_depth}-bit to {filename}")
|
|
|
|
# Build base path
|
|
if file_path:
|
|
full_path = os.path.join(self.output_dir, file_path) if not os.path.isabs(file_path) else file_path
|
|
os.makedirs(full_path, exist_ok=True)
|
|
base_path = os.path.join(full_path, filename)
|
|
else:
|
|
base_path = os.path.join(self.output_dir, filename)
|
|
|
|
# Add version string
|
|
version_str = f"_v{version:03d}" if use_versioning and version >= 0 else ""
|
|
|
|
# Track saved files for preview
|
|
saved_files = []
|
|
|
|
# Process each image - handle single vs sequence mode
|
|
for i, img_tensor in enumerate(images):
|
|
# Prepare image (all formats start from float32 [0,1])
|
|
img_np = self.prepare_image(img_tensor, save_as_grayscale)
|
|
|
|
# Build output path based on mode
|
|
if is_sequence_mode and SequenceHandler.detect_sequence_pattern(filename):
|
|
# Sequence mode with #### pattern
|
|
frame_number = start_frame + (i * frame_step)
|
|
sequence_filename = filename.replace('####', f'{frame_number:04d}')
|
|
out_path = f"{os.path.join(os.path.dirname(base_path), sequence_filename)}{version_str}.{file_type}"
|
|
elif is_sequence_mode:
|
|
# Sequence mode without pattern - use frame numbers
|
|
frame_number = start_frame + (i * frame_step)
|
|
out_path = f"{base_path}_{frame_number:04d}{version_str}.{file_type}"
|
|
else:
|
|
# Single mode - use index for multiple images
|
|
frame_str = f"_{i}" if len(images) > 1 else ""
|
|
out_path = f"{base_path}{version_str}{frame_str}.{file_type}"
|
|
|
|
out_path = self.get_unique_filepath(out_path)
|
|
|
|
# Save based on format
|
|
if file_type == "exr":
|
|
self.save_exr(img_np, out_path, bit_depth, exr_compression)
|
|
elif file_type == "png":
|
|
self.save_png(img_np, out_path, bit_depth)
|
|
elif file_type in ["jpg", "jpeg", "webp"]:
|
|
self.save_opencv_format(img_np, out_path, quality)
|
|
elif file_type == "tiff":
|
|
self.save_tiff(img_np, out_path, bit_depth)
|
|
|
|
# Track saved file info
|
|
saved_files.append({
|
|
"filename": os.path.basename(out_path),
|
|
"fullPath": out_path,
|
|
"format": file_type,
|
|
"bitDepth": bit_depth,
|
|
"index": i
|
|
})
|
|
|
|
# Generate full resolution preview for saved images
|
|
preview_data = generate_preview_for_comfyui(
|
|
images,
|
|
source_path=f"saver_{len(saved_files)}_files",
|
|
is_sequence=is_sequence_mode or len(saved_files) > 1,
|
|
frame_index=0
|
|
)
|
|
|
|
# Log completion
|
|
debug_log(logger, "info", f"Saved {len(saved_files)} files successfully",
|
|
f"Successfully saved {len(saved_files)} files: {[f['filename'] for f in saved_files]}")
|
|
|
|
# Return with preview and saved file info
|
|
result = {
|
|
"ui": {
|
|
"images": preview_data or [],
|
|
"saved_files": saved_files
|
|
}
|
|
}
|
|
|
|
return result
|
|
|
|
except Exception as e:
|
|
debug_log(logger, "error", "Save operation failed", f"Saver error: {str(e)}")
|
|
raise RuntimeError(f"Saver error: {str(e)}") from e
|
|
|
|
# NODE_CLASS_MAPPINGS = {
|
|
# "saver": saver,
|
|
# }
|
|
|
|
# NODE_DISPLAY_NAME_MAPPINGS = {
|
|
# "saver": "Image Saver"
|
|
# }
|