Update fast_saver.py

fast_saver.py

@@ -5,6 +5,7 @@ from PIL import Image
 from PIL.PngImagePlugin import PngInfo
 import concurrent.futures
 import re
+import time
 
 class FastAbsoluteSaver:
     @classmethod
@@ -14,8 +15,9 @@ 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"}),
-                "metadata_key": ("STRING", {"default": "sharpness_score", "label": "Metadata Key Name"}),
+                # NEW: Boolean Switch
+                "filename_with_score": ("BOOLEAN", {"default": False, "label": "Append Score to Filename"}),
             },
             "optional": {
                 "scores_info": ("STRING", {"forceInput": True}),
@@ -27,56 +29,49 @@ class FastAbsoluteSaver:
     OUTPUT_NODE = True
     CATEGORY = "BetaHelper/IO"
 
-    def parse_scores(self, scores_str, batch_size):
+    def parse_info(self, info_str, batch_size):
         """
-        Parses the string "F:10 (Score:500)..." into a list of floats.
+        Extracts both Frame Indices AND Scores.
-        Robust to spaces: handles "Score:500" and "Score: 500"
         """
-        if not scores_str:
+        if not info_str:
-            return [0.0] * batch_size
+            return ([0]*batch_size, [0.0]*batch_size)
 
-        # Regex explanation:
+        matches = re.findall(r"F:(\d+).*?Score:\s*(\d+(\.\d+)?)", info_str)
-        # 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)
         
+        frames = []
         scores = []
-        for match in patterns:
-            try:
-                scores.append(float(match[0]))
-            except ValueError:
-                scores.append(0.0)
         
-        # Fill missing scores with 0.0 if batch size mismatches
+        for m in matches:
-        if len(scores) < batch_size:
+            try:
-            scores.extend([0.0] * (batch_size - len(scores)))
+                frames.append(int(m[0]))       
-        return scores[:batch_size]
+                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 save_single_image(self, tensor_img, full_path, score, key_name):
-        """Worker function to save one image with metadata"""
         try:
             array = 255. * tensor_img.cpu().numpy()
             img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
             
             metadata = PngInfo()
-            
-            # 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")
 
-            # compress_level=1 is fast.
             img.save(full_path, pnginfo=metadata, compress_level=1) 
             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, metadata_key, scores_info=None):
+    def save_images_fast(self, images, output_path, filename_prefix, metadata_key, filename_with_score, scores_info=None):
         
-        # 1. Clean Path
         output_path = output_path.strip('"')
         if not os.path.exists(output_path):
             try:
@@ -84,27 +79,35 @@ class FastAbsoluteSaver:
             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)
+        frame_indices, scores_list = self.parse_info(scores_info, batch_size)
 
         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):
-                import time
+                
-                timestamp = int(time.time())
+                real_frame_num = frame_indices[i]
-                # Added index `i` to filename to ensure uniqueness in same batch
+                current_score = scores_list[i]
-                fname = f"{filename_prefix}_{timestamp}_{i:03d}.png"
+                
+                # BASE NAME: frame_001450
+                base_name = f"{filename_prefix}_{real_frame_num:06d}"
+
+                # OPTION: Append Score -> frame_001450_1500
+                if filename_with_score:
+                    base_name += f"_{int(current_score)}"
+
+                # FALLBACK for missing data
+                if real_frame_num == 0 and scores_info is None:
+                    base_name = f"{filename_prefix}_{int(time.time())}_{i:03d}"
+
+                fname = f"{base_name}.png"
                 full_path = os.path.join(output_path, fname)
                 
-                # Pass the metadata_key to the worker
+                futures.append(executor.submit(self.save_single_image, img_tensor, full_path, current_score, metadata_key))
-                futures.append(executor.submit(self.save_single_image, img_tensor, full_path, scores_list[i], metadata_key))
 
             concurrent.futures.wait(futures)
 
-        print("xx- FastSaver: Save Complete.")
         return {"ui": {"images": []}}