fast_saver.py

@@ -1,12 +1,11 @@
 import os
 import torch
 import numpy as np
-from PIL import Image, ExifTags
+from PIL import Image
 from PIL.PngImagePlugin import PngInfo
 import concurrent.futures
 import re
 import time
-import json
 
 class FastAbsoluteSaver:
     @classmethod
@@ -20,11 +19,14 @@ class FastAbsoluteSaver:
                 # --- FORMAT SWITCH ---
                 "save_format": (["png", "webp"], ),
                 
+                # --- PERFORMANCE ---
+                "max_threads": ("INT", {"default": 0, "min": 0, "max": 128, "step": 1, "label": "Max Threads (0=Auto)"}),
+
                 # --- COMMON OPTIONS ---
                 "filename_with_score": ("BOOLEAN", {"default": False, "label": "Append Score to Filename"}),
                 "metadata_key": ("STRING", {"default": "sharpness_score"}),
 
-                # --- WEBP SPECIFIC (-z 6 -q 100 -lossless) ---
+                # --- WEBP SPECIFIC ---
                 "webp_lossless": ("BOOLEAN", {"default": True, "label": "WebP Lossless"}),
                 "webp_quality": ("INT", {"default": 100, "min": 0, "max": 100, "step": 1, "label": "WebP Quality (-q)"}),
                 "webp_method": ("INT", {"default": 4, "min": 0, "max": 6, "step": 1, "label": "WebP Compression (-z)"}),
@@ -60,56 +62,18 @@ class FastAbsoluteSaver:
             
         return frames[:batch_size], scores[:batch_size]
 
-    def get_webp_exif(self, key, value):
-        """
-        Creates a basic Exif header to store the score in UserComment.
-        ComfyUI standard metadata handling for WebP is complex, 
-        so we use a simple JSON dump inside the UserComment tag (ID 0x9286).
-        """
-        # Create a basic exif dict
-        exif_data = {
-            0x9286: f"{key}: {value}".encode("utf-8") # UserComment
-        }
-        
-        # Convert to bytes manually to avoid requiring 'piexif' library
-        # This is a minimal TIFF header structure for Exif.
-        # If this is too hacky, we can just skip metadata for WebP, 
-        # but this usually works for basic viewers.
-        
-        # ACTUALLY: Pillow's image.save(exif=...) expects raw bytes.
-        # Generating raw Exif bytes from scratch is error-prone.
-        # Simpler Strategy: We will create a fresh Image and modify its info.
-        
-        # Since generating raw Exif without a library is risky, 
-        # we will skip internal metadata for WebP in this "No-Dependency" version
-        # and rely on the filename. 
-        # *However*, if you strictly need it, we return None here and rely on filename
-        # unless you have 'piexif' installed.
-        return None
-
     def save_single_image(self, tensor_img, full_path, score, key_name, fmt, lossless, quality, method):
         try:
             array = 255. * tensor_img.cpu().numpy()
             img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
             
             if fmt == "png":
-                # PNG METADATA (Robust)
                 metadata = PngInfo()
                 metadata.add_text(key_name, str(score))
                 metadata.add_text("software", "ComfyUI_Parallel_Node")
-                # PNG uses compress_level (0-9). Level 1 is fastest.
                 img.save(full_path, format="PNG", pnginfo=metadata, compress_level=1)
             
             elif fmt == "webp":
-                # WEBP SAVING
-                # Pillow options map directly to cwebp parameters:
-                # method=6 -> -z 6 (Slowest, best compression)
-                # quality=100 -> -q 100
-                # lossless=True -> -lossless
-                
-                # Note: WebP metadata in Pillow is finicky. 
-                # We save purely visual data here. 
-                # The score is in the filename (if option selected).
                 img.save(full_path, format="WEBP", 
                          lossless=lossless, 
                          quality=quality, 
@@ -120,8 +84,8 @@ class FastAbsoluteSaver:
             print(f"xx- Error saving {full_path}: {e}")
             return False
 
-    def save_images_fast(self, images, output_path, filename_prefix, save_format, filename_with_score, metadata_key, 
+    def save_images_fast(self, images, output_path, filename_prefix, save_format, max_threads, 
-                         webp_lossless, webp_quality, webp_method, scores_info=None):
+                         filename_with_score, metadata_key, webp_lossless, webp_quality, webp_method, scores_info=None):
         
         output_path = output_path.strip('"')
         if not os.path.exists(output_path):
@@ -130,12 +94,20 @@ class FastAbsoluteSaver:
             except OSError:
                 raise ValueError(f"Could not create directory: {output_path}")
 
+        # --- AUTO-SCALING LOGIC ---
+        if max_threads == 0:
+            # os.cpu_count() returns None on some rare systems, so we default to 4 just in case
+            cpu_cores = os.cpu_count() or 4
+            # For WebP (CPU intensive), stick to core count. 
+            # For PNG (Disk intensive), we could technically go higher, but core count is safe.
+            max_threads = cpu_cores
+        
+        print(f"xx- FastSaver: Using {max_threads} Threads for saving.")
+
         batch_size = len(images)
         frame_indices, scores_list = self.parse_info(scores_info, batch_size)
 
-        print(f"xx- FastSaver: Saving {batch_size} images ({save_format}) to {output_path}...")
+        with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor:
-        
-        with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
             futures = []
             
             for i, img_tensor in enumerate(images):
@@ -151,22 +123,20 @@ class FastAbsoluteSaver:
                 if real_frame_num == 0 and scores_info is None:
                     base_name = f"{filename_prefix}_{int(time.time())}_{i:03d}"
 
-                # Append correct extension
                 ext = ".webp" if save_format == "webp" else ".png"
                 fname = f"{base_name}{ext}"
                 full_path = os.path.join(output_path, fname)
                 
-                # Submit
                 futures.append(executor.submit(
                     self.save_single_image, 
                     img_tensor, 
                     full_path, 
                     current_score, 
                     metadata_key,
-                    save_format,    # fmt
+                    save_format,
-                    webp_lossless,  # lossless
+                    webp_lossless,
-                    webp_quality,   # quality
+                    webp_quality,
-                    webp_method     # method
+                    webp_method
                 ))
 
             concurrent.futures.wait(futures)