Boyang ZHONG
dc72a35f4e
For n > 1, the image endpoints returned one preview grid instead of n images. The endpoints now run one model batch, save each decoded sample as its own file, and return one URL or base64 item for each image. This applies to generations and edits. A keyword-only return_samples field (default False) is added to OutputConfig, SamplingParam and ForwardBatch. Batched image requests set it and do not build preview frames. The worker keeps the sample tensor for these requests, and non-output ranks clear the field. Defaults and positional arguments do not change.
1319 lines
49 KiB
Python
1319 lines
49 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,
|
|
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"
|
|
|
|
|
|
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"
|