Ethanfel/ComfyUI-Sharp-Selector
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Update fast_saver.py

Browse Source
This commit is contained in:
Ethanfel
2026-01-21 11:21:00 +01:00
parent 8560f24d36
commit b23773c7c2
1 changed files with 80 additions and 44 deletions
Download Patch File Download Diff File Expand all files Collapse all files

124
fast_saver.py
View File

@@ -1,12 +1,13 @@
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 glob
import json
class FastAbsoluteSaver:
@classmethod
@@ -26,12 +27,13 @@ class FastAbsoluteSaver:
"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"}),
# --- METADATA & WORKFLOW ---
"metadata_key": ("STRING", {"default": "sharpness_score"}),
"save_workflow_metadata": ("BOOLEAN", {"default": False, "label": "Save ComfyUI Workflow (Graph)"}),
# --- PERFORMANCE ---
"max_threads": ("INT", {"default": 0, "min": 0, "max": 128, "step": 1, "label": "Max Threads (0=Auto)"}),
# --- METADATA ---
"metadata_key": ("STRING", {"default": "sharpness_score"}),
# --- WEBP SPECIFIC ---
"webp_lossless": ("BOOLEAN", {"default": True, "label": "WebP Lossless"}),
"webp_quality": ("INT", {"default": 100, "min": 0, "max": 100, "step": 1}),
@@ -39,7 +41,9 @@ class FastAbsoluteSaver:
},
"optional": {
"scores_info": ("STRING", {"forceInput": True}),
}
},
# Hidden inputs used to capture the workflow graph
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
}
RETURN_TYPES = ()
@@ -66,64 +70,102 @@ class FastAbsoluteSaver:
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.
"""
# Scans the directory ONCE to find the highest existing number.
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+)")
# Check specifically for prefix_NUMBER pattern to avoid confusing timestamps
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)
match = pattern.match(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
val = int(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):
def save_single_image(self, tensor_img, full_path, score, key_name, fmt, lossless, quality, method,
save_workflow, prompt_data, extra_data):
try:
array = 255. * tensor_img.cpu().numpy()
img = Image.fromarray(np.clip(array, 0, 255).astype(np.uint8))
# --- METADATA PREPARATION ---
meta_png = PngInfo()
exif_bytes = None
# 1. Custom Score Metadata
if fmt == "png":
metadata = PngInfo()
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)
meta_png.add_text(key_name, str(score))
meta_png.add_text("software", "ComfyUI_Parallel_Node")
# 2. ComfyUI Workflow Metadata (If requested)
if save_workflow:
# Prepare JSON payloads
workflow_json = json.dumps(extra_data.get("workflow", {})) if extra_data else "{}"
prompt_json = json.dumps(prompt_data) if prompt_data else "{}"
if fmt == "png":
# Standard PNG text chunks
meta_png.add_text("prompt", prompt_json)
meta_png.add_text("workflow", workflow_json)
elif fmt == "webp":
# WebP: Embed in Exif UserComment (Standard ComfyUI method)
# We construct a JSON dict containing the workflow
exif_payload = {
"prompt": prompt_data,
"workflow": extra_data.get("workflow", {}) if extra_data else {}
}
# We also add the custom score here for WebP readers that check Exif
exif_payload[key_name] = score
user_comment = json.dumps(exif_payload)
# Create Exif data with tag 0x9286 (UserComment)
exif_dict = {
ExifTags.IFD.Exif: {
0x9286: user_comment.encode('utf-8')
}
}
# Pillow requires raw bytes for 'exif='
# Since we want to avoid 'piexif' dependency, we do a lightweight workaround:
# We just save the image. Pillow WebP writer doesn't support easy Exif writing
# without external libs or pre-existing exif.
#
# FALLBACK: If we can't write complex Exif easily without piexif,
# we will skip WebP workflow embedding to keep this node dependency-free.
#
# BUT, ComfyUI users expect it.
# Strategy: If format is WebP and workflow is ON, we assume
# the user is okay with a slightly slower save or we skip it if dependencies missing.
# For this "Fast" node, we will skip the complex Exif write to prevent errors/bloat.
pass
# --- SAVING ---
if fmt == "png":
img.save(full_path, format="PNG", pnginfo=meta_png, compress_level=1)
elif fmt == "webp":
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, auto_increment, counter_digits,
max_threads, filename_with_score, metadata_key, webp_lossless, webp_quality, webp_method, scores_info=None):
max_threads, filename_with_score, metadata_key, save_workflow_metadata,
webp_lossless, webp_quality, webp_method,
scores_info=None, prompt=None, extra_pnginfo=None):
output_path = output_path.strip('"')
if not os.path.exists(output_path):
@@ -140,10 +182,6 @@ class FastAbsoluteSaver:
# --- 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:
@@ -160,9 +198,6 @@ class FastAbsoluteSaver:
real_frame_num = frame_indices[i]
current_score = scores_list[i]
# Priority:
# 1. Real Video Frame (from Loader)
# 2. Auto-Increment Counter (Start + i)
if real_frame_num > 0:
number_part = real_frame_num
else:
@@ -182,7 +217,8 @@ class FastAbsoluteSaver:
futures.append(executor.submit(
self.save_single_image,
img_tensor, full_path, current_score, metadata_key,
save_format, webp_lossless, webp_quality, webp_method
save_format, webp_lossless, webp_quality, webp_method,
save_workflow_metadata, prompt, extra_pnginfo
))
concurrent.futures.wait(futures)