118 lines
6.6 KiB
Python
118 lines
6.6 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 os
|
|
import time
|
|
from fractions import Fraction
|
|
|
|
import av
|
|
import folder_paths
|
|
from comfy.cli_args import args
|
|
from comfy_api.latest import io, ui
|
|
|
|
from ..._core import log
|
|
|
|
ENCODERS = ("auto", "gpu", "cpu")
|
|
CODECS = ("h264", "h265")
|
|
# NVENC's constant-quality mode at libx264's default number (crf 23).
|
|
NVENC = {"h265": ("hevc_nvenc", {"preset": "p5", "rc": "vbr", "cq": "23", "b:v": "0"}),
|
|
"h264": ("h264_nvenc", {"preset": "p5", "rc": "vbr", "cq": "23", "b:v": "0"})}
|
|
X264 = ("libx264", {"crf": "23"})
|
|
|
|
|
|
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"):
|
|
"""-> which encoder wrote the file. `frames`: uint8 [T, H, W, 3]. GPU first unless told otherwise;
|
|
a GPU attempt that fails leaves no half file and falls back to CPU H.264, unless "gpu" was demanded."""
|
|
gpu = NVENC[codec]
|
|
if encoder != "cpu" and gpu[0] in av.codecs_available:
|
|
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)
|
|
if encoder == "gpu":
|
|
raise RuntimeError(f"the GPU encoder (NVENC) failed: {exc}") from exc
|
|
log.once("save_nvenc", log.ALERT, "FunPack Save Video", f"GPU encoder unavailable ({exc}); using H.264 on the CPU")
|
|
elif encoder == "gpu":
|
|
raise RuntimeError(f"this PyAV build has no NVENC encoder ({gpu[0]}); pick auto or cpu")
|
|
_write(path, frames, fps, audio, metadata, X264)
|
|
if encoder == "auto":
|
|
log.once("save_no_nvenc", log.ALERT, "FunPack Save Video", f"no NVENC in this PyAV build ({gpu[0]} missing): saved on the CPU")
|
|
return "H.264 on libx264 (CPU, all cores)"
|
|
|
|
|
|
class FunPackSaveVideo(io.ComfyNode):
|
|
@classmethod
|
|
def define_schema(cls) -> io.Schema:
|
|
return io.Schema(
|
|
node_id="FunPackSaveVideo",
|
|
display_name="FunPack Save Video",
|
|
category="FunPack/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="FunPack"),
|
|
io.Combo.Input("encoder", options=list(ENCODERS), default="auto", optional=True,
|
|
tooltip="auto: GPU (NVENC) if available, else H.264 on the CPU. gpu: fail rather than use the CPU. 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."),
|
|
],
|
|
hidden=[io.Hidden.prompt, io.Hidden.extra_pnginfo],
|
|
is_output_node=True,
|
|
outputs=[],
|
|
)
|
|
|
|
@classmethod
|
|
def execute(cls, images, fps=24.0, filename_prefix="FunPack", audio=None, encoder="auto", codec="h264") -> io.NodeOutput:
|
|
frames = (images * 255).clamp(0, 255).byte().cpu().numpy()
|
|
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)
|
|
log.info("FunPack Save Video", f"{file}: {len(frames)} frames in {time.perf_counter() - t:.1f}s with {used}")
|
|
return io.NodeOutput(ui=ui.PreviewVideo([ui.SavedResult(file, subfolder, io.FolderType.output)]))
|