1665 lines
65 KiB
Python
1665 lines
65 KiB
Python
import dataclasses
|
|
import json
|
|
import shutil
|
|
import subprocess
|
|
import os
|
|
from types import SimpleNamespace
|
|
import warnings
|
|
|
|
import numpy as np
|
|
import pytest
|
|
import torch
|
|
from einops import rearrange
|
|
|
|
import fastvideo.entrypoints.video_generator as video_generator_module
|
|
import fastvideo.envs as envs
|
|
import fastvideo.utils as fastvideo_utils
|
|
from fastvideo.api import (
|
|
ConfigValidationError,
|
|
GenerationRequest,
|
|
GenerationResult,
|
|
GeneratorConfig,
|
|
InputConfig,
|
|
SamplingConfig,
|
|
load_run_config,
|
|
)
|
|
from fastvideo.api.sampling_param import SamplingParam
|
|
from fastvideo.entrypoints.video_generator import VideoGenerator, _resolve_output_size
|
|
from fastvideo.fastvideo_args import WorkloadType
|
|
from fastvideo.pipelines import ForwardBatch
|
|
from fastvideo.worker.gpu_worker import Worker
|
|
from fastvideo.worker.ray_distributed_executor import RayDistributedExecutor
|
|
|
|
|
|
def _new_video_generator() -> VideoGenerator:
|
|
# Bypass __init__ since we only test a pure helper method.
|
|
return VideoGenerator.__new__(VideoGenerator)
|
|
|
|
|
|
def _new_runtime_video_generator() -> VideoGenerator:
|
|
generator = _new_video_generator()
|
|
generator.fastvideo_args = SimpleNamespace(
|
|
model_path="test-model",
|
|
prompt_txt=None,
|
|
output_type="pil",
|
|
workload_type=SimpleNamespace(value="t2v"),
|
|
pipeline_config=SimpleNamespace(apply_request_constraints=lambda request, sampling_param: sampling_param),
|
|
)
|
|
generator.executor = SimpleNamespace(
|
|
set_log_queue=lambda queue: None,
|
|
clear_log_queue=lambda: None,
|
|
is_output_rank=True,
|
|
uses_spmd_execution=False,
|
|
broadcast_from_output_rank=lambda value: value,
|
|
)
|
|
generator.config = None
|
|
return generator
|
|
|
|
|
|
def _patch_from_fastvideo_args(monkeypatch):
|
|
captured = {}
|
|
|
|
def fake_from_fastvideo_args(cls, fastvideo_args, *, log_queue=None):
|
|
generator = cls.__new__(cls)
|
|
generator.fastvideo_args = fastvideo_args
|
|
generator.executor = None
|
|
generator.config = None
|
|
captured["fastvideo_args"] = fastvideo_args
|
|
captured["log_queue"] = log_queue
|
|
return generator
|
|
|
|
monkeypatch.setattr(
|
|
VideoGenerator,
|
|
"from_fastvideo_args",
|
|
classmethod(fake_from_fastvideo_args),
|
|
)
|
|
return captured
|
|
|
|
|
|
def _patch_fastvideo_args_from_kwargs(monkeypatch):
|
|
captured = {}
|
|
|
|
def fake_from_kwargs(cls, **kwargs):
|
|
captured["kwargs"] = kwargs
|
|
return SimpleNamespace(
|
|
model_path=kwargs["model_path"],
|
|
num_gpus=kwargs["num_gpus"],
|
|
workload_type=WorkloadType.from_string(kwargs.get("workload_type", "t2v")),
|
|
)
|
|
|
|
monkeypatch.setattr(
|
|
"fastvideo.api.compat.FastVideoArgs.from_kwargs",
|
|
classmethod(fake_from_kwargs),
|
|
)
|
|
return captured
|
|
|
|
|
|
def _patch_sampling_param_from_pretrained(monkeypatch):
|
|
|
|
def fake_from_pretrained(cls, model_path):
|
|
return cls()
|
|
|
|
monkeypatch.setattr(SamplingParam, "from_pretrained", classmethod(fake_from_pretrained))
|
|
|
|
|
|
class _NoCpuMaterializationOutput:
|
|
|
|
@property
|
|
def shape(self):
|
|
raise AssertionError("metadata-only output path should not inspect output shape")
|
|
|
|
def cpu(self):
|
|
raise AssertionError("metadata-only output path should not move output to CPU")
|
|
|
|
def __mul__(self, other):
|
|
raise AssertionError("metadata-only output path should not build frames")
|
|
|
|
|
|
def _single_video_args(output_type="pil"):
|
|
return SimpleNamespace(
|
|
output_type=output_type,
|
|
pin_cpu_memory=False,
|
|
VSA_sparsity=0.0,
|
|
pipeline_config=SimpleNamespace(flow_shift=1.0,
|
|
embedded_cfg_scale=1.0,
|
|
apply_request_constraints=lambda request, sampling_param: sampling_param),
|
|
workload_type=SimpleNamespace(value="t2v"),
|
|
)
|
|
|
|
|
|
def _small_sampling_param(*, save_video=False, return_frames=False):
|
|
return SamplingParam(
|
|
num_frames=2,
|
|
height=16,
|
|
width=16,
|
|
fps=8,
|
|
num_inference_steps=1,
|
|
save_video=save_video,
|
|
return_frames=return_frames,
|
|
)
|
|
|
|
|
|
def _single_video_output_batch(output, *, extra=None):
|
|
return SimpleNamespace(
|
|
output=output,
|
|
extra=extra or {},
|
|
logging_info=None,
|
|
trajectory_latents=None,
|
|
trajectory_timesteps=None,
|
|
trajectory_decoded=None,
|
|
)
|
|
|
|
|
|
def _single_video_generator(output_batch, fastvideo_args):
|
|
generator = _new_video_generator()
|
|
generator.fastvideo_args = fastvideo_args
|
|
generator.executor = SimpleNamespace(
|
|
execute_forward=lambda batch, args: output_batch,
|
|
is_output_rank=True,
|
|
uses_spmd_execution=False,
|
|
broadcast_from_output_rank=lambda value: value,
|
|
)
|
|
generator.config = None
|
|
return generator
|
|
|
|
|
|
def test_run_on_output_rank_non_root_does_not_touch_local_filesystem():
|
|
generator = _new_video_generator()
|
|
operation = pytest.fail
|
|
generator.executor = SimpleNamespace(
|
|
is_output_rank=False,
|
|
uses_spmd_execution=True,
|
|
broadcast_from_output_rank=lambda value: (True, "/rank-zero/output.mp4"),
|
|
)
|
|
|
|
assert generator._run_on_output_rank(operation, description="output-path setup") == "/rank-zero/output.mp4"
|
|
|
|
|
|
def test_run_on_output_rank_propagates_rank_zero_error():
|
|
generator = _new_video_generator()
|
|
generator.executor = SimpleNamespace(
|
|
is_output_rank=True,
|
|
uses_spmd_execution=True,
|
|
broadcast_from_output_rank=lambda value: value,
|
|
)
|
|
|
|
def fail_setup():
|
|
raise PermissionError("output mount is read-only")
|
|
|
|
with pytest.raises(RuntimeError, match="output mount is read-only"):
|
|
generator._run_on_output_rank(fail_setup, description="output-path setup")
|
|
|
|
|
|
def test_resolve_output_size_uses_produced_pixel_geometry() -> None:
|
|
"""Report refined dimensions instead of the base request geometry."""
|
|
samples = torch.empty(1, 3, 5, 64, 96)
|
|
assert _resolve_output_size(samples, (32, 48, 1), pixel_output=True) == (64, 96, 5)
|
|
assert _resolve_output_size(samples, (32, 48, 1), pixel_output=False) == (32, 48, 1)
|
|
|
|
|
|
def test_prepare_output_path_file_sanitization(tmp_path):
|
|
vg = _new_video_generator()
|
|
target_dir = tmp_path / "dir"
|
|
raw_path = target_dir / "inv:al*id?.mp4"
|
|
|
|
result = vg._prepare_output_path(str(raw_path), prompt="ignored")
|
|
|
|
assert os.path.dirname(result) == str(target_dir)
|
|
assert os.path.basename(result) == "invalid.mp4"
|
|
assert os.path.isdir(target_dir)
|
|
|
|
|
|
def test_prepare_output_path_directory_prompt_derived(tmp_path):
|
|
vg = _new_video_generator()
|
|
out_dir = tmp_path / "outputs"
|
|
prompt = "Hello:/\\*?<>| world"
|
|
|
|
result = vg._prepare_output_path(str(out_dir), prompt=prompt)
|
|
|
|
assert os.path.dirname(result) == str(out_dir)
|
|
# spaces are preserved (collapsed) by sanitizer; here it becomes "Hello world.mp4"
|
|
assert os.path.basename(result) == "Hello world.mp4"
|
|
assert os.path.isdir(out_dir)
|
|
|
|
|
|
def test_prepare_output_path_non_mp4_treated_as_dir(tmp_path):
|
|
vg = _new_video_generator()
|
|
weird_dir = tmp_path / "foo.gif"
|
|
prompt = "My Video"
|
|
|
|
result = vg._prepare_output_path(str(weird_dir), prompt=prompt)
|
|
|
|
assert os.path.dirname(result) == str(weird_dir)
|
|
assert os.path.basename(result) == "My Video.mp4"
|
|
assert os.path.isdir(weird_dir)
|
|
|
|
|
|
def test_prepare_output_path_uniqueness_suffix(tmp_path):
|
|
vg = _new_video_generator()
|
|
out_dir = tmp_path / "outputs"
|
|
prompt = "Sample Name"
|
|
|
|
first = vg._prepare_output_path(str(out_dir), prompt=prompt)
|
|
# simulate existing file
|
|
os.makedirs(os.path.dirname(first), exist_ok=True)
|
|
with open(first, "wb") as f:
|
|
f.write(b"")
|
|
|
|
second = vg._prepare_output_path(str(out_dir), prompt=prompt)
|
|
assert os.path.basename(second) == "Sample Name_1.mp4"
|
|
|
|
# simulate second existing file as well
|
|
with open(second, "wb") as f:
|
|
f.write(b"")
|
|
third = vg._prepare_output_path(str(out_dir), prompt=prompt)
|
|
assert os.path.basename(third) == "Sample Name_2.mp4"
|
|
|
|
|
|
def test_prepare_output_path_empty_prompt_fallback(tmp_path):
|
|
vg = _new_video_generator()
|
|
out_dir = tmp_path / "outputs"
|
|
bad_prompt = ":/\\*?<>| .." # sanitizes to empty, should fallback to "video"
|
|
|
|
result = vg._prepare_output_path(str(out_dir), prompt=bad_prompt)
|
|
|
|
assert os.path.dirname(result) == str(out_dir)
|
|
assert os.path.basename(result) == "output.mp4"
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ["generate", "generate_video", "sampling_param"])
|
|
def test_generate_honors_output_video_name(monkeypatch, tmp_path, entry):
|
|
generator = _new_runtime_video_generator()
|
|
_patch_sampling_param_from_pretrained(monkeypatch)
|
|
captured = {}
|
|
|
|
def fake_generate_single_video(prompt, sampling_param, **kwargs):
|
|
captured["prompt"] = prompt
|
|
captured["name"] = sampling_param.output_video_name
|
|
return {"prompts": prompt, "video_path": kwargs["output_path"]}
|
|
|
|
monkeypatch.setattr(generator, "_generate_single_video", fake_generate_single_video)
|
|
output = {"output_path": str(tmp_path), "output_video_name": "my_clip"}
|
|
if entry == "generate":
|
|
result = generator.generate({"prompt": "a red fox", "output": output})
|
|
output_path = result.video_path
|
|
elif entry == "generate_video":
|
|
with pytest.warns(DeprecationWarning):
|
|
result = generator.generate_video(prompt="a red fox", **output)
|
|
output_path = result["video_path"]
|
|
else:
|
|
result = generator._generate_video_impl("a red fox", sampling_param=SamplingParam(**output))
|
|
output_path = result["video_path"]
|
|
|
|
assert output_path == str(tmp_path / "my_clip.mp4")
|
|
assert captured == {"prompt": "a red fox", "name": "my_clip"}
|
|
|
|
|
|
@pytest.mark.parametrize("name,expected", [
|
|
(None, "a red fox.mp4"),
|
|
("", "a red fox.mp4"),
|
|
("my_clip", "my_clip.mp4"),
|
|
("my_clip.mp4", "my_clip.mp4"),
|
|
("my_clip.MP4", "my_clip.mp4"),
|
|
("my.clip.v1", "my.clip.v1.mp4"),
|
|
(" my:clip?*.mp4", "myclip.mp4"),
|
|
(r"../nested\my:clip?.mp4", "nestedmyclip.mp4"),
|
|
(":/\\*?<>| ..", "output.mp4"),
|
|
("a" * 120, "a" * 120 + ".mp4"),
|
|
])
|
|
def test_prepare_output_path_requested_name(tmp_path, name, expected):
|
|
generator = _new_video_generator()
|
|
directory = tmp_path / "outputs"
|
|
|
|
result = generator._prepare_output_path(str(directory), "a red fox", output_video_name=name)
|
|
|
|
assert result == str(directory / expected)
|
|
assert directory.is_dir()
|
|
|
|
|
|
@pytest.mark.parametrize("workload,extension", [("t2v", ".mp4"), ("t2i", ".jpg"), ("t2a", ".wav")])
|
|
def test_prepare_output_path_explicit_file_overrides_requested_name(tmp_path, workload, extension):
|
|
generator = _new_runtime_video_generator()
|
|
generator.fastvideo_args.workload_type = WorkloadType.from_string(workload)
|
|
target = tmp_path / ("explicit" + extension)
|
|
target.write_bytes(b"existing output")
|
|
|
|
result = generator._prepare_output_path(str(target), "prompt", output_video_name="ignored.png")
|
|
|
|
assert result == str(tmp_path / ("explicit_1" + extension))
|
|
assert target.read_bytes() == b"existing output"
|
|
|
|
|
|
@pytest.mark.parametrize("workload,name,expected", [
|
|
("t2i", "my_image", "my_image.png"),
|
|
("t2i", "my_image.png", "my_image.png"),
|
|
("t2i", "my_image.jpg", "my_image.jpg"),
|
|
("t2i", "my_image.JPEG", "my_image.jpeg"),
|
|
("t2i", "my_image.webp", "my_image.webp"),
|
|
("t2a", "my_audio", "my_audio.wav"),
|
|
("t2a", "my_audio.wav", "my_audio.wav"),
|
|
])
|
|
def test_prepare_output_path_requested_name_uses_workload_format(tmp_path, workload, name, expected):
|
|
generator = _new_runtime_video_generator()
|
|
generator.fastvideo_args.workload_type = WorkloadType.from_string(workload)
|
|
|
|
assert generator._prepare_output_path(str(tmp_path), "prompt", output_video_name=name) == str(tmp_path / expected)
|
|
|
|
|
|
def test_prepare_output_path_requested_name_in_extension_directory(tmp_path):
|
|
generator = _new_video_generator()
|
|
directory = tmp_path / "outputs.mp4"
|
|
directory.mkdir()
|
|
existing = directory / "my_clip.mp4"
|
|
existing.write_bytes(b"existing output")
|
|
|
|
result = generator._prepare_output_path(str(directory), "prompt", output_video_name="my_clip.mp4")
|
|
|
|
assert result == str(directory / "my_clip_1.mp4")
|
|
assert existing.read_bytes() == b"existing output"
|
|
|
|
|
|
@pytest.mark.parametrize("batch_source", ["prompt_file", "prompt_list"])
|
|
def test_generate_batch_honors_output_video_name(monkeypatch, tmp_path, batch_source):
|
|
generator = _new_runtime_video_generator()
|
|
_patch_sampling_param_from_pretrained(monkeypatch)
|
|
captured_prompts = []
|
|
|
|
def fake_generate_single_video(prompt, sampling_param, **kwargs):
|
|
captured_prompts.append(prompt)
|
|
assert sampling_param.output_video_name == "my_clip.mp4"
|
|
with open(kwargs["output_path"], "wb") as output:
|
|
output.write(b"generated output")
|
|
return {"prompts": prompt, "video_path": kwargs["output_path"]}
|
|
|
|
monkeypatch.setattr(generator, "_generate_single_video", fake_generate_single_video)
|
|
request = {"output": {"output_path": str(tmp_path), "output_video_name": "my_clip.mp4"}}
|
|
if batch_source == "prompt_file":
|
|
prompt_file = tmp_path / "prompts.txt"
|
|
prompt_file.write_text("first prompt\nsecond prompt\n", encoding="utf-8")
|
|
request["inputs"] = {"prompt_path": str(prompt_file)}
|
|
else:
|
|
request["prompt"] = ["first prompt", "second prompt"]
|
|
|
|
results = generator.generate(request)
|
|
|
|
assert [result.video_path for result in results] == [str(tmp_path / "my_clip.mp4"), str(tmp_path / "my_clip_1.mp4")]
|
|
assert [result.prompt_index for result in results] == [0, 1]
|
|
assert captured_prompts == ["first prompt", "second prompt"]
|
|
|
|
|
|
@pytest.mark.parametrize("name,expected", [(None, "input.mp4"), ("custom", "custom.mp4")])
|
|
def test_generate_optional_prompt_respects_output_video_name(monkeypatch, tmp_path, name, expected):
|
|
generator = _new_runtime_video_generator()
|
|
generator.fastvideo_args.pipeline_config.prompt_optional = True
|
|
monkeypatch.setattr(generator, "_generate_single_video", lambda **kwargs: kwargs)
|
|
sampling = SamplingParam(video_path="/data/input.mp4", output_path=str(tmp_path), output_video_name=name)
|
|
|
|
result = generator._generate_video_impl(prompt=None, sampling_param=sampling)
|
|
|
|
assert result["prompt"] == ""
|
|
assert result["output_path"] == str(tmp_path / expected)
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["a" * 256, "a" * 252, "\u754c" * 84])
|
|
@pytest.mark.parametrize("explicit_path", [False, True])
|
|
def test_prepare_output_path_rejects_overlong_filename(monkeypatch, tmp_path, name, explicit_path):
|
|
generator = _new_video_generator()
|
|
monkeypatch.setattr(video_generator_module.os, "pathconf", lambda *args: 255, raising=False)
|
|
output_path = str(tmp_path / (name + ".mp4")) if explicit_path else str(tmp_path)
|
|
|
|
with pytest.raises(ValueError, match="Output filename exceeds the 255-byte filesystem limit"):
|
|
generator._prepare_output_path(output_path, "prompt", output_video_name=name)
|
|
|
|
|
|
@pytest.mark.parametrize("name", ["a" * 251, "\u754c" * 83 + "ab", "?" * 300 + "clip"])
|
|
def test_prepare_output_path_checks_sanitized_filename_bytes(monkeypatch, tmp_path, name):
|
|
generator = _new_video_generator()
|
|
monkeypatch.setattr(video_generator_module.os, "pathconf", lambda *args: 255, raising=False)
|
|
|
|
result = generator._prepare_output_path(str(tmp_path), "prompt", output_video_name=name)
|
|
|
|
assert os.path.basename(result) == name.replace("?", "") + ".mp4"
|
|
assert len(os.fsencode(os.path.basename(result))) <= 255
|
|
with open(result, "wb") as output:
|
|
output.write(b"generated output")
|
|
|
|
|
|
@pytest.mark.parametrize("existing_count,stem_length,rejected", [(1, 249, False), (1, 251, True), (10, 249, True)])
|
|
def test_prepare_output_path_checks_collision_suffix_bytes(monkeypatch, tmp_path, existing_count, stem_length, rejected):
|
|
generator = _new_video_generator()
|
|
monkeypatch.setattr(video_generator_module.os, "pathconf", lambda *args: 255, raising=False)
|
|
name = "a" * stem_length
|
|
existing_files = [tmp_path / (name + ".mp4")]
|
|
existing_files.extend(tmp_path / f"{name}_{index}.mp4" for index in range(1, existing_count))
|
|
for path in existing_files:
|
|
path.write_bytes(b"existing output")
|
|
|
|
if rejected:
|
|
with pytest.raises(ValueError, match="Output filename exceeds the 255-byte filesystem limit"):
|
|
generator._prepare_output_path(str(tmp_path), "prompt", output_video_name=name)
|
|
else:
|
|
result = generator._prepare_output_path(str(tmp_path), "prompt", output_video_name=name)
|
|
assert result == str(tmp_path / f"{name}_{existing_count}.mp4")
|
|
assert len(os.fsencode(os.path.basename(result))) == 255
|
|
with open(result, "wb") as output:
|
|
output.write(b"generated output")
|
|
assert all(path.read_bytes() == b"existing output" for path in existing_files)
|
|
|
|
|
|
def test_prepare_output_path_respects_filesystem_name_limit(monkeypatch, tmp_path):
|
|
generator = _new_video_generator()
|
|
directory = tmp_path / "new_output_directory"
|
|
|
|
def pathconf(path, key):
|
|
assert path == str(directory)
|
|
assert directory.is_dir()
|
|
assert key == "PC_NAME_MAX"
|
|
return 20
|
|
|
|
monkeypatch.setattr(video_generator_module.os, "pathconf", pathconf, raising=False)
|
|
with pytest.raises(ValueError, match="Output filename exceeds the 20-byte filesystem limit"):
|
|
generator._prepare_output_path(str(directory), "prompt", output_video_name="a" * 17)
|
|
assert generator._prepare_output_path(str(directory), "prompt", output_video_name="a" * 16) == str(
|
|
directory / ("a" * 16 + ".mp4"))
|
|
|
|
|
|
@pytest.mark.parametrize("unavailable", [AttributeError, OSError, ValueError, -1, 0])
|
|
def test_prepare_output_path_filename_limit_fallback(monkeypatch, tmp_path, unavailable):
|
|
generator = _new_video_generator()
|
|
|
|
def pathconf(*args):
|
|
if isinstance(unavailable, type):
|
|
raise unavailable("name limit unavailable")
|
|
return unavailable
|
|
|
|
monkeypatch.setattr(video_generator_module.os, "pathconf", pathconf, raising=False)
|
|
with pytest.raises(ValueError, match="Output filename exceeds the 255-byte filesystem limit"):
|
|
generator._prepare_output_path(str(tmp_path), "prompt", output_video_name="a" * 252)
|
|
assert generator._prepare_output_path(str(tmp_path), "prompt", output_video_name="clip") == str(tmp_path / "clip.mp4")
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ["generate", "generate_video", "prompt_file"])
|
|
@pytest.mark.parametrize("collision", [False, True])
|
|
def test_generate_rejects_overlong_filename_before_model_execution(monkeypatch, tmp_path, entry, collision):
|
|
generator = _new_runtime_video_generator()
|
|
_patch_sampling_param_from_pretrained(monkeypatch)
|
|
monkeypatch.setattr(video_generator_module.os, "pathconf", lambda *args: 255, raising=False)
|
|
executed = []
|
|
|
|
def fake_generate_single_video(**kwargs):
|
|
executed.append(kwargs)
|
|
return {"prompts": kwargs["prompt"], "video_path": kwargs["output_path"]}
|
|
|
|
monkeypatch.setattr(generator, "_generate_single_video", fake_generate_single_video)
|
|
name = "a" * (251 if collision else 256)
|
|
if collision:
|
|
(tmp_path / (name + ".mp4")).write_bytes(b"existing output")
|
|
output = {"output_path": str(tmp_path), "output_video_name": name}
|
|
|
|
if entry == "prompt_file":
|
|
prompt_file = tmp_path / "prompts.txt"
|
|
prompt_file.write_text("first prompt\nsecond prompt\n", encoding="utf-8")
|
|
assert generator.generate({"inputs": {"prompt_path": str(prompt_file)}, "output": output}) == []
|
|
elif entry == "generate":
|
|
with pytest.raises(ValueError, match="Output filename exceeds"):
|
|
generator.generate({"prompt": "prompt", "output": output})
|
|
else:
|
|
with pytest.warns(DeprecationWarning), pytest.raises(ValueError, match="Output filename exceeds"):
|
|
generator.generate_video(prompt="prompt", **output)
|
|
assert not executed
|
|
if collision:
|
|
assert (tmp_path / (name + ".mp4")).read_bytes() == b"existing output"
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ["generate", "generate_video", "prompt_file", "prompt_list"])
|
|
@pytest.mark.parametrize("naming", ["prompt", "custom", "explicit"])
|
|
@pytest.mark.parametrize("output_type,save_video,return_frames,return_samples", [
|
|
("pil", False, True, False),
|
|
("pil", False, False, True),
|
|
("latent", True, False, True),
|
|
], ids=["frames", "samples", "latent"])
|
|
def test_generate_in_memory_preserves_results_with_long_names(
|
|
monkeypatch, tmp_path, entry, naming, output_type, save_video, return_frames, return_samples
|
|
):
|
|
prompt = "\u754c" * 100
|
|
output = torch.ones((1, 16, 1, 2, 2) if output_type == "latent" else (1, 3, 2, 16, 16)) * 0.5
|
|
fastvideo_args = _single_video_args(output_type=output_type)
|
|
fastvideo_args.model_path = "test-model"
|
|
fastvideo_args.prompt_txt = None
|
|
generator = _single_video_generator(_single_video_output_batch(output), fastvideo_args)
|
|
_patch_sampling_param_from_pretrained(monkeypatch)
|
|
directory = tmp_path / "unused_outputs"
|
|
output_path = directory / ("a" * 256 + ".mp4") if naming == "explicit" else directory
|
|
output_options = {
|
|
"output_path": str(output_path),
|
|
"output_video_name": "a" * 256 if naming == "custom" else None,
|
|
"save_video": save_video,
|
|
"return_frames": return_frames,
|
|
"return_samples": return_samples,
|
|
}
|
|
sampling_options = {"height": 16, "width": 16, "num_frames": 2, "num_inference_steps": 1}
|
|
request = {"prompt": prompt, "sampling": sampling_options, "output": output_options}
|
|
if entry == "prompt_file":
|
|
prompt_file = tmp_path / "prompts.txt"
|
|
prompt_file.write_text(f"{prompt}\n{prompt}\n", encoding="utf-8")
|
|
request["prompt"] = None
|
|
request["inputs"] = {"prompt_path": str(prompt_file)}
|
|
elif entry == "prompt_list":
|
|
request["prompt"] = [prompt, prompt]
|
|
|
|
def fail_path_operation(*args, **kwargs):
|
|
pytest.fail("In-memory generation must not prepare an output file")
|
|
|
|
with monkeypatch.context() as patches:
|
|
patches.setattr(video_generator_module.os, "pathconf", fail_path_operation, raising=False)
|
|
patches.setattr(video_generator_module.os, "makedirs", fail_path_operation)
|
|
if entry == "generate_video":
|
|
with pytest.warns(DeprecationWarning):
|
|
results = [generator.generate_video(prompt=prompt, **sampling_options, **output_options)]
|
|
else:
|
|
generated = generator.generate(request)
|
|
results = generated if isinstance(generated, list) else [generated]
|
|
results = [result.to_legacy_dict() for result in results]
|
|
|
|
assert len(results) == (2 if entry in {"prompt_file", "prompt_list"} else 1)
|
|
for result in results:
|
|
assert result["prompts"] == prompt
|
|
torch.testing.assert_close(result["samples"], output)
|
|
assert result["video_path"] is None
|
|
if return_frames:
|
|
assert len(result["frames"]) == 2
|
|
assert result["frames"][0].shape == (16, 16, 3)
|
|
else:
|
|
assert result["frames"] is None
|
|
assert not directory.exists()
|
|
|
|
|
|
@pytest.mark.parametrize("output_type", ["pil", "latent"])
|
|
@pytest.mark.parametrize("save_video", [False, True])
|
|
def test_generate_output_path_uses_effective_request_options(monkeypatch, tmp_path, output_type, save_video):
|
|
output = torch.ones((1, 16, 1, 2, 2) if output_type == "latent" else (1, 3, 2, 16, 16)) * 0.5
|
|
default_args = _single_video_args(output_type="latent" if output_type == "pil" else "pil")
|
|
generator = _single_video_generator(_single_video_output_batch(output), default_args)
|
|
request_args = _single_video_args(output_type=output_type)
|
|
request_args.model_path = "test-model"
|
|
request_args.prompt_txt = None
|
|
sampling_param = _small_sampling_param(save_video=not save_video)
|
|
sampling_param.output_path = str(tmp_path / "outputs")
|
|
monkeypatch.setattr(video_generator_module.os, "pathconf", lambda *args: 255, raising=False)
|
|
|
|
def generate():
|
|
return generator._generate_video_impl(
|
|
prompt="\u754c" * 100,
|
|
sampling_param=sampling_param,
|
|
fastvideo_args=request_args,
|
|
save_video=save_video,
|
|
return_samples=True,
|
|
)
|
|
|
|
if save_video and output_type != "latent":
|
|
|
|
def fail_execute(*args):
|
|
pytest.fail("Invalid output filenames must be rejected before model execution")
|
|
|
|
monkeypatch.setattr(generator.executor, "execute_forward", fail_execute)
|
|
with pytest.raises(ValueError, match="Output filename exceeds the 255-byte filesystem limit"):
|
|
generate()
|
|
else:
|
|
result = generate()
|
|
assert result["video_path"] is None
|
|
torch.testing.assert_close(result["samples"], output)
|
|
assert not (tmp_path / "outputs").exists()
|
|
|
|
|
|
def test_generate_single_video_metadata_only_skips_output_materialization(monkeypatch, tmp_path):
|
|
output_batch = _single_video_output_batch(
|
|
_NoCpuMaterializationOutput(),
|
|
extra={"peak_memory_mb": 42.0},
|
|
)
|
|
fastvideo_args = _single_video_args()
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
empty_calls = []
|
|
real_empty = video_generator_module.torch.empty
|
|
|
|
def record_empty(*args, **kwargs):
|
|
empty_calls.append((args, kwargs))
|
|
return real_empty(*args, **kwargs)
|
|
|
|
def fail_make_grid(*args, **kwargs):
|
|
raise AssertionError("metadata-only output path should not build frame grids")
|
|
|
|
monkeypatch.setattr(video_generator_module.torch, "empty", record_empty)
|
|
monkeypatch.setattr(video_generator_module.torchvision.utils, "make_grid", fail_make_grid)
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="metadata only",
|
|
sampling_param=_small_sampling_param(save_video=False, return_frames=False),
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "unused.mp4"),
|
|
)
|
|
|
|
assert result["samples"] is None
|
|
assert result["frames"] is None
|
|
assert result["video_path"] is None
|
|
assert result["peak_memory_mb"] == 42.0
|
|
assert empty_calls == [((0, ), {"device": "cpu"})]
|
|
|
|
|
|
def test_generate_single_video_return_frames_still_materializes_output(tmp_path):
|
|
output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5
|
|
output_batch = _single_video_output_batch(output)
|
|
fastvideo_args = _single_video_args()
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="return frames",
|
|
sampling_param=_small_sampling_param(save_video=False, return_frames=True),
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "unused.mp4"),
|
|
)
|
|
|
|
torch.testing.assert_close(result["samples"], output)
|
|
assert len(result["frames"]) == 2
|
|
assert result["video_path"] is None
|
|
|
|
|
|
def test_generate_single_video_retries_pageable_output_when_pinned_allocation_fails(monkeypatch, tmp_path):
|
|
output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5
|
|
output_batch = _single_video_output_batch(output)
|
|
fastvideo_args = _single_video_args()
|
|
fastvideo_args.pin_cpu_memory = True
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
real_empty = torch.empty
|
|
allocation_attempts = []
|
|
warnings = []
|
|
|
|
def fake_empty(size, **kwargs):
|
|
allocation_attempts.append(kwargs.get("pin_memory", False))
|
|
if kwargs.get("pin_memory"):
|
|
raise RuntimeError("CUDA error: out of memory")
|
|
return real_empty(size, **kwargs)
|
|
|
|
monkeypatch.setattr(fastvideo_utils, "is_pin_memory_available", lambda: True)
|
|
monkeypatch.setattr(fastvideo_utils.torch, "empty", fake_empty)
|
|
monkeypatch.setattr(fastvideo_utils.logger, "warning", lambda message, *args: warnings.append(message % args))
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="pinned allocation fallback",
|
|
sampling_param=_small_sampling_param(save_video=False, return_frames=True),
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "unused.mp4"),
|
|
)
|
|
|
|
torch.testing.assert_close(result["samples"], output)
|
|
assert allocation_attempts[:2] == [True, False]
|
|
assert len(warnings) == 1
|
|
assert "retrying with pageable memory" in warnings[0]
|
|
|
|
|
|
def test_generate_single_video_frames_match_legacy_cpu_loop(tmp_path):
|
|
"""The on-device quantize path (#1362) must reproduce the legacy
|
|
per-frame CPU loop (make_grid -> permute -> *255 -> uint8) bit-exactly
|
|
for in-range fp32 pixels: same uint8 dtype, same HWC grid layout with
|
|
nrow=6 (batch>1), odd frame count. CPU-only: on CUDA the float->uint8
|
|
cast may differ by <=1 LSB, but on CPU both orderings run identical
|
|
fp32 ops, so exact equality is required."""
|
|
torch.manual_seed(0)
|
|
output = torch.rand((2, 3, 3, 16, 16), dtype=torch.float32)
|
|
output_batch = _single_video_output_batch(output)
|
|
fastvideo_args = _single_video_args()
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
sampling_param = _small_sampling_param(save_video=False, return_frames=True)
|
|
sampling_param.num_frames = 3
|
|
sampling_param.num_videos_per_prompt = 2
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="grid parity",
|
|
sampling_param=sampling_param,
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "unused.mp4"),
|
|
)
|
|
|
|
legacy_frames = []
|
|
for x in rearrange(output, "b c t h w -> t b c h w"):
|
|
grid = video_generator_module.torchvision.utils.make_grid(x, nrow=6)
|
|
grid = grid.permute(1, 2, 0).squeeze(-1)
|
|
legacy_frames.append((grid * 255).to(torch.uint8).contiguous().cpu().numpy())
|
|
|
|
torch.testing.assert_close(result["samples"], output)
|
|
assert len(result["frames"]) == 3
|
|
for got, want in zip(result["frames"], legacy_frames, strict=True):
|
|
assert got.dtype == np.uint8
|
|
assert got.shape == want.shape
|
|
np.testing.assert_array_equal(got, want)
|
|
|
|
|
|
def test_generate_single_video_frames_clamp_out_of_range_pixels(tmp_path):
|
|
"""VAE output slightly outside [0, 1] must saturate at 0/255 in the
|
|
uint8 frames. The pre-#1362 unclamped cast wrapped mod 256 (e.g.
|
|
1.5 -> 126). CPU-only."""
|
|
output = torch.full((1, 3, 2, 16, 16), 1.5, dtype=torch.float32)
|
|
output[:, :, 1] = -0.5
|
|
output_batch = _single_video_output_batch(output)
|
|
fastvideo_args = _single_video_args()
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="clamp",
|
|
sampling_param=_small_sampling_param(save_video=False, return_frames=True),
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "unused.mp4"),
|
|
)
|
|
|
|
frames = result["frames"]
|
|
assert len(frames) == 2
|
|
# make_grid passes a single image through without grid padding, so
|
|
# every pixel comes from the (clamped) output tensor.
|
|
assert frames[0].dtype == np.uint8
|
|
assert frames[0].shape == (16, 16, 3)
|
|
assert (frames[0] == 255).all()
|
|
assert (frames[1] == 0).all()
|
|
|
|
|
|
def test_generate_single_video_accepts_uint8_worker_output(tmp_path):
|
|
"""A worker may quantize decoded pixels to uint8 before they cross the
|
|
executor boundary (the MiniMax-H3 decode stage does, cutting the IPC
|
|
payload 4x). Post-decode must pass those frames through untouched and
|
|
``samples`` must stay on its [0, 1] float contract. CPU-only."""
|
|
torch.manual_seed(0)
|
|
pixels = torch.rand((1, 3, 2, 16, 16), dtype=torch.float32)
|
|
output_u8 = (pixels * 255).clamp_(0, 255).to(torch.uint8)
|
|
output_batch = _single_video_output_batch(output_u8)
|
|
fastvideo_args = _single_video_args()
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="uint8 worker output",
|
|
sampling_param=_small_sampling_param(save_video=False, return_frames=True),
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "unused.mp4"),
|
|
)
|
|
|
|
torch.testing.assert_close(result["samples"], output_u8.float() / 255)
|
|
frames = result["frames"]
|
|
assert len(frames) == 2
|
|
for index, frame in enumerate(frames):
|
|
assert frame.dtype == np.uint8
|
|
np.testing.assert_array_equal(frame, output_u8[0, :, index].permute(1, 2, 0).numpy())
|
|
|
|
|
|
def test_generate_single_video_save_video_still_builds_frames(monkeypatch, tmp_path):
|
|
output = torch.ones((1, 3, 2, 16, 16), dtype=torch.float32) * 0.5
|
|
output_batch = _single_video_output_batch(output)
|
|
fastvideo_args = _single_video_args()
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
saved = {}
|
|
|
|
def fake_mimsave(path, frames, *, fps, format):
|
|
saved["path"] = path
|
|
saved["frame_count"] = len(frames)
|
|
saved["fps"] = fps
|
|
saved["format"] = format
|
|
|
|
monkeypatch.setattr(video_generator_module.imageio, "mimsave", fake_mimsave)
|
|
|
|
output_path = str(tmp_path / "saved.mp4")
|
|
result = generator._generate_single_video(
|
|
prompt="save only",
|
|
sampling_param=_small_sampling_param(save_video=True, return_frames=False),
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=output_path,
|
|
)
|
|
|
|
assert result["samples"] is None
|
|
assert result["frames"] is None
|
|
assert result["video_path"] == output_path
|
|
assert saved == {
|
|
"path": output_path,
|
|
"frame_count": 2,
|
|
"fps": 8,
|
|
"format": "mp4",
|
|
}
|
|
|
|
|
|
def test_generate_single_video_save_only_reports_refined_output_size(monkeypatch, tmp_path):
|
|
"""`GenerationResult.size` must describe the decoded media even when the
|
|
fp32 `samples` mirror is skipped (`return_frames=False`, the CLI save
|
|
flow). Refiner pipelines can change the final pixel geometry, so the size
|
|
has to come from `output_batch.output`, not the base request. CPU-only."""
|
|
# Refiner-style output: request asks for 2 frames of 16x16, pipeline
|
|
# produces 5 frames of 32x48.
|
|
output = torch.full((1, 3, 5, 32, 48), 0.5, dtype=torch.float32)
|
|
output_batch = _single_video_output_batch(output)
|
|
fastvideo_args = _single_video_args()
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
saved = {}
|
|
|
|
def fake_mimsave(path, frames, *, fps, format):
|
|
saved["frame_count"] = len(frames)
|
|
|
|
monkeypatch.setattr(video_generator_module.imageio, "mimsave", fake_mimsave)
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="refined save",
|
|
sampling_param=_small_sampling_param(save_video=True, return_frames=False),
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "refined.mp4"),
|
|
)
|
|
|
|
assert result["samples"] is None
|
|
assert result["frames"] is None
|
|
assert result["size"] == (32, 48, 5)
|
|
assert saved["frame_count"] == 5
|
|
|
|
|
|
def test_generate_single_video_audio_only_metadata_returns_audio_without_frames(tmp_path):
|
|
audio = torch.zeros((16, ), dtype=torch.float32)
|
|
output_batch = _single_video_output_batch(
|
|
_NoCpuMaterializationOutput(),
|
|
extra={
|
|
"audio_only": True,
|
|
"audio": audio,
|
|
"audio_sample_rate": 44100,
|
|
},
|
|
)
|
|
fastvideo_args = _single_video_args()
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="audio only",
|
|
sampling_param=_small_sampling_param(save_video=False, return_frames=False),
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "unused.mp4"),
|
|
)
|
|
|
|
assert result["samples"] is None
|
|
assert result["frames"] is None
|
|
assert result["audio"] is audio
|
|
assert result["audio_sample_rate"] == 44100
|
|
|
|
|
|
def test_generate_single_video_audio_only_save_skips_placeholder_materialization(tmp_path):
|
|
audio = torch.zeros((16, ), dtype=torch.float32)
|
|
output_batch = _single_video_output_batch(
|
|
_NoCpuMaterializationOutput(),
|
|
extra={
|
|
"audio_only": True,
|
|
"audio": audio,
|
|
"audio_sample_rate": 44100,
|
|
},
|
|
)
|
|
fastvideo_args = _single_video_args()
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
output_path = str(tmp_path / "audio.mp4")
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="audio only",
|
|
sampling_param=_small_sampling_param(save_video=True, return_frames=False),
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=output_path,
|
|
)
|
|
|
|
assert result["samples"] is None
|
|
assert result["frames"] is None
|
|
assert result["audio"] is audio
|
|
assert result["audio_sample_rate"] == 44100
|
|
assert result["video_path"] == str(tmp_path / "audio.wav")
|
|
|
|
|
|
def test_generate_single_video_ray_audio_only_save_preserves_worker_metadata(monkeypatch, tmp_path):
|
|
audio = torch.zeros((16, ), dtype=torch.float32)
|
|
worker_output = ForwardBatch(
|
|
data_type="audio",
|
|
output=torch.ones((1, 3, 1, 8, 8)),
|
|
extra={
|
|
"audio_only": True,
|
|
"audio": audio,
|
|
"audio_sample_rate": 44100,
|
|
},
|
|
)
|
|
worker = Worker.__new__(Worker)
|
|
worker.fastvideo_args = SimpleNamespace(is_output_rank=True)
|
|
worker.pipeline = SimpleNamespace(forward=lambda batch, args: worker_output)
|
|
|
|
monkeypatch.setattr(RayDistributedExecutor, "__abstractmethods__", frozenset())
|
|
executor = RayDistributedExecutor.__new__(RayDistributedExecutor)
|
|
executor.shutdown = lambda: None
|
|
|
|
def collective_rpc(method, *, kwargs, **_):
|
|
assert method == "execute_forward"
|
|
return [worker.execute_forward(**kwargs)]
|
|
|
|
executor.collective_rpc = collective_rpc
|
|
fastvideo_args = _single_video_args()
|
|
generator = _new_video_generator()
|
|
generator.fastvideo_args = fastvideo_args
|
|
generator.executor = executor
|
|
generator.config = None
|
|
written = {}
|
|
|
|
def fake_write_pcm_wav(path, samples, sample_rate):
|
|
written.update(path=path, samples=samples, sample_rate=sample_rate)
|
|
|
|
monkeypatch.setattr(generator, "_write_pcm_wav", fake_write_pcm_wav)
|
|
sampling_param = _small_sampling_param(save_video=True, return_frames=False)
|
|
sampling_param.data_type = "audio"
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="audio only",
|
|
sampling_param=sampling_param,
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "audio.mp4"),
|
|
)
|
|
|
|
assert worker_output.output is not None
|
|
assert worker_output.output.numel() == 0
|
|
assert result["samples"] is None
|
|
assert result["frames"] is None
|
|
assert result["audio"] is audio
|
|
assert result["audio_sample_rate"] == 44100
|
|
assert result["video_path"] == str(tmp_path / "audio.wav")
|
|
assert written["path"] == str(tmp_path / "audio.wav")
|
|
assert written["samples"] is audio
|
|
assert written["sample_rate"] == 44100
|
|
|
|
|
|
def test_generate_single_video_latent_metadata_skips_cpu_materialization(tmp_path):
|
|
output_batch = _single_video_output_batch(_NoCpuMaterializationOutput())
|
|
fastvideo_args = _single_video_args(output_type="latent")
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="latent metadata",
|
|
sampling_param=_small_sampling_param(save_video=False, return_frames=False),
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "unused.mp4"),
|
|
)
|
|
|
|
assert result["samples"] is None
|
|
assert result["frames"] is None
|
|
assert result["video_path"] is None
|
|
|
|
|
|
def test_generate_single_video_non_output_rank_suppresses_all_user_payloads(tmp_path):
|
|
audio = torch.ones(8)
|
|
audio_latents = torch.ones(2)
|
|
output_batch = _single_video_output_batch(
|
|
torch.empty(0),
|
|
extra={
|
|
"audio": audio,
|
|
"audio_sample_rate": 24000,
|
|
"ltx2_audio_latents": audio_latents,
|
|
},
|
|
)
|
|
output_batch.trajectory_latents = torch.ones(1)
|
|
output_batch.trajectory_timesteps = [torch.ones(1)]
|
|
output_batch.trajectory_decoded = [torch.ones(1)]
|
|
fastvideo_args = _single_video_args(output_type="latent")
|
|
generator = _single_video_generator(output_batch, fastvideo_args)
|
|
captured = {}
|
|
|
|
def execute_forward(batch, args):
|
|
captured["batch"] = batch
|
|
return output_batch
|
|
|
|
generator.executor = SimpleNamespace(
|
|
execute_forward=execute_forward,
|
|
is_output_rank=False,
|
|
uses_spmd_execution=True,
|
|
broadcast_from_output_rank=lambda value: value,
|
|
)
|
|
sampling = _small_sampling_param(save_video=True, return_frames=True)
|
|
sampling.return_samples = True
|
|
sampling.return_trajectory_latents = True
|
|
sampling.return_trajectory_decoded = True
|
|
|
|
result = generator._generate_single_video(
|
|
prompt="non-output rank",
|
|
sampling_param=sampling,
|
|
fastvideo_args=fastvideo_args,
|
|
output_path=str(tmp_path / "must-not-exist.mp4"),
|
|
)
|
|
|
|
batch = captured["batch"]
|
|
assert not batch.save_video
|
|
assert not batch.return_frames
|
|
assert not batch.return_samples
|
|
assert not batch.return_trajectory_latents
|
|
assert not batch.return_trajectory_decoded
|
|
assert result["samples"] is None
|
|
assert result["frames"] is None
|
|
assert result["audio"] is None
|
|
assert result["audio_sample_rate"] is None
|
|
assert result["ltx2_audio_latents"] is None
|
|
assert result["trajectory"] is None
|
|
assert result["trajectory_timesteps"] is None
|
|
assert result["trajectory_decoded"] is None
|
|
assert result["video_path"] is None
|
|
|
|
|
|
def test_prompt_batch_failure_is_not_swallowed_in_spmd_mode(tmp_path, monkeypatch):
|
|
generator = _new_runtime_video_generator()
|
|
generator.executor = SimpleNamespace(
|
|
set_log_queue=lambda queue: None,
|
|
clear_log_queue=lambda: None,
|
|
is_output_rank=True,
|
|
uses_spmd_execution=True,
|
|
broadcast_from_output_rank=lambda value: value,
|
|
)
|
|
_patch_sampling_param_from_pretrained(monkeypatch)
|
|
prompt_file = tmp_path / "prompts.txt"
|
|
prompt_file.write_text("first prompt\nsecond prompt\n", encoding="utf-8")
|
|
calls = []
|
|
|
|
def fail_first_prompt(prompt, **kwargs):
|
|
calls.append(prompt)
|
|
raise RuntimeError("injected forward failure")
|
|
|
|
monkeypatch.setattr(generator, "_generate_single_video", fail_first_prompt)
|
|
|
|
with pytest.raises(RuntimeError, match="injected forward failure"):
|
|
generator.generate({
|
|
"inputs": {
|
|
"prompt_path": str(prompt_file)
|
|
},
|
|
"output": {
|
|
"output_path": str(tmp_path / "outputs"),
|
|
"save_video": False,
|
|
"return_frames": False,
|
|
},
|
|
})
|
|
assert calls == ["first prompt"]
|
|
|
|
|
|
def test_from_config_normalizes_and_translates(monkeypatch):
|
|
captured = _patch_from_fastvideo_args(monkeypatch)
|
|
_patch_fastvideo_args_from_kwargs(monkeypatch)
|
|
config = GeneratorConfig(model_path="test-model")
|
|
config.engine.num_gpus = 2
|
|
config.pipeline.workload_type = "t2v"
|
|
|
|
generator = VideoGenerator.from_config(config)
|
|
|
|
assert captured["fastvideo_args"].model_path == "test-model"
|
|
assert captured["fastvideo_args"].num_gpus == 2
|
|
assert captured["fastvideo_args"].workload_type.value == "t2v"
|
|
assert generator.config == config
|
|
|
|
|
|
def test_from_config_translates_external_launcher_backend(monkeypatch):
|
|
_patch_from_fastvideo_args(monkeypatch)
|
|
captured = _patch_fastvideo_args_from_kwargs(monkeypatch)
|
|
config = GeneratorConfig(model_path="test-model")
|
|
config.engine.execution_backend = "external_launcher"
|
|
|
|
VideoGenerator.from_config(config)
|
|
|
|
assert captured["kwargs"]["distributed_executor_backend"] == "external_launcher"
|
|
|
|
|
|
def test_from_file_loads_generator_from_run_config(tmp_path, monkeypatch):
|
|
captured = _patch_from_fastvideo_args(monkeypatch)
|
|
_patch_fastvideo_args_from_kwargs(monkeypatch)
|
|
config_path = tmp_path / "run.yaml"
|
|
config_path.write_text(
|
|
"generator:\n"
|
|
" model_path: test-model\n"
|
|
" engine:\n"
|
|
" num_gpus: 3\n"
|
|
"request:\n"
|
|
" prompt: hello\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
VideoGenerator.from_file(str(config_path))
|
|
|
|
assert captured["fastvideo_args"].model_path == "test-model"
|
|
assert captured["fastvideo_args"].num_gpus == 3
|
|
|
|
|
|
def test_from_pretrained_convenience_kwargs_do_not_warn(monkeypatch):
|
|
captured = _patch_from_fastvideo_args(monkeypatch)
|
|
fastvideo_args_capture = _patch_fastvideo_args_from_kwargs(monkeypatch)
|
|
|
|
with warnings.catch_warnings(record=True) as caught:
|
|
warnings.simplefilter("always")
|
|
generator = VideoGenerator.from_pretrained(
|
|
"test-model",
|
|
num_gpus=4,
|
|
use_fsdp_inference=False,
|
|
text_encoder_cpu_offload=True,
|
|
pin_cpu_memory=True,
|
|
dit_cpu_offload=False,
|
|
vae_cpu_offload=False,
|
|
)
|
|
|
|
assert not caught
|
|
assert captured["fastvideo_args"].model_path == "test-model"
|
|
assert captured["fastvideo_args"].num_gpus == 4
|
|
assert fastvideo_args_capture["kwargs"]["use_fsdp_inference"] is False
|
|
assert fastvideo_args_capture["kwargs"]["text_encoder_cpu_offload"] is True
|
|
assert fastvideo_args_capture["kwargs"]["pin_cpu_memory"] is True
|
|
assert fastvideo_args_capture["kwargs"]["dit_cpu_offload"] is False
|
|
assert fastvideo_args_capture["kwargs"]["vae_cpu_offload"] is False
|
|
assert generator.config is not None
|
|
assert generator.config.model_path == "test-model"
|
|
assert generator.config.engine.num_gpus == 4
|
|
|
|
|
|
def test_from_pretrained_accepts_external_launcher_backend(monkeypatch):
|
|
_patch_from_fastvideo_args(monkeypatch)
|
|
captured = _patch_fastvideo_args_from_kwargs(monkeypatch)
|
|
|
|
VideoGenerator.from_pretrained(
|
|
"test-model",
|
|
num_gpus=2,
|
|
sp_size=2,
|
|
distributed_executor_backend="external_launcher",
|
|
)
|
|
|
|
assert captured["kwargs"]["distributed_executor_backend"] == "external_launcher"
|
|
|
|
|
|
def test_from_pretrained_legacy_only_kwargs_warn(monkeypatch):
|
|
captured = _patch_from_fastvideo_args(monkeypatch)
|
|
_patch_fastvideo_args_from_kwargs(monkeypatch)
|
|
|
|
with pytest.warns(DeprecationWarning, match="legacy-only kwargs"):
|
|
generator = VideoGenerator.from_pretrained(
|
|
"test-model",
|
|
num_gpus=4,
|
|
workload_type="t2v",
|
|
)
|
|
|
|
assert captured["fastvideo_args"].model_path == "test-model"
|
|
assert captured["fastvideo_args"].num_gpus == 4
|
|
assert captured["fastvideo_args"].workload_type.value == "t2v"
|
|
assert generator.config is not None
|
|
assert generator.config.pipeline.workload_type == "t2v"
|
|
|
|
|
|
def test_generate_uses_typed_request_path(monkeypatch):
|
|
generator = _new_runtime_video_generator()
|
|
_patch_sampling_param_from_pretrained(monkeypatch)
|
|
captured = {}
|
|
|
|
def fake_generate_video_impl(prompt=None, sampling_param=None, **kwargs):
|
|
captured["prompt"] = prompt
|
|
captured["sampling_param"] = sampling_param
|
|
captured["kwargs"] = kwargs
|
|
return {"prompts": prompt, "video_path": "outputs/test.mp4"}
|
|
|
|
monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl)
|
|
|
|
result = generator.generate(
|
|
GenerationRequest(
|
|
prompt="hello world",
|
|
sampling=SamplingConfig(num_frames=81, height=480, width=832),
|
|
))
|
|
|
|
assert isinstance(result, GenerationResult)
|
|
assert captured["prompt"] == "hello world"
|
|
assert captured["sampling_param"].num_frames == 81
|
|
assert captured["sampling_param"].height == 480
|
|
assert captured["sampling_param"].width == 832
|
|
assert result.video_path == "outputs/test.mp4"
|
|
|
|
|
|
@pytest.mark.parametrize("entry", ["generate", "generate_video"])
|
|
def test_generate_applies_pipeline_config_request_constraints(monkeypatch, entry):
|
|
"""Both entry points run the sampling params that the loaded model's request constraints return."""
|
|
generator = _new_runtime_video_generator()
|
|
_patch_sampling_param_from_pretrained(monkeypatch)
|
|
|
|
def apply_request_constraints(request, sampling_param):
|
|
return dataclasses.replace(sampling_param, num_inference_steps=8)
|
|
|
|
generator.fastvideo_args.pipeline_config = SimpleNamespace(apply_request_constraints=apply_request_constraints)
|
|
captured = {}
|
|
|
|
def fake_generate_video_impl(prompt=None, sampling_param=None, **kwargs):
|
|
captured["sampling_param"] = sampling_param
|
|
return {"prompts": prompt, "video_path": "outputs/test.mp4"}
|
|
|
|
monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl)
|
|
if entry == "generate":
|
|
generator.generate(GenerationRequest(prompt="hello world"))
|
|
else:
|
|
with pytest.warns(DeprecationWarning):
|
|
generator.generate_video(prompt="hello world")
|
|
assert captured["sampling_param"].num_inference_steps == 8
|
|
|
|
|
|
def test_generate_rejects_stage_override_outside_registered_stage(monkeypatch) -> None:
|
|
"""Validate typed stage overrides at the public generation entrypoint."""
|
|
generator = _new_runtime_video_generator()
|
|
monkeypatch.setattr(
|
|
"fastvideo.registry.get_preset_selection",
|
|
lambda _model_path: ("lingbot_video_moe_refiner_t2v", "lingbot_video"),
|
|
)
|
|
with pytest.raises(ConfigValidationError, match="stage_overrides.denoise.t_thresh"):
|
|
generator.generate({
|
|
"prompt": "hello world",
|
|
"stage_overrides": {
|
|
"denoise": {
|
|
"t_thresh": 0.7
|
|
}
|
|
},
|
|
})
|
|
|
|
|
|
def test_generate_preserves_schema_defaults_for_dataclass_request(monkeypatch):
|
|
generator = _new_runtime_video_generator()
|
|
captured = {}
|
|
|
|
def fake_from_pretrained(cls, model_path):
|
|
return cls(
|
|
negative_prompt="model default",
|
|
num_frames=61,
|
|
height=448,
|
|
width=832,
|
|
)
|
|
|
|
def fake_generate_video_impl(prompt=None, sampling_param=None, **kwargs):
|
|
captured["sampling_param"] = sampling_param
|
|
return {"prompts": prompt, "video_path": "outputs/test.mp4"}
|
|
|
|
monkeypatch.setattr(SamplingParam, "from_pretrained", classmethod(fake_from_pretrained))
|
|
monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl)
|
|
|
|
generator.generate(
|
|
GenerationRequest(
|
|
prompt="hello world",
|
|
negative_prompt=None,
|
|
sampling=SamplingConfig(num_frames=125, height=720, width=1280),
|
|
))
|
|
|
|
assert captured["sampling_param"].negative_prompt is None
|
|
assert captured["sampling_param"].num_frames == 125
|
|
assert captured["sampling_param"].height == 720
|
|
assert captured["sampling_param"].width == 1280
|
|
|
|
|
|
def test_generate_mapping_request_preserves_model_defaults_for_omitted_fields(monkeypatch, ):
|
|
generator = _new_runtime_video_generator()
|
|
captured = {}
|
|
|
|
def fake_from_pretrained(cls, model_path):
|
|
return cls(
|
|
negative_prompt="model default",
|
|
num_frames=61,
|
|
height=448,
|
|
width=832,
|
|
fps=16,
|
|
guidance_scale=3.0,
|
|
)
|
|
|
|
def fake_generate_video_impl(prompt=None, sampling_param=None, **kwargs):
|
|
captured["sampling_param"] = sampling_param
|
|
return {"prompts": prompt, "video_path": "outputs/test.mp4"}
|
|
|
|
monkeypatch.setattr(SamplingParam, "from_pretrained", classmethod(fake_from_pretrained))
|
|
monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl)
|
|
|
|
generator.generate({
|
|
"prompt": "hello world",
|
|
})
|
|
|
|
assert captured["sampling_param"].negative_prompt == "model default"
|
|
assert captured["sampling_param"].num_frames == 61
|
|
assert captured["sampling_param"].height == 448
|
|
assert captured["sampling_param"].width == 832
|
|
assert captured["sampling_param"].fps == 16
|
|
assert captured["sampling_param"].guidance_scale == 3.0
|
|
|
|
|
|
def test_generate_honors_post_load_request_mutations(monkeypatch, tmp_path):
|
|
generator = _new_runtime_video_generator()
|
|
captured = {}
|
|
config_path = tmp_path / "run.yaml"
|
|
config_path.write_text(
|
|
"generator:\n"
|
|
" model_path: test-model\n"
|
|
"request:\n"
|
|
" prompt: hello world\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
def fake_from_pretrained(cls, model_path):
|
|
return cls(seed=1024, num_frames=61, height=448, width=832)
|
|
|
|
def fake_generate_video_impl(prompt=None, sampling_param=None, **kwargs):
|
|
captured["sampling_param"] = sampling_param
|
|
return {"prompts": prompt, "video_path": "outputs/test.mp4"}
|
|
|
|
monkeypatch.setattr(SamplingParam, "from_pretrained", classmethod(fake_from_pretrained))
|
|
monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl)
|
|
|
|
config = load_run_config(config_path)
|
|
config.request.sampling.seed = 7
|
|
|
|
generator.generate(config.request)
|
|
|
|
assert captured["sampling_param"].seed == 7
|
|
|
|
|
|
def test_generate_honors_post_load_mutations_matching_schema_defaults(
|
|
monkeypatch,
|
|
tmp_path,
|
|
):
|
|
generator = _new_runtime_video_generator()
|
|
captured = {}
|
|
config_path = tmp_path / "run.yaml"
|
|
config_path.write_text(
|
|
"generator:\n"
|
|
" model_path: test-model\n"
|
|
"request:\n"
|
|
" prompt: hello world\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
def fake_from_pretrained(cls, model_path):
|
|
return cls(guidance_scale=3.0)
|
|
|
|
def fake_generate_video_impl(prompt=None, sampling_param=None, **kwargs):
|
|
captured["sampling_param"] = sampling_param
|
|
return {"prompts": prompt, "video_path": "outputs/test.mp4"}
|
|
|
|
monkeypatch.setattr(SamplingParam, "from_pretrained", classmethod(fake_from_pretrained))
|
|
monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl)
|
|
|
|
config = load_run_config(config_path)
|
|
config.request.sampling.guidance_scale = 1.0
|
|
|
|
generator.generate(config.request)
|
|
|
|
assert captured["sampling_param"].guidance_scale == 1.0
|
|
|
|
|
|
def test_generate_removes_deleted_loaded_stage_overrides(monkeypatch, tmp_path):
|
|
generator = _new_runtime_video_generator()
|
|
captured = {}
|
|
config_path = tmp_path / "run.yaml"
|
|
config_path.write_text(
|
|
"generator:\n"
|
|
" model_path: test-model\n"
|
|
"request:\n"
|
|
" prompt: hello world\n"
|
|
" stage_overrides:\n"
|
|
" refine:\n"
|
|
" t_thresh: 0.8\n",
|
|
encoding="utf-8",
|
|
)
|
|
|
|
def fake_from_pretrained(cls, model_path):
|
|
return cls(t_thresh=0.5)
|
|
|
|
def fake_generate_video_impl(prompt=None, sampling_param=None, **kwargs):
|
|
captured["sampling_param"] = sampling_param
|
|
return {"prompts": prompt, "video_path": "outputs/test.mp4"}
|
|
|
|
monkeypatch.setattr(SamplingParam, "from_pretrained", classmethod(fake_from_pretrained))
|
|
monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl)
|
|
|
|
config = load_run_config(config_path)
|
|
del config.request.stage_overrides["refine"]
|
|
|
|
generator.generate(config.request)
|
|
|
|
assert captured["sampling_param"].t_thresh == 0.5
|
|
|
|
|
|
def test_generate_video_legacy_call_uses_legacy_impl(monkeypatch):
|
|
generator = _new_runtime_video_generator()
|
|
captured = {}
|
|
|
|
def fake_generate_video_impl(
|
|
prompt=None,
|
|
sampling_param=None,
|
|
mouse_cond=None,
|
|
keyboard_cond=None,
|
|
grid_sizes=None,
|
|
**kwargs,
|
|
):
|
|
captured["prompt"] = prompt
|
|
captured["sampling_param"] = sampling_param
|
|
captured["mouse_cond"] = mouse_cond
|
|
captured["keyboard_cond"] = keyboard_cond
|
|
captured["grid_sizes"] = grid_sizes
|
|
captured["kwargs"] = kwargs
|
|
return {"prompts": prompt, "video_path": "outputs/test.mp4"}
|
|
|
|
monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl)
|
|
|
|
with pytest.warns(DeprecationWarning):
|
|
result = generator.generate_video(
|
|
prompt="legacy prompt",
|
|
num_frames=49,
|
|
output_path="outputs/legacy",
|
|
save_video=False,
|
|
log_queue="queue-token",
|
|
)
|
|
|
|
assert captured["prompt"] == "legacy prompt"
|
|
assert captured["sampling_param"].num_frames == 49
|
|
assert captured["sampling_param"].output_path == "outputs/legacy"
|
|
assert captured["sampling_param"].save_video is False
|
|
assert result["video_path"] == "outputs/test.mp4"
|
|
|
|
|
|
def test_generate_video_legacy_call_routes_compat_kwargs(monkeypatch):
|
|
generator = _new_runtime_video_generator()
|
|
generator.fastvideo_args.pipeline_config = SimpleNamespace(
|
|
embedded_cfg_scale=1.0, apply_request_constraints=lambda request, sampling_param: sampling_param)
|
|
captured = {}
|
|
|
|
def fake_generate_video_impl(
|
|
prompt=None,
|
|
sampling_param=None,
|
|
mouse_cond=None,
|
|
keyboard_cond=None,
|
|
grid_sizes=None,
|
|
**kwargs,
|
|
):
|
|
captured["prompt"] = prompt
|
|
captured["sampling_param"] = sampling_param
|
|
captured["kwargs"] = kwargs
|
|
return {"prompts": prompt, "video_path": "outputs/test.mp4"}
|
|
|
|
monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl)
|
|
|
|
with pytest.warns(DeprecationWarning):
|
|
result = generator.generate_video(
|
|
prompt="legacy prompt",
|
|
neg_prompt="custom negative",
|
|
embedded_cfg_scale=7.5,
|
|
)
|
|
|
|
assert captured["prompt"] == "legacy prompt"
|
|
assert captured["sampling_param"].negative_prompt == "custom negative"
|
|
assert captured["kwargs"]["fastvideo_args"].pipeline_config.embedded_cfg_scale == 7.5
|
|
assert not hasattr(captured["sampling_param"], "embedded_cfg_scale")
|
|
assert result["video_path"] == "outputs/test.mp4"
|
|
|
|
|
|
def test_generate_batch_prompt_file_returns_typed_results(tmp_path, monkeypatch):
|
|
generator = _new_runtime_video_generator()
|
|
_patch_sampling_param_from_pretrained(monkeypatch)
|
|
prompt_file = tmp_path / "prompts.txt"
|
|
prompt_file.write_text("first prompt\nsecond prompt\n", encoding="utf-8")
|
|
output_dir = tmp_path / "outputs"
|
|
captured_prompts = []
|
|
|
|
def fake_generate_single_video(prompt, sampling_param=None, **kwargs):
|
|
captured_prompts.append(prompt)
|
|
return {"prompts": prompt, "video_path": kwargs["output_path"]}
|
|
|
|
monkeypatch.setattr(generator, "_generate_single_video", fake_generate_single_video)
|
|
|
|
results = generator.generate({
|
|
"inputs": {
|
|
"prompt_path": str(prompt_file)
|
|
},
|
|
"output": {
|
|
"output_path": str(output_dir),
|
|
"save_video": False,
|
|
"return_frames": False,
|
|
},
|
|
})
|
|
|
|
assert isinstance(results, list)
|
|
assert [result.prompt for result in results] == ["first prompt", "second prompt"]
|
|
assert [result.prompt_index for result in results] == [0, 1]
|
|
assert captured_prompts == ["first prompt", "second prompt"]
|
|
|
|
|
|
def test_generate_batched_request_fans_out_media_inputs(monkeypatch):
|
|
generator = _new_runtime_video_generator()
|
|
_patch_sampling_param_from_pretrained(monkeypatch)
|
|
captured: list[tuple[str | None, str | None, str | None]] = []
|
|
|
|
def fake_generate_video_impl(prompt=None, sampling_param=None, **kwargs):
|
|
captured.append((prompt, sampling_param.image_path, sampling_param.video_path))
|
|
return {"prompts": prompt, "video_path": "outputs/test.mp4"}
|
|
|
|
monkeypatch.setattr(generator, "_generate_video_impl", fake_generate_video_impl)
|
|
|
|
results = generator.generate(
|
|
GenerationRequest(
|
|
prompt=["first prompt", "second prompt"],
|
|
inputs=InputConfig(
|
|
image_path=["first.png", "second.png"],
|
|
video_path=["first.mp4", "second.mp4"],
|
|
),
|
|
))
|
|
|
|
assert [result.prompt for result in results] == ["first prompt", "second prompt"]
|
|
assert captured == [
|
|
("first prompt", "first.png", "first.mp4"),
|
|
("second prompt", "second.png", "second.mp4"),
|
|
]
|
|
|
|
|
|
def test_generate_batched_request_rejects_mismatched_media_inputs(monkeypatch):
|
|
generator = _new_runtime_video_generator()
|
|
_patch_sampling_param_from_pretrained(monkeypatch)
|
|
|
|
with pytest.raises(ValueError, match="image_path"):
|
|
generator.generate(
|
|
GenerationRequest(
|
|
prompt=["first prompt", "second prompt"],
|
|
inputs=InputConfig(image_path=["first.png"]),
|
|
))
|
|
|
|
|
|
@pytest.mark.parametrize("audio_seconds", [1, 2, 3])
|
|
def test_ffmpeg_pipe_preserves_video_duration(tmp_path, env_overrides, audio_seconds):
|
|
ffmpeg = shutil.which("ffmpeg")
|
|
ffprobe = shutil.which("ffprobe")
|
|
if ffmpeg is None or ffprobe is None:
|
|
pytest.skip("ffmpeg and ffprobe are required")
|
|
env_overrides.enter_context(envs.FASTVIDEO_FFMPEG_BIN.override(ffmpeg))
|
|
env_overrides.enter_context(envs.FASTVIDEO_VIDEO_CODEC.override("libx264"))
|
|
env_overrides.enter_context(envs.FASTVIDEO_OUTPUT_PIX_FMT.override("yuv420p"))
|
|
fps, frame_count, sample_rate = 24, 48, 32000
|
|
frames = [np.full((64, 64, 3), index * 4, dtype=np.uint8) for index in range(frame_count)]
|
|
audio = np.zeros((2, audio_seconds * sample_rate), dtype=np.float32)
|
|
output = tmp_path / "video.mp4"
|
|
|
|
assert VideoGenerator._save_video_with_audio_ffmpeg_pipe(str(output), frames, fps, audio, sample_rate)
|
|
probe = json.loads(subprocess.check_output([
|
|
ffprobe, "-v", "error", "-count_frames", "-show_entries",
|
|
"stream=codec_type,nb_read_frames,duration", "-of", "json", str(output),
|
|
], text=True))
|
|
video = next(stream for stream in probe["streams"] if stream["codec_type"] == "video")
|
|
sound = next(stream for stream in probe["streams"] if stream["codec_type"] == "audio")
|
|
assert int(video["nb_read_frames"]) == frame_count
|
|
assert float(video["duration"]) == pytest.approx(frame_count / fps, abs=1 / fps)
|
|
assert float(sound["duration"]) == pytest.approx(frame_count / fps, abs=1 / fps)
|
|
|
|
|
|
@pytest.mark.parametrize("dtype", [torch.float32, torch.uint8])
|
|
def test_return_samples_only_skips_preview_processing(monkeypatch, tmp_path, dtype):
|
|
output = torch.ones((2, 3, 1, 16, 16), dtype=dtype)
|
|
generator = _single_video_generator(_single_video_output_batch(output), _single_video_args())
|
|
monkeypatch.setattr("fastvideo.entrypoints.video_generator.pixels_to_uint8",
|
|
lambda *args: pytest.fail("Samples-only requests must not quantize a preview"))
|
|
monkeypatch.setattr("fastvideo.entrypoints.video_generator.torchvision.utils.make_grid",
|
|
lambda *args, **kwargs: pytest.fail("Samples-only requests must not build a preview grid"))
|
|
sampling = _small_sampling_param()
|
|
sampling.num_frames = 1
|
|
sampling.num_videos_per_prompt = 2
|
|
sampling.return_samples = True
|
|
result = generator._generate_single_video("samples only", sampling,
|
|
output_path=str(tmp_path / "unused.mp4"))
|
|
torch.testing.assert_close(result["samples"], output.float() / (255 if dtype == torch.uint8 else 1))
|
|
assert result["frames"] is None
|
|
assert result["video_path"] is None
|
|
assert result["size"] == (16, 16, 1)
|
|
|
|
|
|
@pytest.mark.parametrize("return_frames", [False, True])
|
|
def test_input_sized_output_skips_request_pixel_preallocation(monkeypatch, tmp_path, return_frames):
|
|
output = torch.full((1, 3, 1, 24, 32), 0.5)
|
|
args = _single_video_args()
|
|
args.pipeline_config.output_shape_from_input = True
|
|
generator = _single_video_generator(_single_video_output_batch(output), args)
|
|
monkeypatch.setattr(
|
|
video_generator_module, "allocate_cpu_tensor_with_pin_fallback",
|
|
lambda *args, **kwargs: pytest.fail("Input-sized outputs must not allocate request-sized pixel buffers"))
|
|
sampling = _small_sampling_param(return_frames=return_frames)
|
|
sampling.num_frames = 1
|
|
sampling.return_samples = True
|
|
result = generator._generate_single_video("input-sized output", sampling,
|
|
output_path=str(tmp_path / "unused.mp4"))
|
|
torch.testing.assert_close(result["samples"], output)
|
|
assert result["size"] == (24, 32, 1)
|
|
if return_frames:
|
|
assert len(result["frames"]) == 1
|
|
else:
|
|
assert result["frames"] is None
|
|
assert result["video_path"] is None
|
|
|
|
|
|
@pytest.mark.parametrize("extension,format_name", [(".png", "PNG"), (".jpg", "JPEG"), (".jpeg", "JPEG"), (".webp", "WEBP")])
|
|
def test_image_output_path_preserves_requested_format(tmp_path, extension, format_name):
|
|
from PIL import Image
|
|
|
|
args = _single_video_args()
|
|
args.workload_type = WorkloadType.from_string("t2i")
|
|
generator = _single_video_generator(_single_video_output_batch(torch.zeros(1, 3, 1, 16, 16)), args)
|
|
target = tmp_path / ("image" + extension)
|
|
output_path = generator._prepare_output_path(str(target), "a cat")
|
|
assert output_path == str(target)
|
|
sampling = _small_sampling_param(save_video=True)
|
|
sampling.num_frames = 1
|
|
result = generator._generate_single_video("a cat", sampling, output_path=output_path)
|
|
assert result["video_path"] == str(target)
|
|
with Image.open(target) as image:
|
|
assert image.format == format_name
|
|
assert image.size == (16, 16)
|
|
|
|
|
|
def test_image_output_directory_defaults_to_png(tmp_path):
|
|
generator = _new_video_generator()
|
|
generator.fastvideo_args = SimpleNamespace(workload_type=WorkloadType.from_string("t2i"))
|
|
assert generator._prepare_output_path(str(tmp_path), "a cat") == str(tmp_path / "a cat.png")
|
|
|
|
|
|
@pytest.mark.parametrize("extension,target_extension,workload", [
|
|
(".png", ".png", "t2i"), (".jpg", ".png", "t2i"), (".jpeg", ".png", "t2i"),
|
|
(".webp", ".png", "t2i"), (".JPEG", ".png", "t2i"),
|
|
(".mp4", ".mp4", "t2v"), (".wav", ".wav", "t2a"),
|
|
])
|
|
def test_existing_image_extension_directory_receives_output(tmp_path, extension, target_extension, workload):
|
|
generator = _new_video_generator()
|
|
generator.fastvideo_args = SimpleNamespace(workload_type=WorkloadType.from_string("t2i"))
|
|
generator.fastvideo_args.workload_type = WorkloadType.from_string(workload)
|
|
directory = tmp_path / ("renders" + extension)
|
|
directory.mkdir()
|
|
existing = directory / ("a cat" + target_extension)
|
|
existing.write_bytes(b"existing image")
|
|
result = generator._prepare_output_path(str(directory), "a cat")
|
|
assert result == str(directory / ("a cat_1" + target_extension))
|
|
assert existing.read_bytes() == b"existing image"
|