From aa8e3417d23a0a2106616de4de868b859e059e84 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Thu, 26 Feb 2026 16:11:26 +0100 Subject: [PATCH] Move FastAbsoluteSaver from Sharp-Selector, add metadata PNG sidecar option and optimize - Move FastAbsoluteSaver node from ComfyUI-Sharp-Selector into this pack - Add save_metadata_png toggle: embeds workflow in a sidecar PNG for webp/video exports, or in the first file for PNG sequences - Batch GPU->CPU tensor transfer (single sync instead of per-image) - Remove dead webp exif code and unused variables/imports Co-Authored-By: Claude Opus 4.6 --- __init__.py | 8 +- fast_saver.py | 436 ++++++++++++++++++++++++++++++++++++++++++++++++++ 2 files changed, 442 insertions(+), 2 deletions(-) create mode 100644 fast_saver.py diff --git a/__init__.py b/__init__.py index 95ff234..690a273 100644 --- a/__init__.py +++ b/__init__.py @@ -10,9 +10,13 @@ from .image_preview import ( NODE_CLASS_MAPPINGS as _image_class_mappings, NODE_DISPLAY_NAME_MAPPINGS as _image_display_mappings, ) +from .fast_saver import ( + NODE_CLASS_MAPPINGS as _saver_class_mappings, + NODE_DISPLAY_NAME_MAPPINGS as _saver_display_mappings, +) -NODE_CLASS_MAPPINGS = {**_json_class_mappings, **_string_class_mappings, **_image_class_mappings} -NODE_DISPLAY_NAME_MAPPINGS = {**_json_display_mappings, **_string_display_mappings, **_image_display_mappings} +NODE_CLASS_MAPPINGS = {**_json_class_mappings, **_string_class_mappings, **_image_class_mappings, **_saver_class_mappings} +NODE_DISPLAY_NAME_MAPPINGS = {**_json_display_mappings, **_string_display_mappings, **_image_display_mappings, **_saver_display_mappings} WEB_DIRECTORY = "./web" diff --git a/fast_saver.py b/fast_saver.py new file mode 100644 index 0000000..4805983 --- /dev/null +++ b/fast_saver.py @@ -0,0 +1,436 @@ +import os +import platform +import torch +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo +import concurrent.futures +import re +import time +import glob +import json +import subprocess +import shutil +import stat +import tempfile +import urllib.request +import zipfile +import tarfile + +_NODE_DIR = os.path.dirname(os.path.abspath(__file__)) +_FFMPEG_DIR = os.path.join(_NODE_DIR, "ffmpeg_bin") + + +def _get_ffmpeg(): + """Find or download a ffmpeg binary. Search order: + 1. Bundled binary in this node's ffmpeg_bin/ folder + 2. imageio_ffmpeg (shipped by VideoHelperSuite) + 3. System PATH + 4. Auto-download a static build into ffmpeg_bin/ + """ + system = platform.system() + exe_name = "ffmpeg.exe" if system == "Windows" else "ffmpeg" + local_bin = os.path.join(_FFMPEG_DIR, exe_name) + + # 1. Already downloaded + if os.path.isfile(local_bin): + return local_bin + + # 2. imageio_ffmpeg + try: + import imageio_ffmpeg + path = imageio_ffmpeg.get_ffmpeg_exe() + if path and os.path.isfile(path): + return path + except Exception: + pass + + # 3. System PATH + system_bin = shutil.which("ffmpeg") + if system_bin: + return system_bin + + # 4. Auto-download static build + print("xx- FastSaver: ffmpeg not found. Downloading static build...") + os.makedirs(_FFMPEG_DIR, exist_ok=True) + + urls = { + ("Linux", "x86_64"): "https://johnvansickle.com/ffmpeg/releases/ffmpeg-release-amd64-static.tar.xz", + ("Linux", "aarch64"): "https://johnvansickle.com/ffmpeg/releases/ffmpeg-release-arm64-static.tar.xz", + ("Windows", "AMD64"): "https://www.gyan.dev/ffmpeg/builds/ffmpeg-release-essentials.zip", + } + + machine = platform.machine() + key = (system, machine) + url = urls.get(key) + if not url: + raise RuntimeError( + f"No automatic ffmpeg download available for {system}/{machine}. " + f"Please install ffmpeg manually or place a binary in {_FFMPEG_DIR}" + ) + + archive_path = os.path.join(_FFMPEG_DIR, "ffmpeg_archive") + try: + urllib.request.urlretrieve(url, archive_path) + + if url.endswith(".tar.xz"): + with tarfile.open(archive_path, "r:xz") as tar: + for member in tar.getmembers(): + if member.name.endswith("/ffmpeg") and member.isfile(): + member.name = exe_name + tar.extract(member, _FFMPEG_DIR) + break + elif url.endswith(".zip"): + with zipfile.ZipFile(archive_path, "r") as zf: + for name in zf.namelist(): + if name.endswith("bin/ffmpeg.exe"): + data = zf.read(name) + with open(local_bin, "wb") as f: + f.write(data) + break + finally: + if os.path.exists(archive_path): + os.remove(archive_path) + + if not os.path.isfile(local_bin): + raise RuntimeError(f"Failed to extract ffmpeg to {local_bin}") + + # Make executable on Linux/Mac + if system != "Windows": + os.chmod(local_bin, os.stat(local_bin).st_mode | stat.S_IEXEC) + + print(f"xx- FastSaver: ffmpeg downloaded to {local_bin}") + return local_bin + +class FastAbsoluteSaver: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", ), + "output_path": ("STRING", {"default": "D:\\Datasets\\Sharp_Output"}), + "filename_prefix": ("STRING", {"default": "frame"}), + + # --- FORMAT SWITCH --- + "save_format": (["png", "webp", "mp4", "webm"], ), + + # --- NAMING CONTROL --- + "use_timestamp": ("BOOLEAN", {"default": False, "label": "Add Timestamp (Unique)"}), + "auto_increment": ("BOOLEAN", {"default": True, "label": "Auto-Increment Counter (Scan Folder)"}), + "counter_digits": ("INT", {"default": 4, "min": 1, "max": 12, "step": 1, "label": "Number Padding (000X)"}), + "filename_with_score": ("BOOLEAN", {"default": False, "label": "Append Score to Filename"}), + + # --- METADATA & WORKFLOW --- + "metadata_key": ("STRING", {"default": "sharpness_score"}), + "save_workflow_metadata": ("BOOLEAN", {"default": False, "label": "Save ComfyUI Workflow (Graph)"}), + "save_metadata_png": ("BOOLEAN", {"default": False, "label": "Embed Workflow in PNG (sidecar or first file)"}), + + # --- PERFORMANCE --- + "max_threads": ("INT", {"default": 0, "min": 0, "max": 128, "step": 1, "label": "Max Threads (0=Auto)"}), + + # --- WEBP SPECIFIC --- + "webp_lossless": ("BOOLEAN", {"default": True, "label": "WebP Lossless"}), + "webp_quality": ("INT", {"default": 100, "min": 0, "max": 100, "step": 1}), + "webp_method": ("INT", {"default": 4, "min": 0, "max": 6, "step": 1}), + + # --- VIDEO SPECIFIC --- + "video_fps": ("INT", {"default": 24, "min": 1, "max": 120, "step": 1, "label": "Video FPS"}), + "video_crf": ("INT", {"default": 18, "min": 0, "max": 51, "step": 1, "label": "Video CRF (0=Lossless, 51=Worst)"}), + "video_pixel_format": (["yuv420p", "yuv444p"], {"label": "Pixel Format"}), + }, + "optional": { + "scores_info": ("STRING", {"forceInput": True}), + }, + # Hidden inputs used to capture the workflow graph + "hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"}, + } + + RETURN_TYPES = () + FUNCTION = "save_images_fast" + OUTPUT_NODE = True + CATEGORY = "utils/io" + + def parse_info(self, info_str, batch_size): + if not info_str: + # No scores connected - return None for scores to indicate "not provided" + return ([0]*batch_size, None) + matches = re.findall(r"F:(\d+).*?Score:\s*(\d+(\.\d+)?)", info_str) + frames = [] + scores = [] + for m in matches: + try: + frames.append(int(m[0])) + scores.append(float(m[1])) + except ValueError: + pass + if len(frames) < batch_size: + missing = batch_size - len(frames) + frames.extend([0] * missing) + scores.extend([0.0] * missing) + return frames[:batch_size], scores[:batch_size] + + def get_start_index(self, output_path, prefix): + # Scans the directory ONCE to find the highest existing number. + print(f"xx- FastSaver: Scanning folder for existing '{prefix}' files...") + files = glob.glob(os.path.join(output_path, f"{prefix}*.*")) + + max_idx = 0 + # Check specifically for prefix_NUMBER pattern to avoid confusing timestamps + pattern = re.compile(rf"{re.escape(prefix)}_(\d+)") + + for f in files: + fname = os.path.basename(f) + match = pattern.match(fname) + if match: + try: + val = int(match.group(1)) + if val > max_idx: + max_idx = val + except ValueError: + continue + print(f"xx- FastSaver: Found highest index {max_idx}. Starting at {max_idx + 1}") + return max_idx + 1 + + def _save_sidecar_png(self, img, png_path, score, key_name, prompt_data, extra_data): + """Save a metadata-carrying PNG sidecar next to a non-PNG export.""" + meta = PngInfo() + if score is not None: + meta.add_text(key_name, str(score)) + meta.add_text("software", "ComfyUI_FastAbsoluteSaver") + workflow_json = json.dumps(extra_data.get("workflow", {})) if extra_data else "{}" + prompt_json = json.dumps(prompt_data) if prompt_data else "{}" + meta.add_text("prompt", prompt_json) + meta.add_text("workflow", workflow_json) + img.save(png_path, format="PNG", pnginfo=meta, compress_level=1) + + def save_single_image(self, img_array, full_path, score, key_name, fmt, lossless, quality, method, + save_workflow, prompt_data, extra_data, force_png_metadata=False): + try: + img = Image.fromarray(img_array) + + # --- METADATA PREPARATION --- + meta_png = PngInfo() + + if fmt == "png": + if score is not None: + meta_png.add_text(key_name, str(score)) + meta_png.add_text("software", "ComfyUI_Parallel_Node") + + # ComfyUI Workflow Metadata (if requested, or forced via sidecar toggle for first PNG) + if (save_workflow or force_png_metadata) and fmt == "png": + workflow_json = json.dumps(extra_data.get("workflow", {})) if extra_data else "{}" + prompt_json = json.dumps(prompt_data) if prompt_data else "{}" + meta_png.add_text("prompt", prompt_json) + meta_png.add_text("workflow", workflow_json) + + # --- SAVING --- + if fmt == "png": + img.save(full_path, format="PNG", pnginfo=meta_png, compress_level=1) + + elif fmt == "webp": + img.save(full_path, format="WEBP", lossless=lossless, quality=quality, method=method) + + return True + except Exception as e: + print(f"xx- Error saving {full_path}: {e}") + return False + + def save_video(self, frames_np, output_path, filename_prefix, use_timestamp, fps, crf, pixel_format, video_format, + scores_list=None, metadata_key="sharpness_score", save_workflow=False, prompt_data=None, extra_data=None): + """Save image batch as a video file using ffmpeg. frames_np is a list/array of uint8 numpy arrays.""" + ffmpeg_path = _get_ffmpeg() + + ts_str = f"_{int(time.time())}" if use_timestamp else "" + ext = ".mp4" if video_format == "mp4" else ".webm" + out_file = os.path.join(output_path, f"{filename_prefix}{ts_str}{ext}") + + batch_size = len(frames_np) + h, w = frames_np[0].shape[0], frames_np[0].shape[1] + + # --- BUILD METADATA FILE (avoids arg-too-long for large workflows) --- + meta_lines = [";FFMETADATA1"] + meta_lines.append("software=ComfyUI_FastAbsoluteSaver") + + if scores_list: + avg_score = sum(scores_list) / len(scores_list) + meta_lines.append(f"{metadata_key}_avg={avg_score:.2f}") + meta_lines.append(f"{metadata_key}_all={','.join(f'{s:.2f}' for s in scores_list)}") + + if save_workflow: + if prompt_data: + # Escape ffmetadata special chars: =, ;, #, \ and newlines + prompt_str = json.dumps(prompt_data).replace("\\", "\\\\").replace("=", "\\=").replace(";", "\\;").replace("#", "\\#").replace("\n", "\\\n") + meta_lines.append(f"prompt={prompt_str}") + if extra_data: + workflow = extra_data.get("workflow", {}) + if workflow: + workflow_str = json.dumps(workflow).replace("\\", "\\\\").replace("=", "\\=").replace(";", "\\;").replace("#", "\\#").replace("\n", "\\\n") + meta_lines.append(f"workflow={workflow_str}") + + self._meta_tmpfile = tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False, encoding="utf-8") + self._meta_tmpfile.write("\n".join(meta_lines)) + self._meta_tmpfile.close() + meta_file = self._meta_tmpfile.name + + if video_format == "mp4": + codec = "libx264" + cmd = [ + ffmpeg_path, "-y", + "-f", "rawvideo", "-pix_fmt", "rgb24", + "-s", f"{w}x{h}", "-r", str(fps), + "-i", "-", + "-i", meta_file, "-map_metadata", "1", + "-c:v", codec, "-crf", str(crf), + "-pix_fmt", pixel_format, + "-movflags", "+faststart", + out_file + ] + else: # webm + codec = "libvpx-vp9" + cmd = [ + ffmpeg_path, "-y", + "-f", "rawvideo", "-pix_fmt", "rgb24", + "-s", f"{w}x{h}", "-r", str(fps), + "-i", "-", + "-i", meta_file, "-map_metadata", "1", + "-c:v", codec, "-crf", str(crf), "-b:v", "0", + "-pix_fmt", pixel_format, + out_file + ] + + print(f"xx- FastSaver: Encoding {batch_size} frames to {out_file} ({codec}, crf={crf}, {fps}fps)...") + + proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE) + try: + for frame in frames_np: + proc.stdin.write(frame.tobytes()) + proc.stdin.close() + except BrokenPipeError: + pass + stderr = proc.stderr.read() + proc.wait() + + # Clean up metadata temp file + try: + os.remove(meta_file) + except OSError: + pass + + if proc.returncode != 0: + raise RuntimeError(f"ffmpeg failed: {stderr.decode()}") + + print(f"xx- FastSaver: Video saved to {out_file}") + return out_file + + def save_images_fast(self, images, output_path, filename_prefix, save_format, use_timestamp, auto_increment, counter_digits, + max_threads, filename_with_score, metadata_key, save_workflow_metadata, save_metadata_png, + webp_lossless, webp_quality, webp_method, + video_fps, video_crf, video_pixel_format, + scores_info=None, prompt=None, extra_pnginfo=None): + + output_path = output_path.strip('"') + if not os.path.exists(output_path): + try: + os.makedirs(output_path, exist_ok=True) + except OSError: + raise ValueError(f"Could not create directory: {output_path}") + + if images is None or len(images) == 0: + raise ValueError("No images provided to FastAbsoluteSaver.") + + # Batch GPU->CPU transfer once (avoids N individual sync transfers) + images_np = (255.0 * images.cpu().numpy()).clip(0, 255).astype(np.uint8) + batch_size = len(images_np) + + # --- VIDEO PATH (check early, before image-specific logic) --- + if save_format in ("mp4", "webm"): + _, scores_list = self.parse_info(scores_info, batch_size) + out_file = self.save_video(images_np, output_path, filename_prefix, use_timestamp, + video_fps, video_crf, video_pixel_format, save_format, + scores_list=scores_list, metadata_key=metadata_key, + save_workflow=save_workflow_metadata, prompt_data=prompt, + extra_data=extra_pnginfo) + # Save metadata sidecar PNG next to the video file + if save_metadata_png: + png_path = os.path.splitext(out_file)[0] + ".png" + img = Image.fromarray(images_np[0]) + avg_score = (sum(scores_list) / len(scores_list)) if scores_list else None + self._save_sidecar_png(img, png_path, avg_score, metadata_key, prompt, extra_pnginfo) + print(f"xx- FastSaver: Metadata PNG sidecar saved to {png_path}") + return {"ui": {"images": []}} + + if max_threads == 0: + max_threads = os.cpu_count() or 4 + + frame_indices, scores_list = self.parse_info(scores_info, batch_size) + + # --- INDEX LOGIC --- + start_counter = 0 + using_real_frames = any(idx > 0 for idx in frame_indices) + + if auto_increment and not use_timestamp and not using_real_frames: + start_counter = self.get_start_index(output_path, filename_prefix) + + ts_str = f"_{int(time.time())}" if use_timestamp else "" + + print(f"xx- FastSaver: Saving {batch_size} images to {output_path}...") + + first_image_path = None + + with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor: + futures = [] + + for i, img_array in enumerate(images_np): + real_frame_num = frame_indices[i] + current_score = scores_list[i] if scores_list else None + + if real_frame_num > 0: + number_part = real_frame_num + else: + number_part = start_counter + i + + fmt_str = f"{{:0{counter_digits}d}}" + number_str = fmt_str.format(number_part) + + base_name = f"{filename_prefix}{ts_str}_{number_str}" + + if filename_with_score and current_score is not None: + base_name += f"_{int(current_score)}" + + ext = ".webp" if save_format == "webp" else ".png" + full_path = os.path.join(output_path, f"{base_name}{ext}") + + if i == 0: + first_image_path = full_path + + # For PNG sequences: embed workflow metadata in the first file only + force_meta = (save_metadata_png and save_format == "png" and i == 0) + + futures.append(executor.submit( + self.save_single_image, + img_array, full_path, current_score, metadata_key, + save_format, webp_lossless, webp_quality, webp_method, + save_workflow_metadata, prompt, extra_pnginfo, + force_png_metadata=force_meta + )) + + concurrent.futures.wait(futures) + + # Save a single metadata sidecar PNG using the first image (skip for PNG sequences - handled above) + if save_metadata_png and save_format != "png" and first_image_path is not None: + png_path = os.path.splitext(first_image_path)[0] + ".png" + img = Image.fromarray(images_np[0]) + first_score = scores_list[0] if scores_list else None + self._save_sidecar_png(img, png_path, first_score, metadata_key, prompt, extra_pnginfo) + print(f"xx- FastSaver: Metadata PNG sidecar saved to {png_path}") + + return {"ui": {"images": []}} + + +NODE_CLASS_MAPPINGS = { + "FastAbsoluteSaver": FastAbsoluteSaver, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "FastAbsoluteSaver": "Fast Absolute Saver (Metadata)", +}