From 178247c79fbcd101d0db9cf702990854852eee84 Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Mon, 19 Jan 2026 23:04:07 +0100 Subject: [PATCH 1/4] Add fast_saver.py --- fast_saver.py | 120 ++++++++++++++++++++++++++++++++++++++++++++++++++ 1 file changed, 120 insertions(+) create mode 100644 fast_saver.py diff --git a/fast_saver.py b/fast_saver.py new file mode 100644 index 0000000..9863aea --- /dev/null +++ b/fast_saver.py @@ -0,0 +1,120 @@ +import os +import torch +import numpy as np +from PIL import Image +from PIL.PngImagePlugin import PngInfo +import concurrent.futures +import re + +class FastAbsoluteSaver: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "images": ("IMAGE", ), + "output_path": ("STRING", {"default": "D:\\Datasets\\Sharp_Output"}), + "filename_prefix": ("STRING", {"default": "frame"}), + }, + "optional": { + # We take the string output from your Parallel Loader here + "scores_info": ("STRING", {"forceInput": True}), + } + } + + RETURN_TYPES = () + FUNCTION = "save_images_fast" + OUTPUT_NODE = True + CATEGORY = "BetaHelper/IO" + + def parse_scores(self, scores_str, batch_size): + """ + Parses the string "F:10 (Score:500), F:12 (Score:800)..." into a list of floats. + If inputs don't match, returns a list of 0.0. + """ + if not scores_str: + return [0.0] * batch_size + + # Regex to find 'Score:NUMBER' or just numbers + # Matches your specific format: (Score: 123) + patterns = re.findall(r"Score:(\d+(\.\d+)?)", scores_str) + + scores = [] + for match in patterns: + # match is a tuple due to the group inside regex, index 0 is the full number + try: + scores.append(float(match[0])) + except ValueError: + scores.append(0.0) + + # Validation: If we found more or fewer scores than images, pad or truncate + if len(scores) < batch_size: + scores.extend([0.0] * (batch_size - len(scores))) + return scores[:batch_size] + + def save_single_image(self, tensor_img, full_path, score): + """Worker function to save one image with metadata""" + try: + # 1. Convert Tensor to Pillow + array = 255. * tensor_img.cpu().numpy() + img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8)) + + # 2. Add Metadata + metadata = PngInfo() + metadata.add_text("sharpness_score", str(score)) + # You can add more keys here if needed + metadata.add_text("software", "ComfyUI_Parallel_Node") + + # 3. Save (Optimized) + img.save(full_path, pnginfo=metadata, compress_level=1) + # compress_level=1 is FAST. Default is 6 (slow). 0 is uncompressed (huge files). + + return True + except Exception as e: + print(f"xx- Error saving {full_path}: {e}") + return False + + def save_images_fast(self, images, output_path, filename_prefix, scores_info=None): + + # 1. Clean Path + 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}") + + # 2. Parse Scores + batch_size = len(images) + scores_list = self.parse_scores(scores_info, batch_size) + + # 3. Parallel Saving + # We use a ThreadPool to save files concurrently. + # This saturates the SSD write speed, mimicking VHS performance. + print(f"xx- FastSaver: Saving {batch_size} images to {output_path}...") + + with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor: + futures = [] + + for i, img_tensor in enumerate(images): + # Construct filename: prefix_00001.png + # We use a unique counter or just the batch index + # Ideally, we should use the Frame Index if we can extract it, + # but for now we use simple batch increment to avoid overwriting. + + # If we want unique filenames based on existing files, it slows things down. + # We will assume the user manages folders or prefixes well. + import time + timestamp = int(time.time()) + fname = f"{filename_prefix}_{timestamp}_{i:05d}.png" + full_path = os.path.join(output_path, fname) + + # Submit to thread + futures.append(executor.submit(self.save_single_image, img_tensor, full_path, scores_list[i])) + + # Wait for all to finish + concurrent.futures.wait(futures) + + print("xx- FastSaver: Save Complete.") + + # Return nothing to UI to prevent Lag + return {"ui": {"images": []}} \ No newline at end of file From 44f3130a1519a4f0696512ebc8d6448b252975fc Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Mon, 19 Jan 2026 23:05:03 +0100 Subject: [PATCH 2/4] Update __init__.py --- __init__.py | 6 ++++-- 1 file changed, 4 insertions(+), 2 deletions(-) diff --git a/__init__.py b/__init__.py index 0401603..7c44eca 100644 --- a/__init__.py +++ b/__init__.py @@ -4,13 +4,15 @@ from .parallel_loader import ParallelSharpnessLoader NODE_CLASS_MAPPINGS = { "SharpnessAnalyzer": SharpnessAnalyzer, "SharpFrameSelector": SharpFrameSelector, - "ParallelSharpnessLoader": ParallelSharpnessLoader + "ParallelSharpnessLoader": ParallelSharpnessLoader, + "FastAbsoluteSaver": FastAbsoluteSaver } NODE_DISPLAY_NAME_MAPPINGS = { "SharpnessAnalyzer": "1. Sharpness Analyzer", "SharpFrameSelector": "2. Sharp Frame Selector", - "ParallelSharpnessLoader": "3. Parallel Video Loader (Sharpness)" + "ParallelSharpnessLoader": "3. Parallel Video Loader (Sharpness)", + "FastAbsoluteSaver": "Fast Absolute Saver (Metadata)" } __all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] \ No newline at end of file From 099ce948ae9de4ec114e7bb60f839f4d5d5022bd Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Mon, 19 Jan 2026 23:09:21 +0100 Subject: [PATCH 3/4] Update __init__.py --- __init__.py | 1 + 1 file changed, 1 insertion(+) diff --git a/__init__.py b/__init__.py index 7c44eca..df11bdd 100644 --- a/__init__.py +++ b/__init__.py @@ -1,5 +1,6 @@ from .sharp_node import SharpnessAnalyzer, SharpFrameSelector from .parallel_loader import ParallelSharpnessLoader +from .fast_saver import FastAbsoluteSaver # <--- Added this missing import NODE_CLASS_MAPPINGS = { "SharpnessAnalyzer": SharpnessAnalyzer, From 24a59a6da2a93ea5b2e20f72a0c6df5f81559a8a Mon Sep 17 00:00:00 2001 From: Ethanfel Date: Mon, 19 Jan 2026 23:44:40 +0100 Subject: [PATCH 4/4] Update fast_saver.py --- fast_saver.py | 56 +++++++++++++++++++++------------------------------ 1 file changed, 23 insertions(+), 33 deletions(-) diff --git a/fast_saver.py b/fast_saver.py index 9863aea..d5b05a9 100644 --- a/fast_saver.py +++ b/fast_saver.py @@ -14,9 +14,10 @@ class FastAbsoluteSaver: "images": ("IMAGE", ), "output_path": ("STRING", {"default": "D:\\Datasets\\Sharp_Output"}), "filename_prefix": ("STRING", {"default": "frame"}), + # NEW: User can define the metadata key name + "metadata_key": ("STRING", {"default": "sharpness_score", "label": "Metadata Key Name"}), }, "optional": { - # We take the string output from your Parallel Loader here "scores_info": ("STRING", {"forceInput": True}), } } @@ -28,52 +29,52 @@ class FastAbsoluteSaver: def parse_scores(self, scores_str, batch_size): """ - Parses the string "F:10 (Score:500), F:12 (Score:800)..." into a list of floats. - If inputs don't match, returns a list of 0.0. + Parses the string "F:10 (Score:500)..." into a list of floats. + Robust to spaces: handles "Score:500" and "Score: 500" """ if not scores_str: return [0.0] * batch_size - # Regex to find 'Score:NUMBER' or just numbers - # Matches your specific format: (Score: 123) - patterns = re.findall(r"Score:(\d+(\.\d+)?)", scores_str) + # Regex explanation: + # Score:\s* -> Matches "Score:" followed by optional spaces + # (\d+(\.\d+)?) -> Matches integer or float (Capture Group 1) + patterns = re.findall(r"Score:\s*(\d+(\.\d+)?)", scores_str) scores = [] for match in patterns: - # match is a tuple due to the group inside regex, index 0 is the full number try: scores.append(float(match[0])) except ValueError: scores.append(0.0) - # Validation: If we found more or fewer scores than images, pad or truncate + # Fill missing scores with 0.0 if batch size mismatches if len(scores) < batch_size: scores.extend([0.0] * (batch_size - len(scores))) return scores[:batch_size] - def save_single_image(self, tensor_img, full_path, score): + def save_single_image(self, tensor_img, full_path, score, key_name): """Worker function to save one image with metadata""" try: - # 1. Convert Tensor to Pillow array = 255. * tensor_img.cpu().numpy() img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8)) - # 2. Add Metadata metadata = PngInfo() - metadata.add_text("sharpness_score", str(score)) - # You can add more keys here if needed + + # Use the user-defined key. + # If you want to force no spaces, uncomment the line below: + # key_name = key_name.replace(" ", "_") + + metadata.add_text(key_name, str(score)) metadata.add_text("software", "ComfyUI_Parallel_Node") - # 3. Save (Optimized) + # compress_level=1 is fast. img.save(full_path, pnginfo=metadata, compress_level=1) - # compress_level=1 is FAST. Default is 6 (slow). 0 is uncompressed (huge files). - return True except Exception as e: print(f"xx- Error saving {full_path}: {e}") return False - def save_images_fast(self, images, output_path, filename_prefix, scores_info=None): + def save_images_fast(self, images, output_path, filename_prefix, metadata_key, scores_info=None): # 1. Clean Path output_path = output_path.strip('"') @@ -87,34 +88,23 @@ class FastAbsoluteSaver: batch_size = len(images) scores_list = self.parse_scores(scores_info, batch_size) - # 3. Parallel Saving - # We use a ThreadPool to save files concurrently. - # This saturates the SSD write speed, mimicking VHS performance. print(f"xx- FastSaver: Saving {batch_size} images to {output_path}...") + # 3. Parallel Saving with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor: futures = [] for i, img_tensor in enumerate(images): - # Construct filename: prefix_00001.png - # We use a unique counter or just the batch index - # Ideally, we should use the Frame Index if we can extract it, - # but for now we use simple batch increment to avoid overwriting. - - # If we want unique filenames based on existing files, it slows things down. - # We will assume the user manages folders or prefixes well. import time timestamp = int(time.time()) - fname = f"{filename_prefix}_{timestamp}_{i:05d}.png" + # Added index `i` to filename to ensure uniqueness in same batch + fname = f"{filename_prefix}_{timestamp}_{i:03d}.png" full_path = os.path.join(output_path, fname) - # Submit to thread - futures.append(executor.submit(self.save_single_image, img_tensor, full_path, scores_list[i])) + # Pass the metadata_key to the worker + futures.append(executor.submit(self.save_single_image, img_tensor, full_path, scores_list[i], metadata_key)) - # Wait for all to finish concurrent.futures.wait(futures) print("xx- FastSaver: Save Complete.") - - # Return nothing to UI to prevent Lag return {"ui": {"images": []}} \ No newline at end of file