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

74
fast_saver.py
View File

@@ -1,12 +1,11 @@
import os
import torch
import numpy as np
from PIL import Image, ExifTags
from PIL import Image
from PIL.PngImagePlugin import PngInfo
import concurrent.futures
import re
import time
import json
class FastAbsoluteSaver:
@classmethod
@@ -20,11 +19,14 @@ class FastAbsoluteSaver:
# --- FORMAT SWITCH ---
"save_format": (["png", "webp"], ),
# --- PERFORMANCE ---
"max_threads": ("INT", {"default": 0, "min": 0, "max": 128, "step": 1, "label": "Max Threads (0=Auto)"}),
# --- 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 SPECIFIC ---
"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)"}),
@@ -60,56 +62,18 @@ class FastAbsoluteSaver:
return frames[:batch_size], scores[:batch_size]
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))
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,
@@ -120,8 +84,8 @@ class FastAbsoluteSaver:
print(f"xx- Error saving {full_path}: {e}")
return False
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):
def save_images_fast(self, images, output_path, filename_prefix, save_format, max_threads,
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):
@@ -130,12 +94,20 @@ class FastAbsoluteSaver:
except OSError:
raise ValueError(f"Could not create directory: {output_path}")
# --- AUTO-SCALING LOGIC ---
if max_threads == 0:
# os.cpu_count() returns None on some rare systems, so we default to 4 just in case
cpu_cores = os.cpu_count() or 4
# For WebP (CPU intensive), stick to core count.
# For PNG (Disk intensive), we could technically go higher, but core count is safe.
max_threads = cpu_cores
print(f"xx- FastSaver: Using {max_threads} Threads for saving.")
batch_size = len(images)
frame_indices, scores_list = self.parse_info(scores_info, batch_size)
print(f"xx- FastSaver: Saving {batch_size} images ({save_format}) to {output_path}...")
with concurrent.futures.ThreadPoolExecutor(max_workers=16) as executor:
with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor:
futures = []
for i, img_tensor in enumerate(images):
@@ -151,22 +123,20 @@ class FastAbsoluteSaver:
if real_frame_num == 0 and scores_info is None:
base_name = f"{filename_prefix}_{int(time.time())}_{i:03d}"
# Append correct extension
ext = ".webp" if save_format == "webp" else ".png"
fname = f"{base_name}{ext}"
full_path = os.path.join(output_path, fname)
# 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
save_format,
webp_lossless,
webp_quality,
webp_method
))
concurrent.futures.wait(futures)