EthanfelandClaude Opus 4.8
dbba6cd5f0
The nvenc_* formats require an ffmpeg built with NVENC. Static builds (e.g. johnvansickle, which _get_ffmpeg may resolve) lack it, producing a cryptic "Unknown encoder 'av1_nvenc'" crash. - _get_ffmpeg(required_encoder): prefer the first existing ffmpeg that actually provides the encoder, so a NVENC-capable system ffmpeg wins over a static build. - save_video: if the hardware encoder is unavailable anywhere, fall back to the CPU codec for the same container (av1_nvenc-webm -> VP9 webm, etc.) with a loud warning instead of losing the run's output. Co-Authored-By: Claude Opus 4.8 <noreply@anthropic.com>
756 lines
32 KiB
Python
756 lines
32 KiB
Python
import os
|
|
import platform
|
|
import torch
|
|
import numpy as np
|
|
from PIL import Image
|
|
from PIL.PngImagePlugin import PngInfo
|
|
import concurrent.futures
|
|
import re
|
|
import time
|
|
import glob
|
|
import json
|
|
import subprocess
|
|
import shutil
|
|
import stat
|
|
import tempfile
|
|
import urllib.request
|
|
import zipfile
|
|
import tarfile
|
|
import wave
|
|
|
|
_NODE_DIR = os.path.dirname(os.path.abspath(__file__))
|
|
_FFMPEG_DIR = os.path.join(_NODE_DIR, "ffmpeg_bin")
|
|
|
|
|
|
_ENCODER_CACHE = {}
|
|
|
|
|
|
def _ffmpeg_has_encoder(ffmpeg_path, encoder):
|
|
"""Return True if this ffmpeg binary provides the named encoder (cached)."""
|
|
key = (ffmpeg_path, encoder)
|
|
if key in _ENCODER_CACHE:
|
|
return _ENCODER_CACHE[key]
|
|
found = False
|
|
try:
|
|
out = subprocess.run([ffmpeg_path, "-hide_banner", "-encoders"],
|
|
capture_output=True, text=True, timeout=15)
|
|
found = encoder in out.stdout.split()
|
|
except Exception:
|
|
found = False
|
|
_ENCODER_CACHE[key] = found
|
|
return found
|
|
|
|
|
|
def _existing_ffmpeg_paths():
|
|
"""Existing ffmpeg binaries in priority order (bundled, imageio, PATH). No download."""
|
|
system = platform.system()
|
|
exe_name = "ffmpeg.exe" if system == "Windows" else "ffmpeg"
|
|
local_bin = os.path.join(_FFMPEG_DIR, exe_name)
|
|
|
|
paths = []
|
|
def _add(p):
|
|
if p and os.path.isfile(p) and p not in paths:
|
|
paths.append(p)
|
|
|
|
_add(local_bin)
|
|
try:
|
|
import imageio_ffmpeg
|
|
_add(imageio_ffmpeg.get_ffmpeg_exe())
|
|
except Exception:
|
|
pass
|
|
_add(shutil.which("ffmpeg"))
|
|
return paths
|
|
|
|
|
|
def _get_ffmpeg(required_encoder=None):
|
|
"""Find or download a ffmpeg binary. Search order:
|
|
1. Bundled binary in this node's ffmpeg_bin/ folder
|
|
2. imageio_ffmpeg (shipped by VideoHelperSuite)
|
|
3. System PATH
|
|
4. Auto-download a static build into ffmpeg_bin/
|
|
|
|
If required_encoder is given (e.g. 'av1_nvenc'), prefer the first existing
|
|
binary that actually provides it; otherwise return the highest-priority
|
|
existing binary, downloading a static build only if none exist.
|
|
"""
|
|
system = platform.system()
|
|
exe_name = "ffmpeg.exe" if system == "Windows" else "ffmpeg"
|
|
local_bin = os.path.join(_FFMPEG_DIR, exe_name)
|
|
|
|
candidates = _existing_ffmpeg_paths()
|
|
|
|
if required_encoder:
|
|
for path in candidates:
|
|
if _ffmpeg_has_encoder(path, required_encoder):
|
|
return path
|
|
# None of the existing binaries has it; return the default below and let
|
|
# the caller decide how to fall back.
|
|
|
|
if candidates:
|
|
return candidates[0]
|
|
|
|
# 4. Auto-download static build
|
|
print("xx- FastSaver: ffmpeg not found. Downloading static build...")
|
|
os.makedirs(_FFMPEG_DIR, exist_ok=True)
|
|
|
|
urls = {
|
|
("Linux", "x86_64"): "https://johnvansickle.com/ffmpeg/releases/ffmpeg-release-amd64-static.tar.xz",
|
|
("Linux", "aarch64"): "https://johnvansickle.com/ffmpeg/releases/ffmpeg-release-arm64-static.tar.xz",
|
|
("Windows", "AMD64"): "https://www.gyan.dev/ffmpeg/builds/ffmpeg-release-essentials.zip",
|
|
}
|
|
|
|
machine = platform.machine()
|
|
key = (system, machine)
|
|
url = urls.get(key)
|
|
if not url:
|
|
raise RuntimeError(
|
|
f"No automatic ffmpeg download available for {system}/{machine}. "
|
|
f"Please install ffmpeg manually or place a binary in {_FFMPEG_DIR}"
|
|
)
|
|
|
|
archive_path = os.path.join(_FFMPEG_DIR, "ffmpeg_archive")
|
|
try:
|
|
urllib.request.urlretrieve(url, archive_path)
|
|
|
|
if url.endswith(".tar.xz"):
|
|
with tarfile.open(archive_path, "r:xz") as tar:
|
|
for member in tar.getmembers():
|
|
if member.name.endswith("/ffmpeg") and member.isfile():
|
|
member.name = exe_name
|
|
tar.extract(member, _FFMPEG_DIR)
|
|
break
|
|
elif url.endswith(".zip"):
|
|
with zipfile.ZipFile(archive_path, "r") as zf:
|
|
for name in zf.namelist():
|
|
if name.endswith("bin/ffmpeg.exe"):
|
|
data = zf.read(name)
|
|
with open(local_bin, "wb") as f:
|
|
f.write(data)
|
|
break
|
|
finally:
|
|
if os.path.exists(archive_path):
|
|
os.remove(archive_path)
|
|
|
|
if not os.path.isfile(local_bin):
|
|
raise RuntimeError(f"Failed to extract ffmpeg to {local_bin}")
|
|
|
|
# Make executable on Linux/Mac
|
|
if system != "Windows":
|
|
os.chmod(local_bin, os.stat(local_bin).st_mode | stat.S_IEXEC)
|
|
|
|
print(f"xx- FastSaver: ffmpeg downloaded to {local_bin}")
|
|
return local_bin
|
|
|
|
_COLOR_MGMT = [
|
|
"-vf", "scale=out_color_matrix=bt709",
|
|
"-color_range", "tv", "-colorspace", "bt709",
|
|
"-color_primaries", "bt709", "-color_trc", "bt709",
|
|
]
|
|
|
|
VIDEO_FORMATS = {
|
|
"mp4": {"ext": ".mp4", "codec": ["-c:v", "libx264"],
|
|
"quality": "crf", "color_mgmt": True, "acodec": "aac",
|
|
"extra": ["-movflags", "+faststart"]},
|
|
"h265-mp4": {"ext": ".mp4", "codec": ["-c:v", "libx265", "-vtag", "hvc1",
|
|
"-preset", "medium", "-x265-params", "log-level=quiet"],
|
|
"quality": "crf", "color_mgmt": True, "acodec": "aac",
|
|
"extra": ["-movflags", "+faststart"]},
|
|
"av1-mp4": {"ext": ".mp4", "codec": ["-c:v", "libsvtav1"],
|
|
"quality": "crf", "color_mgmt": True, "acodec": "aac",
|
|
"extra": ["-movflags", "+faststart"]},
|
|
"webm": {"ext": ".webm", "codec": ["-c:v", "libvpx-vp9"],
|
|
"quality": "crf", "zero_bitrate": True, "color_mgmt": True,
|
|
"acodec": "libopus"},
|
|
"gif": {"ext": ".gif", "special": "gif"},
|
|
"ffv1-mkv": {"ext": ".mkv", "codec": ["-c:v", "ffv1", "-level", "3",
|
|
"-coder", "1", "-context", "1", "-g", "1",
|
|
"-slices", "16", "-slicecrc", "1"],
|
|
"quality": "lossless", "acodec": "flac"},
|
|
"prores-mov": {"ext": ".mov", "codec": ["-c:v", "prores_ks"],
|
|
"quality": "profile", "color_mgmt": True, "acodec": "pcm_s16le"},
|
|
"nvenc_h264-mp4":{"ext": ".mp4", "codec": ["-c:v", "h264_nvenc"],
|
|
"quality": "bitrate", "color_mgmt": True, "acodec": "aac",
|
|
"extra": ["-movflags", "+faststart"]},
|
|
"nvenc_hevc-mp4":{"ext": ".mp4", "codec": ["-c:v", "hevc_nvenc", "-vtag", "hvc1"],
|
|
"quality": "bitrate", "color_mgmt": True, "acodec": "aac",
|
|
"extra": ["-movflags", "+faststart"]},
|
|
"nvenc_av1-mp4": {"ext": ".mp4", "codec": ["-c:v", "av1_nvenc"],
|
|
"quality": "bitrate", "color_mgmt": True, "acodec": "aac",
|
|
"extra": ["-movflags", "+faststart"]},
|
|
# GPU webm: NVENC has no VP9 encoder, so use AV1 (a valid WebM codec) on the GPU.
|
|
# WebM containers only allow Opus/Vorbis audio, not AAC.
|
|
"nvenc_av1-webm":{"ext": ".webm", "codec": ["-c:v", "av1_nvenc"],
|
|
"quality": "bitrate", "color_mgmt": True, "acodec": "libopus"},
|
|
}
|
|
|
|
# If a hardware (NVENC) format is picked but the encoder is missing from the
|
|
# resolved ffmpeg, fall back to the CPU codec that produces the SAME container.
|
|
_HW_CPU_FALLBACK = {
|
|
"nvenc_h264-mp4": "mp4", # h264_nvenc -> libx264
|
|
"nvenc_hevc-mp4": "h265-mp4", # hevc_nvenc -> libx265
|
|
"nvenc_av1-mp4": "av1-mp4", # av1_nvenc -> libsvtav1
|
|
"nvenc_av1-webm": "webm", # av1_nvenc -> libvpx-vp9
|
|
}
|
|
|
|
|
|
def _latent_sidecar_path(media_path):
|
|
return os.path.splitext(media_path)[0] + ".latent"
|
|
|
|
|
|
def _load_latent_file(latent_path):
|
|
try:
|
|
return torch.load(latent_path, map_location="cpu", weights_only=False)
|
|
except TypeError:
|
|
return torch.load(latent_path, map_location="cpu")
|
|
|
|
|
|
class FastAbsoluteSaver:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"images": ("IMAGE", ),
|
|
"output_path": ("STRING", {"default": "D:\\Datasets\\Sharp_Output"}),
|
|
"filename_prefix": ("STRING", {"default": "frame"}),
|
|
|
|
# --- FORMAT SWITCH ---
|
|
"save_format": (["png", "webp", "mp4", "webm", "h265-mp4", "av1-mp4", "gif",
|
|
"ffv1-mkv", "prores-mov", "nvenc_h264-mp4", "nvenc_hevc-mp4", "nvenc_av1-mp4",
|
|
"nvenc_av1-webm"], ),
|
|
|
|
# --- NAMING CONTROL ---
|
|
"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"}),
|
|
|
|
# --- METADATA & WORKFLOW ---
|
|
"metadata_key": ("STRING", {"default": "sharpness_score"}),
|
|
"save_workflow_metadata": ("BOOLEAN", {"default": False, "label": "Save ComfyUI Workflow (Graph)"}),
|
|
"save_metadata_png": ("BOOLEAN", {"default": False, "label": "Embed Workflow in PNG (sidecar or first file)"}),
|
|
"save_latent": ("BOOLEAN", {"default": True, "label": "Save Latent Sidecar"}),
|
|
|
|
# --- PERFORMANCE ---
|
|
"max_threads": ("INT", {"default": 0, "min": 0, "max": 128, "step": 1, "label": "Max Threads (0=Auto)"}),
|
|
|
|
# --- WEBP SPECIFIC ---
|
|
"webp_lossless": ("BOOLEAN", {"default": True, "label": "WebP Lossless"}),
|
|
"webp_quality": ("INT", {"default": 100, "min": 0, "max": 100, "step": 1}),
|
|
"webp_method": ("INT", {"default": 4, "min": 0, "max": 6, "step": 1}),
|
|
|
|
# --- VIDEO SPECIFIC ---
|
|
"video_fps": ("INT", {"default": 24, "min": 1, "max": 120, "step": 1, "label": "Video FPS"}),
|
|
"video_crf": ("INT", {"default": 18, "min": 0, "max": 51, "step": 1, "label": "Video CRF (0=Lossless, 51=Worst)"}),
|
|
"video_pixel_format": (["yuv420p", "yuv422p", "yuv444p", "yuv420p10le", "rgb24", "bgra"], {"label": "Pixel Format"}),
|
|
"video_bitrate": ("INT", {"default": 10, "min": 1, "max": 999, "step": 1, "label": "Video Bitrate (Mbps, NVENC)"}),
|
|
"prores_profile": (["lt", "standard", "hq", "4444", "4444xq"], {"label": "ProRes Profile"}),
|
|
"gif_dither": (["sierra2_4a", "floyd_steinberg", "bayer", "sierra2", "sierra3", "burkes", "atkinson", "heckbert", "none"], {"label": "GIF Dither Algorithm"}),
|
|
},
|
|
"optional": {
|
|
"scores_info": ("STRING", {"forceInput": True}),
|
|
"audio": ("AUDIO", ),
|
|
"latent": ("LATENT", ),
|
|
},
|
|
# Hidden inputs used to capture the workflow graph
|
|
"hidden": {"prompt": "PROMPT", "extra_pnginfo": "EXTRA_PNGINFO"},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("latent",)
|
|
FUNCTION = "save_images_fast"
|
|
OUTPUT_NODE = True
|
|
CATEGORY = "JSON Dynamic/io"
|
|
|
|
def parse_info(self, info_str, batch_size):
|
|
if not info_str:
|
|
# No scores connected - return None for scores to indicate "not provided"
|
|
return ([0]*batch_size, None)
|
|
matches = re.findall(r"F:(\d+).*?Score:\s*(\d+(\.\d+)?)", info_str)
|
|
frames = []
|
|
scores = []
|
|
for m in matches:
|
|
try:
|
|
frames.append(int(m[0]))
|
|
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.
|
|
print(f"xx- FastSaver: Scanning folder for existing '{prefix}' files...")
|
|
files = glob.glob(os.path.join(output_path, f"{prefix}*.*"))
|
|
|
|
max_idx = 0
|
|
# 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)
|
|
match = pattern.match(fname)
|
|
if match:
|
|
try:
|
|
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_sidecar_png(self, img, png_path, score, key_name, prompt_data, extra_data):
|
|
"""Save a metadata-carrying PNG sidecar next to a non-PNG export."""
|
|
meta = PngInfo()
|
|
if score is not None:
|
|
meta.add_text(key_name, str(score))
|
|
meta.add_text("software", "ComfyUI_FastAbsoluteSaver")
|
|
workflow_json = json.dumps(extra_data.get("workflow", {})) if extra_data else "{}"
|
|
prompt_json = json.dumps(prompt_data) if prompt_data else "{}"
|
|
meta.add_text("prompt", prompt_json)
|
|
meta.add_text("workflow", workflow_json)
|
|
img.save(png_path, format="PNG", pnginfo=meta, compress_level=1)
|
|
|
|
def _save_latent_sidecar(self, latent, media_path):
|
|
latent_path = _latent_sidecar_path(media_path)
|
|
torch.save(latent, latent_path)
|
|
print(f"xx- FastSaver: Latent sidecar saved to {latent_path}")
|
|
return latent_path
|
|
|
|
def save_single_image(self, img_array, full_path, score, key_name, fmt, lossless, quality, method,
|
|
save_workflow, prompt_data, extra_data, force_png_metadata=False):
|
|
try:
|
|
img = Image.fromarray(img_array)
|
|
|
|
# --- METADATA PREPARATION ---
|
|
meta_png = PngInfo()
|
|
|
|
if fmt == "png":
|
|
if score is not None:
|
|
meta_png.add_text(key_name, str(score))
|
|
meta_png.add_text("software", "ComfyUI_Parallel_Node")
|
|
|
|
# ComfyUI Workflow Metadata (if requested, or forced via sidecar toggle for first PNG)
|
|
if (save_workflow or force_png_metadata) and fmt == "png":
|
|
workflow_json = json.dumps(extra_data.get("workflow", {})) if extra_data else "{}"
|
|
prompt_json = json.dumps(prompt_data) if prompt_data else "{}"
|
|
meta_png.add_text("prompt", prompt_json)
|
|
meta_png.add_text("workflow", workflow_json)
|
|
|
|
# --- 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 _write_temp_wav(self, audio):
|
|
"""Write a ComfyUI AUDIO dict ({waveform, sample_rate}) to a temp WAV file.
|
|
waveform is a tensor of shape [batch, channels, samples] in float [-1, 1].
|
|
Returns the file path, or None if no usable audio is present."""
|
|
try:
|
|
waveform = audio.get("waveform")
|
|
sample_rate = int(audio.get("sample_rate", 0))
|
|
if waveform is None or sample_rate <= 0:
|
|
return None
|
|
|
|
# Take the first item in the batch -> [channels, samples]
|
|
wf = waveform[0]
|
|
if hasattr(wf, "cpu"):
|
|
wf = wf.cpu().numpy()
|
|
wf = np.asarray(wf, dtype=np.float32)
|
|
if wf.ndim == 1:
|
|
wf = wf[np.newaxis, :]
|
|
channels = wf.shape[0]
|
|
if channels == 0 or wf.shape[1] == 0:
|
|
return None
|
|
|
|
# Interleave channels: [channels, samples] -> [samples, channels]
|
|
interleaved = np.clip(wf.T, -1.0, 1.0)
|
|
pcm = (interleaved * 32767.0).astype(np.int16)
|
|
|
|
tmp = tempfile.NamedTemporaryFile(suffix=".wav", delete=False)
|
|
tmp.close()
|
|
with wave.open(tmp.name, "wb") as w:
|
|
w.setnchannels(channels)
|
|
w.setsampwidth(2)
|
|
w.setframerate(sample_rate)
|
|
w.writeframes(pcm.tobytes())
|
|
return tmp.name
|
|
except Exception as e:
|
|
print(f"xx- FastSaver: Could not prepare audio, skipping: {e}")
|
|
return None
|
|
|
|
def save_video(self, frames_np, output_path, filename_prefix, use_timestamp, fps, crf, pixel_format, video_format,
|
|
auto_increment=False, counter_digits=4,
|
|
scores_list=None, metadata_key="sharpness_score", save_workflow=False, prompt_data=None, extra_data=None,
|
|
bitrate=10, prores_profile="hq", gif_dither="sierra2_4a", audio=None):
|
|
"""Save image batch as a video file using ffmpeg. frames_np is a list/array of uint8 numpy arrays."""
|
|
fmt = VIDEO_FORMATS[video_format]
|
|
|
|
# Resolve ffmpeg, preferring one that provides this format's hardware encoder.
|
|
codec_list = fmt.get("codec", [])
|
|
required_encoder = None
|
|
if "-c:v" in codec_list:
|
|
enc = codec_list[codec_list.index("-c:v") + 1]
|
|
if any(tag in enc for tag in ("nvenc", "qsv", "vaapi", "amf")):
|
|
required_encoder = enc
|
|
|
|
ffmpeg_path = _get_ffmpeg(required_encoder)
|
|
|
|
# If the hardware encoder still isn't available anywhere, fall back to the
|
|
# CPU codec for the same container so the run's output isn't lost.
|
|
if required_encoder and not _ffmpeg_has_encoder(ffmpeg_path, required_encoder):
|
|
fallback = _HW_CPU_FALLBACK.get(video_format)
|
|
if fallback:
|
|
print(f"xx- FastSaver: WARNING: '{required_encoder}' is not available in "
|
|
f"ffmpeg ({ffmpeg_path}); falling back to CPU format '{fallback}' "
|
|
f"(slower). Install an ffmpeg built with NVENC for GPU encoding.")
|
|
video_format = fallback
|
|
fmt = VIDEO_FORMATS[video_format]
|
|
|
|
ext = fmt["ext"]
|
|
if use_timestamp:
|
|
ts_str = f"_{int(time.time())}"
|
|
out_file = os.path.join(output_path, f"{filename_prefix}{ts_str}{ext}")
|
|
elif auto_increment:
|
|
start_idx = self.get_start_index(output_path, filename_prefix)
|
|
fmt_str = f"{{:0{counter_digits}d}}"
|
|
out_file = os.path.join(output_path, f"{filename_prefix}_{fmt_str.format(start_idx)}{ext}")
|
|
else:
|
|
out_file = os.path.join(output_path, f"{filename_prefix}{ext}")
|
|
|
|
batch_size = len(frames_np)
|
|
h, w = frames_np[0].shape[0], frames_np[0].shape[1]
|
|
|
|
# --- GIF SPECIAL CASE ---
|
|
if fmt.get("special") == "gif":
|
|
if audio is not None:
|
|
print("xx- FastSaver: GIF format cannot carry audio, ignoring audio input.")
|
|
filter_str = (
|
|
"[0:v] split [a][b]; [a] palettegen=reserve_transparent=on"
|
|
":transparency_color=ffffff [p]; [b][p] paletteuse=dither=" + gif_dither
|
|
)
|
|
cmd = [ffmpeg_path, "-y", "-f", "rawvideo", "-pix_fmt", "rgb24",
|
|
"-s", f"{w}x{h}", "-r", str(fps), "-i", "-",
|
|
"-filter_complex", filter_str, out_file]
|
|
|
|
print(f"xx- FastSaver: Encoding {batch_size} frames to {out_file} (gif, {fps}fps)...")
|
|
|
|
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
try:
|
|
for frame in frames_np:
|
|
proc.stdin.write(frame.tobytes())
|
|
proc.stdin.close()
|
|
except BrokenPipeError:
|
|
pass
|
|
stderr = proc.stderr.read()
|
|
proc.wait()
|
|
|
|
if proc.returncode != 0:
|
|
raise RuntimeError(f"ffmpeg failed: {stderr.decode()}")
|
|
|
|
print(f"xx- FastSaver: Video saved to {out_file}")
|
|
return out_file
|
|
|
|
# --- BUILD METADATA FILE (avoids arg-too-long for large workflows) ---
|
|
meta_lines = [";FFMETADATA1"]
|
|
meta_lines.append("software=ComfyUI_FastAbsoluteSaver")
|
|
|
|
if scores_list:
|
|
avg_score = sum(scores_list) / len(scores_list)
|
|
meta_lines.append(f"{metadata_key}_avg={avg_score:.2f}")
|
|
meta_lines.append(f"{metadata_key}_all={','.join(f'{s:.2f}' for s in scores_list)}")
|
|
|
|
if save_workflow:
|
|
if prompt_data:
|
|
# Escape ffmetadata special chars: =, ;, #, \ and newlines
|
|
prompt_str = json.dumps(prompt_data).replace("\\", "\\\\").replace("=", "\\=").replace(";", "\\;").replace("#", "\\#").replace("\n", "\\\n")
|
|
meta_lines.append(f"prompt={prompt_str}")
|
|
if extra_data:
|
|
workflow = extra_data.get("workflow", {})
|
|
if workflow:
|
|
workflow_str = json.dumps(workflow).replace("\\", "\\\\").replace("=", "\\=").replace(";", "\\;").replace("#", "\\#").replace("\n", "\\\n")
|
|
meta_lines.append(f"workflow={workflow_str}")
|
|
|
|
self._meta_tmpfile = tempfile.NamedTemporaryFile(mode="w", suffix=".txt", delete=False, encoding="utf-8")
|
|
self._meta_tmpfile.write("\n".join(meta_lines))
|
|
self._meta_tmpfile.close()
|
|
meta_file = self._meta_tmpfile.name
|
|
|
|
# --- AUDIO (optional) ---
|
|
audio_file = self._write_temp_wav(audio) if audio is not None else None
|
|
|
|
# --- BUILD FFMPEG COMMAND ---
|
|
# Input order: 0=rawvideo (stdin), 1=metadata, 2=audio (if present)
|
|
cmd = [ffmpeg_path, "-y",
|
|
"-f", "rawvideo", "-pix_fmt", "rgb24",
|
|
"-s", f"{w}x{h}", "-r", str(fps), "-i", "-",
|
|
"-i", meta_file]
|
|
if audio_file:
|
|
cmd.extend(["-i", audio_file])
|
|
cmd.extend(["-map_metadata", "1"])
|
|
|
|
# Codec args
|
|
cmd.extend(fmt["codec"])
|
|
|
|
# Quality control
|
|
quality_mode = fmt.get("quality", "crf")
|
|
if quality_mode == "crf":
|
|
cmd.extend(["-crf", str(crf)])
|
|
if fmt.get("zero_bitrate"):
|
|
cmd.extend(["-b:v", "0"])
|
|
elif quality_mode == "bitrate":
|
|
cmd.extend(["-b:v", f"{bitrate}M"])
|
|
elif quality_mode == "profile":
|
|
# ffmpeg's prores_ks accepts "xq" not "4444xq"
|
|
ffmpeg_profile = "xq" if prores_profile == "4444xq" else prores_profile
|
|
cmd.extend(["-profile:v", ffmpeg_profile])
|
|
|
|
# Pixel format
|
|
if "pix_fmt" in fmt:
|
|
cmd.extend(["-pix_fmt", fmt["pix_fmt"]])
|
|
elif quality_mode == "profile":
|
|
# ProRes auto pixel format based on profile
|
|
if prores_profile in ("4444", "4444xq"):
|
|
cmd.extend(["-pix_fmt", "yuv444p10le"])
|
|
else:
|
|
cmd.extend(["-pix_fmt", "yuv422p10le"])
|
|
else:
|
|
# NVENC doesn't support yuv420p10le — remap to p010le
|
|
pf = pixel_format
|
|
if "nvenc" in video_format and pf == "yuv420p10le":
|
|
pf = "p010le"
|
|
cmd.extend(["-pix_fmt", pf])
|
|
|
|
# Color management
|
|
if fmt.get("color_mgmt"):
|
|
cmd.extend(_COLOR_MGMT)
|
|
|
|
# Extra args (e.g. -movflags +faststart)
|
|
if "extra" in fmt:
|
|
cmd.extend(fmt["extra"])
|
|
|
|
# Audio: map the video stream explicitly and mux in the audio track
|
|
if audio_file:
|
|
acodec = fmt.get("acodec", "aac")
|
|
cmd.extend(["-map", "0:v:0", "-map", "2:a:0", "-c:a", acodec, "-shortest"])
|
|
|
|
cmd.append(out_file)
|
|
|
|
codec_label = fmt["codec"][fmt["codec"].index("-c:v") + 1] if "-c:v" in fmt["codec"] else video_format
|
|
print(f"xx- FastSaver: Encoding {batch_size} frames to {out_file} ({codec_label}, {fps}fps)...")
|
|
|
|
proc = subprocess.Popen(cmd, stdin=subprocess.PIPE, stderr=subprocess.PIPE)
|
|
try:
|
|
for frame in frames_np:
|
|
proc.stdin.write(frame.tobytes())
|
|
proc.stdin.close()
|
|
except BrokenPipeError:
|
|
pass
|
|
stderr = proc.stderr.read()
|
|
proc.wait()
|
|
|
|
# Clean up temp files
|
|
try:
|
|
os.remove(meta_file)
|
|
except OSError:
|
|
pass
|
|
if audio_file:
|
|
try:
|
|
os.remove(audio_file)
|
|
except OSError:
|
|
pass
|
|
|
|
if proc.returncode != 0:
|
|
raise RuntimeError(f"ffmpeg failed: {stderr.decode()}")
|
|
|
|
print(f"xx- FastSaver: Video saved to {out_file}")
|
|
return out_file
|
|
|
|
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, save_workflow_metadata, save_metadata_png, save_latent,
|
|
webp_lossless, webp_quality, webp_method,
|
|
video_fps, video_crf, video_pixel_format,
|
|
video_bitrate, prores_profile, gif_dither,
|
|
scores_info=None, audio=None, latent=None, prompt=None, extra_pnginfo=None):
|
|
|
|
output_path = output_path.strip('"')
|
|
if not os.path.isabs(output_path):
|
|
import folder_paths
|
|
output_path = os.path.join(folder_paths.base_path, output_path)
|
|
if not os.path.exists(output_path):
|
|
try:
|
|
os.makedirs(output_path, exist_ok=True)
|
|
except OSError:
|
|
raise ValueError(f"Could not create directory: {output_path}")
|
|
|
|
if images is None or len(images) == 0:
|
|
raise ValueError("No images provided to FastAbsoluteSaver.")
|
|
|
|
# Batch GPU->CPU transfer once (avoids N individual sync transfers)
|
|
images_np = (255.0 * images.cpu().numpy()).clip(0, 255).astype(np.uint8)
|
|
batch_size = len(images_np)
|
|
|
|
# --- VIDEO PATH (check early, before image-specific logic) ---
|
|
if save_format in VIDEO_FORMATS:
|
|
_, scores_list = self.parse_info(scores_info, batch_size)
|
|
out_file = self.save_video(images_np, output_path, filename_prefix, use_timestamp,
|
|
video_fps, video_crf, video_pixel_format, save_format,
|
|
auto_increment=auto_increment, counter_digits=counter_digits,
|
|
scores_list=scores_list, metadata_key=metadata_key,
|
|
save_workflow=save_workflow_metadata, prompt_data=prompt,
|
|
extra_data=extra_pnginfo,
|
|
bitrate=video_bitrate, prores_profile=prores_profile,
|
|
gif_dither=gif_dither, audio=audio)
|
|
if latent is not None and save_latent:
|
|
self._save_latent_sidecar(latent, out_file)
|
|
# Save metadata sidecar PNG next to the video file
|
|
if save_metadata_png:
|
|
png_path = os.path.splitext(out_file)[0] + ".png"
|
|
img = Image.fromarray(images_np[0])
|
|
avg_score = (sum(scores_list) / len(scores_list)) if scores_list else None
|
|
self._save_sidecar_png(img, png_path, avg_score, metadata_key, prompt, extra_pnginfo)
|
|
print(f"xx- FastSaver: Metadata PNG sidecar saved to {png_path}")
|
|
return {"ui": {"images": []}, "result": (latent,)}
|
|
|
|
if max_threads == 0:
|
|
max_threads = os.cpu_count() or 4
|
|
|
|
frame_indices, scores_list = self.parse_info(scores_info, batch_size)
|
|
|
|
# --- INDEX LOGIC ---
|
|
start_counter = 0
|
|
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}...")
|
|
|
|
first_image_path = None
|
|
saved_image_paths = []
|
|
|
|
with concurrent.futures.ThreadPoolExecutor(max_workers=max_threads) as executor:
|
|
futures = []
|
|
|
|
for i, img_array in enumerate(images_np):
|
|
real_frame_num = frame_indices[i]
|
|
current_score = scores_list[i] if scores_list else None
|
|
|
|
if real_frame_num > 0:
|
|
number_part = real_frame_num
|
|
else:
|
|
number_part = start_counter + i
|
|
|
|
fmt_str = f"{{:0{counter_digits}d}}"
|
|
number_str = fmt_str.format(number_part)
|
|
|
|
base_name = f"{filename_prefix}{ts_str}_{number_str}"
|
|
|
|
if filename_with_score and current_score is not None:
|
|
base_name += f"_{int(current_score)}"
|
|
|
|
ext = ".webp" if save_format == "webp" else ".png"
|
|
full_path = os.path.join(output_path, f"{base_name}{ext}")
|
|
|
|
if i == 0:
|
|
first_image_path = full_path
|
|
|
|
# For PNG sequences: embed workflow metadata in the first file only
|
|
force_meta = (save_metadata_png and save_format == "png" and i == 0)
|
|
|
|
future = executor.submit(
|
|
self.save_single_image,
|
|
img_array, full_path, current_score, metadata_key,
|
|
save_format, webp_lossless, webp_quality, webp_method,
|
|
save_workflow_metadata, prompt, extra_pnginfo,
|
|
force_png_metadata=force_meta
|
|
)
|
|
futures.append((full_path, future))
|
|
|
|
concurrent.futures.wait([future for _, future in futures])
|
|
|
|
for full_path, future in futures:
|
|
if future.result():
|
|
saved_image_paths.append(full_path)
|
|
|
|
if latent is not None and save_latent:
|
|
for image_path in saved_image_paths:
|
|
self._save_latent_sidecar(latent, image_path)
|
|
|
|
# Save a single metadata sidecar PNG using the first image (skip for PNG sequences - handled above)
|
|
if save_metadata_png and save_format != "png" and first_image_path is not None:
|
|
png_path = os.path.splitext(first_image_path)[0] + ".png"
|
|
img = Image.fromarray(images_np[0])
|
|
first_score = scores_list[0] if scores_list else None
|
|
self._save_sidecar_png(img, png_path, first_score, metadata_key, prompt, extra_pnginfo)
|
|
print(f"xx- FastSaver: Metadata PNG sidecar saved to {png_path}")
|
|
|
|
return {"ui": {"images": []}, "result": (latent,)}
|
|
|
|
|
|
class JDL_LoadLatentAbsolute:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {
|
|
"required": {
|
|
"path": ("STRING", {"default": "", "multiline": False}),
|
|
},
|
|
}
|
|
|
|
RETURN_TYPES = ("LATENT",)
|
|
RETURN_NAMES = ("latent",)
|
|
FUNCTION = "load_latent"
|
|
CATEGORY = "JSON Dynamic/io"
|
|
|
|
def load_latent(self, path):
|
|
latent_path = path.strip().strip('"')
|
|
if not os.path.isabs(latent_path):
|
|
raise ValueError(f"Load Latent Absolute requires an absolute path: {path}")
|
|
if not os.path.isfile(latent_path):
|
|
raise FileNotFoundError(f"Latent file not found: {latent_path}")
|
|
|
|
latent = _load_latent_file(latent_path)
|
|
if not isinstance(latent, dict) or "samples" not in latent:
|
|
raise ValueError(f"Expected a latent dict with a 'samples' key: {latent_path}")
|
|
return (latent,)
|
|
|
|
@classmethod
|
|
def IS_CHANGED(s, path):
|
|
latent_path = path.strip().strip('"')
|
|
if not os.path.isfile(latent_path):
|
|
return f"missing:{latent_path}"
|
|
stat_result = os.stat(latent_path)
|
|
return f"{latent_path}:{stat_result.st_mtime_ns}:{stat_result.st_size}"
|
|
|
|
@classmethod
|
|
def VALIDATE_INPUTS(s, path):
|
|
latent_path = path.strip().strip('"')
|
|
if not os.path.isabs(latent_path):
|
|
return f"Path must be absolute: {path}"
|
|
if not os.path.isfile(latent_path):
|
|
return f"Latent file not found: {latent_path}"
|
|
return True
|
|
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"FastAbsoluteSaver": FastAbsoluteSaver,
|
|
"JDL_LoadLatentAbsolute": JDL_LoadLatentAbsolute,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"FastAbsoluteSaver": "Fast Absolute Saver (Metadata)",
|
|
"JDL_LoadLatentAbsolute": "Load Latent Absolute",
|
|
}
|