diff --git a/fast_saver.py b/fast_saver.py index ef7850a..dbd3651 100644 --- a/fast_saver.py +++ b/fast_saver.py @@ -6,6 +6,7 @@ from PIL.PngImagePlugin import PngInfo import concurrent.futures import re import time +import glob class FastAbsoluteSaver: @classmethod @@ -20,8 +21,9 @@ class FastAbsoluteSaver: "save_format": (["png", "webp"], ), # --- NAMING CONTROL --- - "use_timestamp": ("BOOLEAN", {"default": True, "label": "Add Timestamp (Unique)"}), - "counter_digits": ("INT", {"default": 3, "min": 1, "max": 12, "step": 1, "label": "Number Padding (00X)"}), + "use_timestamp": ("BOOLEAN", {"default": False, "label": "Add Timestamp (Unique)"}), + "auto_increment": ("BOOLEAN", {"default": True, "label": "Auto-Increment Counter (Scan Folder)"}), + "counter_digits": ("INT", {"default": 4, "min": 1, "max": 12, "step": 1, "label": "Number Padding (000X)"}), "filename_with_score": ("BOOLEAN", {"default": False, "label": "Append Score to Filename"}), # --- PERFORMANCE --- @@ -48,7 +50,6 @@ class FastAbsoluteSaver: def parse_info(self, info_str, batch_size): 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 = [] @@ -58,14 +59,52 @@ class FastAbsoluteSaver: 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 get_start_index(self, output_path, prefix): + """ + Scans the directory ONCE to find the highest existing number. + Returns the next available index. + """ + print(f"xx- FastSaver: Scanning folder for existing '{prefix}' files...") + # Get all files starting with prefix + files = glob.glob(os.path.join(output_path, f"{prefix}*.*")) + + max_idx = 0 + pattern = re.compile(rf"{re.escape(prefix)}_?(\d+)") + + for f in files: + fname = os.path.basename(f) + # Try to match the last number group + match = pattern.search(fname) + if match: + try: + # We look for the last numeric group in the filename + # This logic handles frame_001.png or frame_001_score.png + groups = re.findall(r"(\d+)", fname) + if groups: + # Usually the counter is the first or second number + # Simplified: Just grab the first number found after prefix + val = int(groups[-1] if len(groups) == 1 else groups[0]) + # If filename has timestamp, this logic gets tricky, + # but auto_increment usually implies NO timestamp. + + # Better approach: Check specifically for prefix_NUMBER + clean_match = re.match(rf"{re.escape(prefix)}_(\d+)", fname) + if clean_match: + val = int(clean_match.group(1)) + if val > max_idx: + max_idx = val + except ValueError: + continue + + print(f"xx- FastSaver: Found highest index {max_idx}. Starting at {max_idx + 1}") + return max_idx + 1 + def save_single_image(self, tensor_img, full_path, score, key_name, fmt, lossless, quality, method): try: array = 255. * tensor_img.cpu().numpy() @@ -76,19 +115,14 @@ class FastAbsoluteSaver: metadata.add_text(key_name, str(score)) metadata.add_text("software", "ComfyUI_Parallel_Node") img.save(full_path, format="PNG", pnginfo=metadata, compress_level=1) - elif fmt == "webp": - img.save(full_path, format="WEBP", - lossless=lossless, - quality=quality, - method=method) - + 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, save_format, use_timestamp, counter_digits, + def save_images_fast(self, images, output_path, filename_prefix, save_format, use_timestamp, auto_increment, counter_digits, max_threads, filename_with_score, metadata_key, webp_lossless, webp_quality, webp_method, scores_info=None): output_path = output_path.strip('"') @@ -104,7 +138,17 @@ class FastAbsoluteSaver: batch_size = len(images) frame_indices, scores_list = self.parse_info(scores_info, batch_size) - # Pre-calculate timestamp once for the whole batch if needed + # --- INDEX LOGIC --- + start_counter = 0 + # Only scan if: + # 1. User wants Auto-Increment + # 2. We are NOT using Timestamps (which are naturally unique) + # 3. We are NOT using Frame Numbers (because overwriting frame 100 with frame 100 is usually desired) + using_real_frames = any(idx > 0 for idx in frame_indices) + + if auto_increment and not use_timestamp and not using_real_frames: + start_counter = self.get_start_index(output_path, filename_prefix) + ts_str = f"_{int(time.time())}" if use_timestamp else "" print(f"xx- FastSaver: Saving {batch_size} images to {output_path}...") @@ -116,20 +160,17 @@ class FastAbsoluteSaver: real_frame_num = frame_indices[i] current_score = scores_list[i] - # Logic: If we have real frame numbers (from Loader), use them. - # If NOT (or if frame is 0), use the loop index 'i' (0, 1, 2...) + # Priority: + # 1. Real Video Frame (from Loader) + # 2. Auto-Increment Counter (Start + i) if real_frame_num > 0: number_part = real_frame_num else: - number_part = i + number_part = start_counter + i - # Format string using dynamic padding size (e.g. :05d) fmt_str = f"{{:0{counter_digits}d}}" number_str = fmt_str.format(number_part) - # Construct Name: prefix + timestamp + number - # Case 1: frame_173000_001.png (Timestamp ON) - # Case 2: frame_001.png (Timestamp OFF) base_name = f"{filename_prefix}{ts_str}_{number_str}" if filename_with_score: