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-20 00:20:21 +01:00
parent 1bde14bd97
commit b54b4329ca
1 changed files with 82 additions and 21 deletions
Download Patch File Download Diff File Expand all files Collapse all files

101
fast_saver.py
View File

@@ -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()
metadata.add_text(key_name, str(score))
metadata.add_text("software", "ComfyUI_Parallel_Node")
if fmt == "png":
# PNG METADATA (Robust)
metadata = PngInfo()
metadata.add_text(key_name, str(score))
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)
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, 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)