Update fast_saver.py

fast_saver.py

@@ -1,11 +1,12 @@
 import os
 import torch
 import numpy as np
-from PIL import Image
+from PIL import Image, ExifTags
 from PIL.PngImagePlugin import PngInfo
 import concurrent.futures
 import re
 import time
+import json
 
 class FastAbsoluteSaver:
     @classmethod
@@ -15,9 +16,18 @@ class FastAbsoluteSaver:
                 "images": ("IMAGE", ),
                 "output_path": ("STRING", {"default": "D:\\Datasets\\Sharp_Output"}),
                 "filename_prefix": ("STRING", {"default": "frame"}),
-                "metadata_key": ("STRING", {"default": "sharpness_score"}),
+                
-                # NEW: Boolean Switch
+                # --- FORMAT SWITCH ---
+                "save_format": (["png", "webp"], ),
+                
+                # --- 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_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)"}),
             },
             "optional": {
                 "scores_info": ("STRING", {"forceInput": True}),
@@ -30,17 +40,12 @@ class FastAbsoluteSaver:
     CATEGORY = "BetaHelper/IO"
 
     def parse_info(self, info_str, batch_size):
-        """
-        Extracts both Frame Indices AND Scores.
-        """
         if not info_str:
             return ([0]*batch_size, [0.0]*batch_size)
 
         matches = re.findall(r"F:(\d+).*?Score:\s*(\d+(\.\d+)?)", info_str)
-        
         frames = []
         scores = []
-        
         for m in matches:
             try:
                 frames.append(int(m[0]))       
@@ -55,22 +60,68 @@ class FastAbsoluteSaver:
             
         return frames[:batch_size], scores[:batch_size]
 
-    def save_single_image(self, tensor_img, full_path, score, key_name):
+    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))
             
-            metadata = PngInfo()
+            if fmt == "png":
-            metadata.add_text(key_name, str(score))
+                # PNG METADATA (Robust)
-            metadata.add_text("software", "ComfyUI_Parallel_Node")
+                metadata = PngInfo()
-
+                metadata.add_text(key_name, str(score))
-            img.save(full_path, pnginfo=metadata, compress_level=1) 
+                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, 
+                         method=method) 
+            
             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, filename_with_score, scores_info=None):
+    def save_images_fast(self, images, output_path, filename_prefix, save_format, 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):
@@ -82,7 +133,7 @@ class FastAbsoluteSaver:
         batch_size = len(images)
         frame_indices, scores_list = self.parse_info(scores_info, batch_size)
 
-        print(f"xx- FastSaver: Saving {batch_size} images to {output_path}...")
+        print(f"xx- FastSaver: Saving {batch_size} images ({save_format}) to {output_path}...")
         
         with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
             futures = []
@@ -92,21 +143,31 @@ class FastAbsoluteSaver:
                 real_frame_num = frame_indices[i]
                 current_score = scores_list[i]
                 
-                # 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"
+                # Append correct extension
+                ext = ".webp" if save_format == "webp" else ".png"
+                fname = f"{base_name}{ext}"
                 full_path = os.path.join(output_path, fname)
                 
-                futures.append(executor.submit(self.save_single_image, img_tensor, full_path, current_score, metadata_key))
+                # Submit
+                futures.append(executor.submit(
+                    self.save_single_image, 
+                    img_tensor, 
+                    full_path, 
+                    current_score, 
+                    metadata_key,
+                    save_format,    # fmt
+                    webp_lossless,  # lossless
+                    webp_quality,   # quality
+                    webp_method     # method
+                ))
 
             concurrent.futures.wait(futures)