279 lines
14 KiB
Python
279 lines
14 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
|
|
|
|
|
|
def test_the_carried_picture_is_the_samplers_not_the_post_look(tmp_path, monkeypatch):
|
|
import torch
|
|
images = torch.rand(3, 16, 16, 3)
|
|
seen = []
|
|
monkeypatch.setattr(nodes, "keep_last_frame", lambda file, frame: seen.append(frame))
|
|
monkeypatch.setattr(nodes, "encode", lambda *a, **k: "libx264")
|
|
monkeypatch.setattr(nodes.folder_paths, "get_output_directory", lambda: str(tmp_path))
|
|
monkeypatch.setattr(nodes.folder_paths, "get_save_image_path", lambda *a: (str(tmp_path), "FunUI", 1, "", "FunUI"))
|
|
monkeypatch.setattr(nodes.FunUISaveVideo, "hidden", type("H", (), {"extra_pnginfo": None, "prompt": None})(), raising=False)
|
|
nodes.FunUISaveVideo.execute(images, settings={"output_save": {"saturation": 0.2, "smoothing": "strong"}})
|
|
assert np.array_equal(seen[0], nodes.to_frames(images[-1:])[-1]) # not the desaturated, smoothed one the file got
|
|
|
|
|
|
def test_a_failed_keep_never_leaves_an_older_videos_picture_under_the_name(tmp_path, monkeypatch):
|
|
monkeypatch.setattr(nodes.folder_paths, "get_temp_directory", lambda: str(tmp_path))
|
|
old = tmp_path / nodes.last_frame_name("FunUI_00001_.mp4")
|
|
old.write_bytes(b"picture of a deleted video")
|
|
monkeypatch.setattr("PIL.Image.Image.save", lambda *a, **k: (_ for _ in ()).throw(OSError("disk full")))
|
|
nodes.keep_last_frame("FunUI_00001_.mp4", np.zeros((4, 4, 3), dtype=np.uint8))
|
|
assert not old.exists()
|
|
|
|
|
|
# --- smoothing (edge anti-aliasing) ----------------------------------------------------------------------------------
|
|
|
|
def _staircase():
|
|
"""A hard diagonal edge (black / white), the picture a jagged edge makes: [T, H, W, 3] floats."""
|
|
import torch
|
|
yy, xx = torch.meshgrid(torch.arange(32), torch.arange(32), indexing="ij")
|
|
return (xx > yy).float()[None, :, :, None].repeat(2, 1, 1, 3)
|
|
|
|
|
|
def test_smoothing_off_is_the_old_path_and_the_edge_is_softened_only_where_it_is():
|
|
import torch
|
|
x = _staircase()
|
|
assert nodes.smooth(x, "off") is x
|
|
light, strong = nodes.smooth(x, "light"), nodes.smooth(x, "strong")
|
|
assert ((light - x).abs() > 0.05).any() and ((strong - x).abs() > (light - x).abs()).any()
|
|
assert torch.equal(light[:, 0, -1], x[:, 0, -1]) and torch.equal(light[:, -1, 0], x[:, -1, 0]) # flat corners untouched
|
|
flat = torch.full((1, 16, 16, 3), 0.4)
|
|
assert torch.allclose(nodes.smooth(flat, "strong"), flat)
|
|
assert torch.equal(light[0], light[1]) # frames never mix: spatial only
|
|
|
|
|
|
def test_smoothing_follows_the_device_and_dtype_of_its_input():
|
|
import torch
|
|
for dtype in (torch.float16, torch.bfloat16, torch.float32):
|
|
x = _staircase().to(dtype)
|
|
assert nodes.smooth(x, "strong").dtype == dtype
|
|
meta = torch.empty(2, 32, 32, 3, device="meta", dtype=torch.float16) # the cloud stand-in for "some other device"
|
|
out = nodes.smooth(meta, "light")
|
|
assert out.device.type == "meta" and out.dtype == torch.float16 and out.shape == meta.shape
|
|
|
|
|
|
def test_chunked_conversion_equals_one_go_and_off_equals_the_plain_cast():
|
|
import torch
|
|
x = torch.rand(nodes._CHUNK * 2 + 3, 16, 16, 3)
|
|
plain = (x * 255).round().clamp(0, 255).byte().numpy()
|
|
assert (nodes.to_frames(x) == plain).all()
|
|
whole = (nodes.smooth(x, "strong") * 255).round().clamp(0, 255).byte().numpy()
|
|
assert (nodes.to_frames(x, {"smoothing": "strong"}) == whole).all()
|
|
|
|
|
|
def _ends(spec):
|
|
"""The values a declared setting is tried at: every choice, both switch positions, a slider's two ends and its default, a text or colour its default and one more."""
|
|
if spec["type"] == "enum":
|
|
return [o["value"] for o in spec["options"]]
|
|
if spec["type"] == "bool":
|
|
return [False, True]
|
|
if spec["type"] == "text":
|
|
return [spec["default"], "FunUI"]
|
|
if spec["type"] == "color":
|
|
return [spec["default"], "#ff8800"]
|
|
return [spec["min"], spec["default"], spec["max"]]
|
|
|
|
|
|
def _combos():
|
|
"""Registry-driven: every on/off combination of the picture's look settings, each with the watermark off and with it on at every
|
|
setting's other values (one at a time against an all-on mark, so the count stays small: the marks do not touch each other)."""
|
|
import itertools
|
|
from modules.output import save
|
|
look = {k: v for k, v in save.SETTINGS.items() if not k.startswith("watermark")}
|
|
mark = {k: v for k, v in save.SETTINGS.items() if k.startswith("watermark")}
|
|
looks = [dict(c) for c in itertools.product(*[[(k, v) for v in _ends(spec)] for k, spec in look.items()])]
|
|
on = {"watermark": "text", "watermark_text": "FunUI", "watermark_outline": 6.0}
|
|
marks = [{}, on]
|
|
for key, spec in mark.items():
|
|
marks += [{**on, key: v} for v in _ends(spec) if v != on.get(key)]
|
|
return [{**a, **b} for a in looks for b in ([{}, on])] + [{**looks[-1], **b} for b in marks]
|
|
|
|
|
|
def test_every_post_combination_saves_and_every_encoder_codec_quality_takes_the_all_on_look(tmp_path):
|
|
# registry-driven: no NVENC here, so "auto"/"gpu" are the NVENC -> CPU fallback
|
|
import itertools
|
|
x = _dull()
|
|
for n, values in enumerate(_combos()):
|
|
path = str(tmp_path / f"c{n}.mp4")
|
|
assert nodes.encode(path, nodes.to_frames(x, values), 24.0, encoder="auto") and os.path.getsize(path) > 0
|
|
everything = {"smoothing": "strong", "auto_contrast": True, "auto_brightness": True, "saturation": 1.5, "watermark": "text", "watermark_text": "FunUI", "watermark_outline": 6.0}
|
|
frames = nodes.to_frames(x, everything)
|
|
for encoder, codec, quality in itertools.product(nodes.ENCODERS, nodes.CODECS, nodes.QUALITIES):
|
|
path = str(tmp_path / f"{encoder}_{codec}_{quality}.mp4")
|
|
assert nodes.encode(path, frames, 24.0, encoder=encoder, codec=codec, quality=quality)
|
|
with av.open(path) as f:
|
|
assert f.streams.video[0].width == 16
|
|
|
|
|
|
def test_the_node_reads_its_settings_from_the_pipelines_settings_object():
|
|
assert nodes._values(None) == nodes.DEFAULTS
|
|
assert nodes._values({"output_save": {"smoothing": "light", "saturation": 1.2}})["smoothing"] == "light"
|
|
assert nodes._values({"output_save": {"smoothing": "bogus"}})["smoothing"] == "off" # wrong value: default, said
|
|
|
|
|
|
def test_post_settings_are_declared_off_and_never_become_node_sockets():
|
|
from modules.output import save
|
|
assert {k: v["default"] for k, v in save.SETTINGS.items()} == nodes.DEFAULTS
|
|
assert "smoothing" not in [i.id for i in nodes.FunUISaveVideo.define_schema().inputs]
|
|
assert "settings" in [i.id for i in nodes.FunUISaveVideo.define_schema().inputs]
|
|
|
|
|
|
def _dull():
|
|
import torch
|
|
torch.manual_seed(0)
|
|
return (0.3 + 0.2 * torch.rand(20, 16, 16, 3)) # a flat, grey, dark clip: levels 0.3..0.5
|
|
|
|
|
|
def test_auto_contrast_stretches_one_fixed_curve_for_the_whole_clip():
|
|
import torch
|
|
x = _dull()
|
|
look = nodes.clip_look(x, {"auto_contrast": True})
|
|
assert look["lo"] == 0.15 and look["hi"] == 0.85 # capped: a dull clip is stretched only so far
|
|
out = nodes.finish(x, {"auto_contrast": True}, look)
|
|
assert out.std() > x.std()
|
|
same = nodes.finish(x[:3], {"auto_contrast": True}, look) # a chunk gets the same curve as the whole: nothing is per frame
|
|
assert torch.equal(same, out[:3])
|
|
assert nodes.clip_look(x, {}) == {"lo": 0.0, "hi": 1.0, "gamma": 1.0}
|
|
|
|
|
|
def test_auto_brightness_moves_the_midtones_toward_the_middle_within_limits():
|
|
import torch
|
|
dark, bright = _dull() * 0.3, 0.7 + 0.3 * _dull()
|
|
gd, gb = nodes.clip_look(dark, {"auto_brightness": True})["gamma"], nodes.clip_look(bright, {"auto_brightness": True})["gamma"]
|
|
assert gd == 0.75 and gb > 1.0 and gb <= 1.33
|
|
assert nodes.finish(dark, {"auto_brightness": True}, {"lo": 0.0, "hi": 1.0, "gamma": gd}).mean() > dark.mean()
|
|
assert nodes.clip_look(torch.zeros(0, 4, 4, 3), {"auto_contrast": True})["gamma"] == 1.0 # an empty clip is not measured
|
|
|
|
|
|
def test_saturation_mutes_and_boosts_and_neutral_changes_nothing():
|
|
import torch
|
|
x = torch.rand(2, 8, 8, 3)
|
|
look = {"lo": 0.0, "hi": 1.0, "gamma": 1.0}
|
|
grey = nodes.finish(x, {"saturation": 0.5}, look)
|
|
assert (grey.max(-1).values - grey.min(-1).values).mean() < (x.max(-1).values - x.min(-1).values).mean()
|
|
assert torch.equal(nodes.finish(x, {"saturation": 1.0}, look), x)
|
|
|
|
|
|
def test_every_look_follows_the_device_and_dtype_of_its_input():
|
|
import torch
|
|
values = {"smoothing": "strong", "auto_contrast": True, "auto_brightness": True, "saturation": 1.3}
|
|
for dtype in (torch.float16, torch.bfloat16, torch.float32):
|
|
x = _dull().to(dtype)
|
|
out = nodes.finish(x, values, nodes.clip_look(x, values))
|
|
assert out.dtype == dtype and out.shape == x.shape
|
|
meta = torch.empty(2, 16, 16, 3, device="meta", dtype=torch.float16)
|
|
out = nodes.finish(meta, {**values}, {"lo": 0.1, "hi": 0.9, "gamma": 0.9})
|
|
assert out.device.type == "meta" and out.dtype == torch.float16
|
|
|
|
|
|
def test_saved_pipelines_get_the_settings_link_and_a_saved_pipeline_keeps_its_own():
|
|
from modules.output import save
|
|
slots = [{"id": "s", "node": "FunUIModifierSettings", "inputs": {}}, {"id": "v", "node": "FunUISaveVideo", "inputs": {"images": ["x", 0]}}]
|
|
assert save.upgrade(slots, lambda n: [])[1]["inputs"]["settings"] == ["s", 0]
|
|
own = [slots[0], {"id": "v", "node": "FunUISaveVideo", "inputs": {"settings": ["other", 0]}}]
|
|
assert save.upgrade(own, lambda n: [])[1]["inputs"]["settings"] == ["other", 0]
|
|
assert "settings" not in save.upgrade(slots[1:], lambda n: [])[0]["inputs"] # nothing to link to
|