266 lines
15 KiB
Python
266 lines
15 KiB
Python
"""Save a video as MP4 into ComfyUI's output folder, fast: H.264 on the GPU's encoder by default.
|
|
|
|
Core's SaveVideo encodes with libx264 through PyAV, and FFmpeg's library default is ONE thread -- on a
|
|
rental's CPU that is where a run's last half-minute went. Here the GPU's own encoder (NVENC) does it when
|
|
this PyAV build and the card have one. H.264 by default: it plays in every browser, Firefox included, and
|
|
NVENC makes it as fast as H.265. H.265 (smaller files, tagged hvc1) is an option for Chrome/Safari. With no
|
|
NVENC it is libx264 on every core: CPU H.265 is several times slower. Same names and metadata as SaveVideo. Takes the frames and sound directly: core's
|
|
CreateVideo step in between only bundled them.
|
|
"""
|
|
|
|
import json
|
|
import math
|
|
import os
|
|
import time
|
|
from fractions import Fraction
|
|
|
|
import av
|
|
import folder_paths
|
|
import numpy as np
|
|
import torch
|
|
import torch.nn.functional as F
|
|
from comfy.cli_args import args
|
|
from comfy_api.latest import io, ui
|
|
|
|
from ..._core import log, registry as registry_mod, schema as schema_mod, video_pipe
|
|
from . import watermark
|
|
|
|
ENCODERS = ("auto", "gpu", "cpu")
|
|
CODECS = ("h264", "h265")
|
|
# Constant-quality mode. The number is set by `quality` (QUALITY below). NVENC needs a lower number than
|
|
# libx264 for the same look, so its table is its own. 23 was the old fixed value: fine for a phone clip,
|
|
# but it smears fine texture and shows blocks and mosquito noise on a high-resolution diffusion picture.
|
|
QUALITIES = ("standard", "high", "maximum")
|
|
QUALITY = {"nvenc": {"standard": "23", "high": "16", "maximum": "10"},
|
|
"x264": {"standard": "23", "high": "16", "maximum": "9"}}
|
|
NVENC = {"h265": ("hevc_nvenc", {"preset": "p6", "rc": "vbr", "cq": "16", "b:v": "0"}),
|
|
"h264": ("h264_nvenc", {"preset": "p6", "rc": "vbr", "cq": "16", "b:v": "0"})}
|
|
X264 = ("libx264", {"crf": "16", "preset": "medium"})
|
|
# The finishing look, from Engine > Post (see __init__.SETTINGS), applied to the frames on their own device before encoding,
|
|
# so NVENC, its CPU fallback and every codec get the same picture. Order: contrast, brightness, saturation, edge smoothing.
|
|
# Contrast and brightness are measured ONCE on a sample of the whole clip and applied as one fixed curve: a per-frame
|
|
# auto-level would pump and flicker as the picture changes.
|
|
SMOOTHINGS = ("off", "light", "strong")
|
|
DEFAULTS = {"smoothing": "off", "auto_contrast": False, "auto_brightness": False, "saturation": 1.0,
|
|
"watermark": "off", "watermark_text": "", "watermark_colour": "#ffffff", "watermark_picture": "", "watermark_corner": "bottom-right",
|
|
"watermark_opacity": 0.6, "watermark_size": 18.0, "watermark_margin": 2.0,
|
|
"watermark_outline": 0.0, "watermark_outline_colour": "#000000", "watermark_outline_opacity": 1.0}
|
|
_AMOUNT = {"light": 0.5, "strong": 1.0}
|
|
_CHUNK = 8 # frames per step: the float working copy stays small on a long 4K clip
|
|
_SAMPLE_FRAMES, _SAMPLE_PIXELS = 16, 1_000_000
|
|
_TARGET_MEDIAN, _MAX_STRETCH = 0.45, 0.15 # brightness aims the median luma here; contrast moves each end by at most this
|
|
|
|
|
|
def _luma(f):
|
|
"""f: [T, 3, H, W] -> [T, 1, H, W]."""
|
|
return f[:, 0:1] * 0.299 + f[:, 1:2] * 0.587 + f[:, 2:3] * 0.114
|
|
|
|
|
|
def clip_look(images, values):
|
|
"""-> {"lo", "hi", "gamma"} as plain floats (no tensors, so nothing can sit on the wrong device) measured on a sample of the clip."""
|
|
look = {"lo": 0.0, "hi": 1.0, "gamma": 1.0}
|
|
if not (values.get("auto_contrast") or values.get("auto_brightness")) or len(images) == 0:
|
|
return look
|
|
n = len(images)
|
|
pick = sorted({round(i * (n - 1) / max(_SAMPLE_FRAMES - 1, 1)) for i in range(min(_SAMPLE_FRAMES, n))})
|
|
luma = _luma(images[pick].permute(0, 3, 1, 2).float()).flatten()
|
|
luma = luma[::max(1, luma.numel() // _SAMPLE_PIXELS)]
|
|
if values.get("auto_contrast"):
|
|
lo, hi = (float(v) for v in torch.quantile(luma, luma.new_tensor([0.01, 0.99])))
|
|
look["lo"], look["hi"] = min(lo, _MAX_STRETCH), max(hi, 1.0 - _MAX_STRETCH)
|
|
luma = ((luma - look["lo"]) / (look["hi"] - look["lo"])).clamp(0, 1)
|
|
if values.get("auto_brightness"):
|
|
mid = min(max(float(luma.median()), 0.02), 0.98)
|
|
look["gamma"] = min(max(math.log(_TARGET_MEDIAN) / math.log(mid), 0.75), 1.33)
|
|
return look
|
|
|
|
|
|
def smooth(x, level):
|
|
"""x: [T, H, W, 3] floats in 0..1 -> the same shape, dtype and device with edges anti-aliased (FXAA idea: only pixels on a
|
|
contrasty edge are blended with their 3x3 neighbours; `level` "off": x itself)."""
|
|
amount = _AMOUNT.get(level)
|
|
if amount is None:
|
|
return x
|
|
f = x.permute(0, 3, 1, 2).float()
|
|
p = F.pad(_luma(f), (1, 1, 1, 1), mode="replicate")
|
|
contrast = F.max_pool2d(p, 3, 1) + F.max_pool2d(-p, 3, 1) # local max - local min
|
|
edge = ((contrast - 0.04) / 0.16).clamp(0, 1) * amount
|
|
blur = F.avg_pool2d(F.pad(f, (1, 1, 1, 1), mode="replicate"), 3, 1)
|
|
return (f + (blur - f) * edge).permute(0, 2, 3, 1).to(x.dtype)
|
|
|
|
|
|
def finish(x, values, look):
|
|
"""One chunk x [T, H, W, 3] through the whole look; the same shape, dtype and device back."""
|
|
sat = float(values.get("saturation", 1.0))
|
|
if look["lo"] == 0.0 and look["hi"] == 1.0 and look["gamma"] == 1.0 and sat == 1.0:
|
|
return smooth(x, values.get("smoothing", "off"))
|
|
f = x.permute(0, 3, 1, 2).float()
|
|
if look["lo"] != 0.0 or look["hi"] != 1.0:
|
|
f = ((f - look["lo"]) / (look["hi"] - look["lo"])).clamp(0, 1)
|
|
if look["gamma"] != 1.0:
|
|
f = f.clamp(0, 1) ** look["gamma"]
|
|
if sat != 1.0:
|
|
grey = _luma(f)
|
|
f = (grey + (f - grey) * sat).clamp(0, 1)
|
|
return smooth(f.permute(0, 2, 3, 1).to(x.dtype), values.get("smoothing", "off"))
|
|
|
|
|
|
def to_frames(images, values=None):
|
|
"""-> uint8 [T, H, W, 3] on the CPU, `values` being the Engine > Post settings. Rounded: a bare cast cuts every value down, a steady half-level darkening."""
|
|
values = values or {}
|
|
look = clip_look(images, values)
|
|
mark = watermark.build(values, images.shape[2], images.shape[1]) # last: nothing before it can touch the mark
|
|
out = []
|
|
for x in images.split(_CHUNK):
|
|
f = finish(x, values, look)
|
|
y = (f * 255).round().clamp(0, 255).byte().cpu().numpy()
|
|
if mark is not None: # only the mark's corner is blended and pasted: no second copy of the chunk
|
|
part = watermark.corner(f, mark)
|
|
y[:, mark.y:mark.y + part.shape[1], mark.x:mark.x + part.shape[2], :] = (watermark.blend(part, mark) * 255).round().clamp(0, 255).byte().cpu().numpy()
|
|
out.append(y)
|
|
return np.concatenate(out)
|
|
|
|
|
|
def _with_quality(codec, quality):
|
|
"""-> the (name, options) pair with its constant-quality number set for `quality` (unknown names: high)."""
|
|
name, options = codec
|
|
table = QUALITY["nvenc" if "nvenc" in name else "x264"]
|
|
number = table.get(quality, table["high"])
|
|
options = dict(options)
|
|
for key in ("cq", "crf"):
|
|
if key in options:
|
|
options[key] = number
|
|
return name, options
|
|
# FFmpeg's "minimum required driver" line compares the driver library the PROCESS loaded, not what nvidia-smi reports
|
|
_HINT = (" (nvidia-smi can look fine while FFmpeg loads an older or missing libnvidia-encode.so.1; "
|
|
"in a container check it is there and that the NVIDIA 'video' capability is on)")
|
|
|
|
|
|
last_frame_name = video_pipe.last_frame_name
|
|
|
|
|
|
def keep_last_frame(video_file, frame):
|
|
"""Save the last decoded picture exactly as the sampler made it: the next shot can continue from this instead of from
|
|
the encoded video (which has already lost some of it). Never worth failing the save over."""
|
|
try:
|
|
from PIL import Image
|
|
folder = folder_paths.get_temp_directory()
|
|
os.makedirs(folder, exist_ok=True)
|
|
target = os.path.join(folder, last_frame_name(video_file))
|
|
if os.path.exists(target):
|
|
os.remove(target) # a failed save below must not leave an older video's picture under this name
|
|
Image.fromarray(frame).save(target, compress_level=1)
|
|
except Exception as exc: # noqa: BLE001 -- the video is saved; the next shot falls back to the video's own last picture
|
|
log.info("FunUI Save Video", f"could not keep the lossless last frame ({exc})")
|
|
|
|
|
|
def _write(path, frames, fps, audio, metadata, codec):
|
|
name, options = codec
|
|
# use_metadata_tags keeps the workflow in the file; faststart puts the index first, which the preview player needs to seek
|
|
with av.open(path, mode="w", format="mp4", options={"movflags": "use_metadata_tags+faststart"}) as out:
|
|
for key, value in (metadata or {}).items():
|
|
out.metadata[key] = json.dumps(value)
|
|
stream = out.add_stream(name, rate=Fraction(round(fps * 1000), 1000))
|
|
stream.width, stream.height, stream.pix_fmt = frames.shape[2], frames.shape[1], "yuv420p"
|
|
stream.options = options
|
|
stream.thread_type, stream.codec_context.thread_count = "AUTO", 0 # every core (FFmpeg's default is one)
|
|
if "hevc" in name or "265" in name:
|
|
stream.codec_context.codec_tag = "hvc1" # hev1 (the default) is refused by Safari and Chrome on a Mac
|
|
sound = None
|
|
if audio is not None: # both tracks before the first write: the header is fixed then
|
|
rate, wave = int(audio["sample_rate"]), audio["waveform"]
|
|
wave = wave[0, :, :int(rate * len(frames) / fps + 0.999)].float().cpu().contiguous().numpy()
|
|
layout = {1: "mono", 2: "stereo"}.get(wave.shape[0], "stereo")
|
|
sound = out.add_stream("aac", rate=rate, layout=layout)
|
|
for frame in frames:
|
|
out.mux(stream.encode(av.VideoFrame.from_ndarray(frame, format="rgb24").reformat(format="yuv420p")))
|
|
out.mux(stream.encode(None))
|
|
if sound is not None:
|
|
chunk = av.AudioFrame.from_ndarray(wave, format="fltp", layout=layout)
|
|
chunk.sample_rate, chunk.pts = rate, 0
|
|
out.mux(sound.encode(chunk))
|
|
out.mux(sound.encode(None))
|
|
|
|
|
|
def encode(path, frames, fps, audio=None, metadata=None, encoder="auto", codec="h264", quality="high"):
|
|
"""-> which encoder wrote the file. `frames`: uint8 [T, H, W, 3]. GPU first unless "cpu" is chosen;
|
|
a GPU attempt that fails leaves no half file and falls back to CPU H.264 with a Log alert saying why."""
|
|
if audio is not None and not bool(audio["waveform"].isfinite().all()): # NaN/Inf sound (typically the model loaded in fp16): AAC refuses it, so the video is kept without it
|
|
audio = None
|
|
log.alert("FunUI Save Video", "The audio contains NaN/Inf values, so you are most likely using the wrong precision for this model (fp16?). The audio has been removed; the video was saved without it.")
|
|
gpu = _with_quality(NVENC[codec], quality)
|
|
why = None
|
|
if encoder != "cpu":
|
|
if gpu[0] not in av.codecs_available:
|
|
why = f"this PyAV build has no NVENC encoder ({gpu[0]} missing)"
|
|
else:
|
|
try:
|
|
_write(path, frames, fps, audio, metadata, gpu)
|
|
return f"{codec.upper()} on NVENC (GPU)"
|
|
except Exception as exc: # noqa: BLE001 -- no encoder on this card, or a driver refusal
|
|
if os.path.exists(path):
|
|
os.remove(path)
|
|
why = f"the GPU encoder (NVENC) could not start: {exc}{_HINT}"
|
|
_write(path, frames, fps, audio, metadata, _with_quality(X264, quality))
|
|
if why: # the video is saved either way; never lose a generation to the encoder
|
|
log.once("save_nvenc", log.ALERT, "FunUI Save Video", f"{why}. Saved as H.264 on the CPU instead.")
|
|
return "H.264 on libx264 (CPU, all cores)"
|
|
|
|
|
|
def _values(settings):
|
|
"""The Engine > Post settings from the pipeline's settings object, checked against their declaration (a wrong one: default, said)."""
|
|
spec = registry_mod.current().specs.get("output_save")
|
|
if spec is None:
|
|
return dict(DEFAULTS)
|
|
pick = spec.settings.get("watermark_picture")
|
|
if pick is not None:
|
|
pick["options"] = watermark.picture_options() # read now: a picture added since Settings was opened still counts
|
|
clean, problems = schema_mod.check_values(spec, (settings or {}).get("output_save"))
|
|
for problem in problems or ():
|
|
log.once(f"save_video:{problem}", log.ALERT, "FunUI Save Video", f"setting ignored, default used: {problem}")
|
|
return clean
|
|
|
|
|
|
class FunUISaveVideo(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="FunUISaveVideo",
|
|
display_name="FunUI Save Video",
|
|
category="FunUI/Output",
|
|
description="Save as MP4: on the GPU's encoder (H.264 or H.265) when there is one, otherwise H.264 on every CPU core.",
|
|
inputs=[
|
|
io.Image.Input("images"),
|
|
io.Audio.Input("audio", optional=True),
|
|
io.Float.Input("fps", default=24.0, min=1.0, max=240.0, step=0.01),
|
|
io.String.Input("filename_prefix", default="FunUI"),
|
|
io.Combo.Input("encoder", options=list(ENCODERS), default="auto", optional=True,
|
|
tooltip="auto and gpu: the GPU encoder (NVENC) when it works, otherwise H.264 on the CPU so the video is still saved. cpu: H.264 (libx264) on every core."),
|
|
io.Combo.Input("codec", options=list(CODECS), default="h264", optional=True,
|
|
tooltip="For the GPU encoder. H.264: plays in every browser. H.265: smaller files, but Firefox may not play it."),
|
|
io.Combo.Input("quality", options=list(QUALITIES), default="high", optional=True,
|
|
tooltip="How closely the file keeps the picture. standard: small files, but fine texture smears and blocks show on a sharp high-resolution picture. high: visually clean for a bigger file. maximum: near lossless, large."),
|
|
io.Custom("FUNUI_SETTINGS").Input("settings", optional=True),
|
|
],
|
|
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
|
|
is_output_node=True,
|
|
outputs=[],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, images, fps=24.0, filename_prefix="FunUI", audio=None, encoder="auto", codec="h264", quality="high", settings=None) -> io.NodeOutput:
|
|
look = _values(settings)
|
|
frames = to_frames(images, look)
|
|
width, height = frames.shape[2], frames.shape[1]
|
|
folder, filename, counter, subfolder, _ = folder_paths.get_save_image_path(
|
|
filename_prefix, folder_paths.get_output_directory(), width, height)
|
|
metadata = None
|
|
if not args.disable_metadata:
|
|
metadata = {**(cls.hidden.extra_pnginfo or {}), **({"prompt": cls.hidden.prompt} if cls.hidden.prompt is not None else {})} or None
|
|
file = f"{filename}_{counter:05}_.mp4"
|
|
t = time.perf_counter()
|
|
used = encode(os.path.join(folder, file), frames, float(fps), audio, metadata, encoder, codec, quality)
|
|
log.info("FunUI Save Video", f"{file}: {len(frames)} frames in {time.perf_counter() - t:.1f}s with {used}, quality {quality}" + "".join(f", {k} {v}" for k, v in look.items() if v != DEFAULTS[k]))
|
|
keep_last_frame(file, to_frames(images[-1:])[-1]) # the sampler's own picture: the Post look belongs to the file, not to the next shot's start
|
|
return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
|