Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
48 lines
2.1 KiB
Python
48 lines
2.1 KiB
Python
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_demanded_without_one_says_so(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(nodes, "NVENC", {"h264": ("no_such_encoder_x", {})})
|
|
with pytest.raises(RuntimeError, match="NVENC"):
|
|
nodes.encode(str(tmp_path / "c.mp4"), _clip(5), 24.0, encoder="gpu")
|
|
|
|
|
|
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"
|