Claude
31871a5219
Save Video keeps the last decoded frame losslessly in temp; the carry uses it when the clip ends with the render. Co-Authored-By: Claude Sonnet 5.5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01Tw9fjYBYFyV8YG5YDfYPRb
103 lines
4.8 KiB
Python
103 lines
4.8 KiB
Python
import os
|
|
|
|
import av
|
|
import numpy as np
|
|
import pytest
|
|
|
|
from modules.output.save import nodes
|
|
|
|
|
|
def _clip(n=30):
|
|
rng = np.random.default_rng(0)
|
|
return rng.integers(0, 255, (n, 64, 96, 3), dtype=np.uint8)
|
|
|
|
|
|
def test_cpu_encode_writes_a_playable_h264_mp4_with_sound_and_metadata(tmp_path):
|
|
import torch
|
|
path = str(tmp_path / "a.mp4")
|
|
audio = {"waveform": torch.zeros(1, 2, 44100 * 2), "sample_rate": 44100}
|
|
used = nodes.encode(path, _clip(), 24.0, audio, {"prompt": {"1": "x"}}, encoder="cpu")
|
|
assert "libx264" in used
|
|
with av.open(path) as f:
|
|
v = f.streams.video[0]
|
|
assert v.codec_context.name == "h264" and v.frames == 30 and f.streams.audio
|
|
assert "prompt" in f.metadata
|
|
head = open(path, "rb").read(4096)
|
|
assert head.find(b"moov") != -1 and (head.find(b"mdat") == -1 or head.find(b"moov") < head.find(b"mdat")), "index first: seekable while streaming"
|
|
|
|
|
|
def test_auto_falls_back_to_the_cpu_when_there_is_no_gpu_encoder(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(nodes, "NVENC", {"h264": ("no_such_encoder_x", {})})
|
|
monkeypatch.setattr(av, "codecs_available", av.codecs_available | {"no_such_encoder_x"})
|
|
path = str(tmp_path / "b.mp4")
|
|
assert "libx264" in nodes.encode(path, _clip(10), 24.0, encoder="auto")
|
|
|
|
|
|
def test_gpu_choice_still_saves_on_the_cpu_when_the_encoder_is_missing(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(nodes, "NVENC", {"h264": ("no_such_encoder_x", {})})
|
|
path = str(tmp_path / "c.mp4")
|
|
assert "libx264" in nodes.encode(path, _clip(5), 24.0, encoder="gpu")
|
|
assert os.path.getsize(path) > 0
|
|
|
|
|
|
def test_a_refused_gpu_open_leaves_no_half_file_and_the_cpu_saves_it(tmp_path, monkeypatch):
|
|
# a codec that exists but refuses to open with these options, as NVENC does on a driver mismatch
|
|
monkeypatch.setattr(nodes, "NVENC", {"h264": ("libx264", {"preset": "no-such-preset-x", "crf": "abc"})})
|
|
alerts = []
|
|
monkeypatch.setattr(nodes.log, "once", lambda key, level, who, msg: alerts.append(msg))
|
|
path = str(tmp_path / "e.mp4")
|
|
assert "libx264" in nodes.encode(path, _clip(5), 24.0, encoder="gpu")
|
|
assert os.path.getsize(path) > 0
|
|
assert alerts and "Saved as H.264 on the CPU" in alerts[0]
|
|
|
|
|
|
def test_h265_is_tagged_hvc1_so_a_mac_browser_plays_it(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(nodes, "NVENC", {"h265": ("libx265", {"crf": "23"})}) # the software HEVC encoder stands in for NVENC here
|
|
if "libx265" not in av.codecs_available:
|
|
pytest.skip("no libx265 in this PyAV build")
|
|
path = str(tmp_path / "d.mp4")
|
|
assert "H265" in nodes.encode(path, _clip(10), 24.0, encoder="gpu", codec="h265")
|
|
with av.open(path) as f:
|
|
assert f.streams.video[0].codec_context.name == "hevc" and f.streams.video[0].codec_tag == "hvc1"
|
|
|
|
|
|
def test_nan_sound_is_dropped_and_the_video_still_saves(tmp_path):
|
|
import torch
|
|
path = str(tmp_path / "n.mp4")
|
|
wave = torch.zeros(1, 2, 44100 * 2)
|
|
wave[0, 0, 5] = float("nan")
|
|
nodes.encode(path, _clip(), 24.0, {"waveform": wave, "sample_rate": 44100}, None, encoder="cpu")
|
|
with av.open(path) as f:
|
|
assert f.streams.video[0].frames == 30 and not f.streams.audio
|
|
|
|
|
|
def test_quality_sets_the_constant_quality_number_and_high_is_the_default():
|
|
for name, key in (("libx264", "crf"), ("h264_nvenc", "cq")):
|
|
base = (name, {key: "23"})
|
|
assert nodes._with_quality(base, "standard")[1][key] == "23"
|
|
assert int(nodes._with_quality(base, "high")[1][key]) < 23
|
|
assert int(nodes._with_quality(base, "maximum")[1][key]) < int(nodes._with_quality(base, "high")[1][key])
|
|
assert nodes._with_quality(base, "nonsense")[1][key] == nodes._with_quality(base, "high")[1][key]
|
|
assert nodes.X264[1]["crf"] == "16"
|
|
|
|
|
|
def test_higher_quality_keeps_more_of_the_picture(tmp_path):
|
|
sizes = {}
|
|
for q in ("standard", "high"):
|
|
path = str(tmp_path / f"{q}.mp4")
|
|
nodes.encode(path, _clip(20), 24.0, encoder="cpu", quality=q)
|
|
sizes[q] = os.path.getsize(path)
|
|
assert sizes["high"] > sizes["standard"]
|
|
|
|
|
|
def test_the_last_decoded_picture_is_kept_exactly_for_the_next_shot(tmp_path, monkeypatch):
|
|
from PIL import Image
|
|
monkeypatch.setattr(nodes.folder_paths, "get_temp_directory", lambda: str(tmp_path / "temp"))
|
|
frame = np.random.default_rng(1).integers(0, 256, (12, 16, 3), dtype=np.uint8)
|
|
nodes.keep_last_frame("FunUI_00007_.mp4", frame)
|
|
name = tmp_path / "temp" / nodes.last_frame_name("FunUI_00007_.mp4")
|
|
assert name.name == "funui_last_FunUI_00007_.png"
|
|
assert np.array_equal(np.array(Image.open(name).convert("RGB")), frame) # lossless: every value as the sampler made it
|
|
monkeypatch.setattr(nodes.folder_paths, "get_temp_directory", lambda: (_ for _ in ()).throw(OSError("no temp")))
|
|
nodes.keep_last_frame("x.mp4", frame) # a failure here never fails the save
|