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