Files
olivv-cs-ComfyUI-FunPack/modules/output/save/tests/test_save.py
T

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"