Update fast_saver.py

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": []}}