From 709457c71d855c288c5fa4fa6654dae3127b7ca5 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 4 Nov 2026 20:22:57 +0100 Subject: [PATCH 1/7] [bugfix] Honor explicit ThunderKittens build requests on aarch64 (#2062) --- fastvideo-kernel/CMakeLists.txt | 6 ++++++ fastvideo-kernel/README.md | 4 ++-- fastvideo-kernel/build.sh | 8 +++----- fastvideo-kernel/tests/test_build_arch.py | 23 ++++++++++++++++++++++- 4 files changed, 33 insertions(+), 8 deletions(-) diff --git a/fastvideo-kernel/CMakeLists.txt b/fastvideo-kernel/CMakeLists.txt index b84b32327..8849c0a96 100644 --- a/fastvideo-kernel/CMakeLists.txt +++ b/fastvideo-kernel/CMakeLists.txt @@ -303,6 +303,12 @@ endif() # there, and tk's base_types.cuh brace-initializes signed-char vector members # from char (narrowing error). Skip TK until upstream is aarch64-clean. if(ENABLE_TK_KERNELS AND CMAKE_SYSTEM_PROCESSOR MATCHES "^(aarch64|arm64)$") + if(FASTVIDEO_KERNEL_BUILD_TK STREQUAL "ON") + message(FATAL_ERROR + "fastvideo-kernel: FASTVIDEO_KERNEL_BUILD_TK=ON was requested, but the " + "ThunderKittens headers do not compile on ${CMAKE_SYSTEM_PROCESSOR}. " + "Use FASTVIDEO_KERNEL_BUILD_TK=AUTO or OFF to build with Triton fallbacks.") + endif() message(STATUS "ThunderKittens kernels: forced OFF on ${CMAKE_SYSTEM_PROCESSOR} (tk headers are not aarch64-clean)") set(ENABLE_TK_KERNELS OFF) endif() diff --git a/fastvideo-kernel/README.md b/fastvideo-kernel/README.md index b5f08734f..f332e2fed 100644 --- a/fastvideo-kernel/README.md +++ b/fastvideo-kernel/README.md @@ -9,7 +9,7 @@ Compiled CUDA extensions (CMake, see the build summary printed at the end of eve | Extension | Kernels | Sources | GPU arch | Build gate | |---|---|---|---|---| | `fastvideo_kernel._C.fastvideo_kernel_ops` | TurboDiffusion INT8 GEMM, quant, RMSNorm, LayerNorm | `csrc/turbodiffusion/` | every arch in `TORCH_CUDA_ARCH_LIST` | always built | -| same extension, optional part | ThunderKittens sliding-tile attention (`sta_fwd`) and VSA block-sparse (`block_sparse_fwd/bwd`) | `csrc/attention/*_h100.cu` | Hopper `sm_90a` only | `FASTVIDEO_KERNEL_BUILD_TK` (AUTO = ON iff `9.0a` is in the arch list; always OFF on aarch64 hosts — TK headers don't compile there) | +| same extension, optional part | ThunderKittens sliding-tile attention (`sta_fwd`) and VSA block-sparse (`block_sparse_fwd/bwd`) | `csrc/attention/*_h100.cu` | Hopper `sm_90a` only | `FASTVIDEO_KERNEL_BUILD_TK` (AUTO enables `9.0a` on supported hosts; ON requires TK and fails on aarch64 compiled-extension builds, where its headers do not compile) | | same extension, optional part | MiniMax-H3 block-sparse VSA forward (64- and 128-token blocks) | `csrc/attention/block_sparse*.cu` | Data-center Blackwell `sm_100a`/`sm_103a` | ON iff `10.0a` or `10.3a` is in `TORCH_CUDA_ARCH_LIST` | | same extension, optional part | fused NVLink Ulysses all-to-all | `csrc/comm/ulysses_all_to_all.cu` | CUDA | `FASTVIDEO_KERNEL_BUILD_ULYSSES_A2A` (AUTO = ON with NCCL 2.29+ device headers and library; always OFF on ROCm) | | `fp4attn_cuda`, `fp4quant_cuda` | FP4 attention + quantization ("attn_qat_infer", modified SageAttention3) | `attn_qat_infer/` | consumer/workstation Blackwell `sm_120a` (CUDA ≥ 12.8) + `sm_121a` (DGX Spark GB10; CUDA ≥ 13) | `FASTVIDEO_KERNEL_BUILD_ATTN_QAT_INFER` (AUTO = ON iff `12.0a` or `12.1a` is in the arch list) | @@ -46,7 +46,7 @@ Notes: ## Installation ### Standard Installation (Local Development) -This will automatically detect your GPU architecture. If an NVIDIA Hopper (H100/sm_90a) GPU is detected, ThunderKittens kernels will be enabled. Otherwise, they will be skipped, and the package will use Triton fallbacks at runtime. +This detects the GPU architecture and enables ThunderKittens kernels for Hopper (H100/sm_90a) on supported hosts. On aarch64, TK is unavailable and the package uses Triton fallbacks. For a compiled-extension build, set `CMAKE_ARGS="-DFASTVIDEO_KERNEL_BUILD_TK=ON"` to require TK; the build fails on aarch64 instead of silently omitting it. Before installation, set CUDA toolchain paths: diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh index 569417297..71745e10c 100755 --- a/fastvideo-kernel/build.sh +++ b/fastvideo-kernel/build.sh @@ -210,11 +210,9 @@ if [ "${GPU_BACKEND}" = "CUDA" ] && ! (( PYTHON_ONLY )); then fi if ! has_cmake_arg "FASTVIDEO_KERNEL_BUILD_TK"; then - if [ "${cc_major}" = "9" ] && [ "${cc_minor}" = "0" ]; then - CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=ON" - else - CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=OFF" - fi + # Let CMake detect both the target GPU arch and host support. Explicit ON + # remains a request to fail if TK cannot be built on the host. + CMAKE_ARGS="${CMAKE_ARGS:-} -DFASTVIDEO_KERNEL_BUILD_TK=AUTO" fi fi diff --git a/fastvideo-kernel/tests/test_build_arch.py b/fastvideo-kernel/tests/test_build_arch.py index 83aa945dc..45a200709 100644 --- a/fastvideo-kernel/tests/test_build_arch.py +++ b/fastvideo-kernel/tests/test_build_arch.py @@ -14,7 +14,9 @@ def _write_executable(path: Path, contents: str) -> None: path.chmod(0o755) -def _run_build(tmp_path: Path, capability: str, cuda_release: str) -> tuple[subprocess.CompletedProcess[str], Path]: +def _run_build( + tmp_path: Path, capability: str, cuda_release: str, cmake_args: str | None = None +) -> tuple[subprocess.CompletedProcess[str], Path]: kernel_dir = tmp_path / "fastvideo-kernel" kernel_dir.mkdir() shutil.copy2(BUILD_SCRIPT, kernel_dir / "build.sh") @@ -46,6 +48,8 @@ def _run_build(tmp_path: Path, capability: str, cuda_release: str) -> tuple[subp "FV_TEST_RESULT": str(result_dir), "PATH": f"{fake_bin}{os.pathsep}{env['PATH']}", }) + if cmake_args is not None: + env["CMAKE_ARGS"] = cmake_args result = subprocess.run( ["bash", "build.sh"], cwd=kernel_dir, @@ -82,3 +86,20 @@ def test_build_script_rejects_sm121_before_cmake_on_cuda_12(tmp_path: Path) -> N assert result.returncode != 0 assert "sm_121a requires CUDA Toolkit 13.0+" in result.stderr + + +@pytest.mark.parametrize("capability", ["9.0", "10.3"]) +def test_build_script_defers_tk_detection_to_cmake(tmp_path: Path, capability: str) -> None: + result, result_dir = _run_build(tmp_path, capability, "13.0") + + assert result.returncode == 0, result.stderr + assert "-DFASTVIDEO_KERNEL_BUILD_TK=AUTO" in (result_dir / "cmake_args").read_text() + + +@pytest.mark.parametrize("option", ["ON", "OFF"]) +def test_build_script_preserves_explicit_tk_option(tmp_path: Path, option: str) -> None: + result, result_dir = _run_build(tmp_path, "9.0", "13.0", f"-DFASTVIDEO_KERNEL_BUILD_TK={option}") + + assert result.returncode == 0, result.stderr + assert (result_dir / "cmake_args").read_text().count("-DFASTVIDEO_KERNEL_BUILD_TK=") == 1 + assert f"-DFASTVIDEO_KERNEL_BUILD_TK={option}" in (result_dir / "cmake_args").read_text() From 702e62f2b296678626e8288bfb81d81959f80b45 Mon Sep 17 00:00:00 2001 From: li-lizhe <82501526+li-lizhe@users.noreply.github.com> Date: Thu, 5 Nov 2026 04:28:49 +0800 Subject: [PATCH 2/7] [bugfix] SP warmup: use the rank's own device instead of hardcoded CUDA (#1988) Signed-off-by: li-lizhe <147392333@qq.com> Co-authored-by: aryan5v Co-authored-by: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> --- .buildkite/scripts/unit_test.sh | 1 + fastvideo/distributed/communication_op.py | 13 ++- .../distributed/test_sp_warmup_device.py | 85 +++++++++++++++++++ 3 files changed, 96 insertions(+), 3 deletions(-) create mode 100644 fastvideo/tests/distributed/test_sp_warmup_device.py diff --git a/.buildkite/scripts/unit_test.sh b/.buildkite/scripts/unit_test.sh index 8d6a61bc1..0723cc97b 100644 --- a/.buildkite/scripts/unit_test.sh +++ b/.buildkite/scripts/unit_test.sh @@ -24,6 +24,7 @@ exec pytest \ ./fastvideo/tests/stages/ \ ./fastvideo/tests/ops/ \ ./fastvideo/tests/worker/ \ + ./fastvideo/tests/distributed/test_sp_warmup_device.py \ ./fastvideo/tests/training/test_runner.py \ ./fastvideo/tests/training/test_trackers.py \ ./fastvideo/tests/training/test_ltx2_rope_fps.py \ diff --git a/fastvideo/distributed/communication_op.py b/fastvideo/distributed/communication_op.py index 502284f06..eb6d207a7 100644 --- a/fastvideo/distributed/communication_op.py +++ b/fastvideo/distributed/communication_op.py @@ -99,7 +99,8 @@ def warmup_sequence_parallel_communication(device: torch.device | None = None) - trigger the initialization upfront, before the first real forward pass. Args: - device: Device to use for warmup tensors. If None, uses CUDA device 0. + device: Device to use for warmup tensors. If None, uses the device type of the + accelerator the rank was placed on. """ global _sp_warmup_done @@ -114,8 +115,14 @@ def warmup_sequence_parallel_communication(device: torch.device | None = None) - _sp_warmup_done = True return + from fastvideo.platforms import current_platform + device_module = current_platform.get_torch_device() + if device is None: - device = torch.device("cuda") + # The warmup tensors go through the SP collectives, so they have to live on the + # accelerator this rank runs on: a "cuda" default is not a valid device on a build + # without CUDA (Ascend NPU, MPS, ...) and the warmup would fail to allocate. + device = torch.device(current_platform.device_type) logger.info("Warming up sequence parallel communication (SP=%d)...", sp_world_size) @@ -149,7 +156,7 @@ def warmup_sequence_parallel_communication(device: torch.device | None = None) - _ = sequence_model_parallel_all_gather(dummy3, dim=2) # Synchronize to ensure warmup completes - torch.cuda.synchronize(device) + device_module.synchronize(device) # Clean up del dummy, dummy2, dummy3 diff --git a/fastvideo/tests/distributed/test_sp_warmup_device.py b/fastvideo/tests/distributed/test_sp_warmup_device.py new file mode 100644 index 000000000..1fa797da1 --- /dev/null +++ b/fastvideo/tests/distributed/test_sp_warmup_device.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 + +from __future__ import annotations + +import torch + +import fastvideo.platforms as platforms +from fastvideo.distributed import communication_op + + +class _RecordingSpGroup: + def __init__(self) -> None: + self.tensors: list[torch.Tensor] = [] + + def all_to_all_4D(self, tensor: torch.Tensor, scatter_dim: int = 2, gather_dim: int = 1) -> torch.Tensor: + self.tensors.append(tensor) + return tensor + + def all_gather(self, tensor: torch.Tensor, dim: int = -1) -> torch.Tensor: + self.tensors.append(tensor) + return tensor + + +class _RecordingDeviceModule: + def __init__(self) -> None: + self.synchronized: list = [] + + def synchronize(self, device=None) -> None: + self.synchronized.append(device) + + +class _RecordingPlatform: + def __init__(self, device_type: str) -> None: + self.device_type = device_type + self.device_module = _RecordingDeviceModule() + + def get_torch_device(self) -> _RecordingDeviceModule: + return self.device_module + + +def _patch_warmup_inputs(monkeypatch, platform: _RecordingPlatform, group: _RecordingSpGroup) -> None: + # Set the lazy source, not the public attribute: a pinned current_platform + # would leak into later tests that replace _current_platform. + monkeypatch.setattr(platforms, "_current_platform", platform) + if "current_platform" in vars(platforms): + monkeypatch.delitem(vars(platforms), "current_platform") + monkeypatch.setattr(communication_op, "model_parallel_is_initialized", lambda: True) + monkeypatch.setattr(communication_op, "get_sp_world_size", lambda: 2) + monkeypatch.setattr(communication_op, "get_sp_group", lambda: group) + monkeypatch.setattr(communication_op, "_sp_warmup_done", False) + + +def test_warmup_uses_the_active_platform_device(monkeypatch) -> None: + platform = _RecordingPlatform(device_type="cpu") + group = _RecordingSpGroup() + _patch_warmup_inputs(monkeypatch, platform, group) + + communication_op.warmup_sequence_parallel_communication() + + assert group.tensors, "the warmup must exercise the SP collectives" + assert {tensor.device.type for tensor in group.tensors} == {"cpu"} + assert platform.device_module.synchronized == [torch.device("cpu")] + + +def test_warmup_keeps_an_explicit_device(monkeypatch) -> None: + platform = _RecordingPlatform(device_type="cuda") + group = _RecordingSpGroup() + _patch_warmup_inputs(monkeypatch, platform, group) + + communication_op.warmup_sequence_parallel_communication(torch.device("cpu")) + + assert {tensor.device.type for tensor in group.tensors} == {"cpu"} + assert platform.device_module.synchronized == [torch.device("cpu")] + + +def test_warmup_is_skipped_without_sequence_parallelism(monkeypatch) -> None: + platform = _RecordingPlatform(device_type="cpu") + group = _RecordingSpGroup() + _patch_warmup_inputs(monkeypatch, platform, group) + monkeypatch.setattr(communication_op, "get_sp_world_size", lambda: 1) + + communication_op.warmup_sequence_parallel_communication() + + assert group.tensors == [] + assert platform.device_module.synchronized == [] From ebe8fb489cc0c9d153a227732eeda10c3d9fa4e4 Mon Sep 17 00:00:00 2001 From: 0z5a Date: Wed, 4 Nov 2026 21:40:47 +0100 Subject: [PATCH 3/7] [bugfix] Fix Python-only kernel build default (#2064) --- fastvideo-kernel/build.sh | 2 +- fastvideo-kernel/tests/test_build_arch.py | 30 +++++++++++++++++++++-- 2 files changed, 29 insertions(+), 3 deletions(-) diff --git a/fastvideo-kernel/build.sh b/fastvideo-kernel/build.sh index 71745e10c..a0d6a6797 100755 --- a/fastvideo-kernel/build.sh +++ b/fastvideo-kernel/build.sh @@ -80,7 +80,7 @@ extension_mode_from_cmake_args() { && printf '%s' "${BASH_REMATCH[2]^^}" } if (( PYTHON_ONLY )); then - _explicit_mode="$(extension_mode_from_cmake_args)" + _explicit_mode="$(extension_mode_from_cmake_args || true)" if [[ -n "${_explicit_mode}" && "${_explicit_mode}" != "OFF" ]]; then echo "ERROR: --python-only conflicts with CMAKE_ARGS -DFASTVIDEO_KERNEL_BUILD_EXTENSION=${_explicit_mode}" >&2 exit 2 diff --git a/fastvideo-kernel/tests/test_build_arch.py b/fastvideo-kernel/tests/test_build_arch.py index 45a200709..14ca51adc 100644 --- a/fastvideo-kernel/tests/test_build_arch.py +++ b/fastvideo-kernel/tests/test_build_arch.py @@ -15,7 +15,11 @@ def _write_executable(path: Path, contents: str) -> None: def _run_build( - tmp_path: Path, capability: str, cuda_release: str, cmake_args: str | None = None + tmp_path: Path, + capability: str, + cuda_release: str, + cmake_args: str | None = None, + extra_args: list[str] | None = None, ) -> tuple[subprocess.CompletedProcess[str], Path]: kernel_dir = tmp_path / "fastvideo-kernel" kernel_dir.mkdir() @@ -51,7 +55,7 @@ def _run_build( if cmake_args is not None: env["CMAKE_ARGS"] = cmake_args result = subprocess.run( - ["bash", "build.sh"], + ["bash", "build.sh", *(extra_args or [])], cwd=kernel_dir, env=env, capture_output=True, @@ -103,3 +107,25 @@ def test_build_script_preserves_explicit_tk_option(tmp_path: Path, option: str) assert result.returncode == 0, result.stderr assert (result_dir / "cmake_args").read_text().count("-DFASTVIDEO_KERNEL_BUILD_TK=") == 1 assert f"-DFASTVIDEO_KERNEL_BUILD_TK={option}" in (result_dir / "cmake_args").read_text() + + +@pytest.mark.parametrize("cmake_args", [None, "-DOTHER_OPTION=ON", "-DFASTVIDEO_KERNEL_BUILD_EXTENSION=OFF"]) +@pytest.mark.parametrize("build_args", [["--python-only"], ["--rocm", "--python-only"]]) +def test_python_only_build_defaults_extension_to_off( + tmp_path: Path, cmake_args: str | None, build_args: list[str] +) -> None: + result, result_dir = _run_build(tmp_path, "9.0", "13.0", cmake_args, build_args) + + assert result.returncode == 0, result.stderr + assert (result_dir / "cmake_args").read_text().count("-DFASTVIDEO_KERNEL_BUILD_EXTENSION=OFF") == 1 + if cmake_args == "-DOTHER_OPTION=ON": + assert cmake_args in (result_dir / "cmake_args").read_text() + if "--rocm" in build_args: + assert "-DGPU_BACKEND=ROCM" in (result_dir / "cmake_args").read_text() + + +def test_python_only_build_rejects_extension_on(tmp_path: Path) -> None: + result, _ = _run_build(tmp_path, "9.0", "13.0", "-DFASTVIDEO_KERNEL_BUILD_EXTENSION=ON", ["--python-only"]) + + assert result.returncode != 0 + assert "--python-only conflicts" in result.stderr From 40651713d50da7065dde9a6f67c955915f462c4c Mon Sep 17 00:00:00 2001 From: Kostiantyn Dementiev <68417998+konstde00@users.noreply.github.com> Date: Thu, 5 Nov 2026 01:14:33 +0200 Subject: [PATCH 4/7] [bugfix] Count the wait for a concurrency slot in bench latency (#1960) --- fastvideo/entrypoints/cli/bench_serving.py | 42 ++++++++--- .../test_bench_serving_counts_queue_wait.py | 75 +++++++++++++++++++ 2 files changed, 105 insertions(+), 12 deletions(-) create mode 100644 fastvideo/tests/entrypoints/test_bench_serving_counts_queue_wait.py diff --git a/fastvideo/entrypoints/cli/bench_serving.py b/fastvideo/entrypoints/cli/bench_serving.py index a51fe7974..33054eb51 100644 --- a/fastvideo/entrypoints/cli/bench_serving.py +++ b/fastvideo/entrypoints/cli/bench_serving.py @@ -22,6 +22,7 @@ import re import time import uuid from abc import ABC, abstractmethod +from collections.abc import Awaitable, Callable from dataclasses import dataclass, field from typing import Any @@ -62,6 +63,7 @@ class RequestFuncOutput: start_time: float = 0.0 response_body: dict[str, Any] = field(default_factory=dict) peak_memory_mb: float = 0.0 + queue_wait: float = 0.0 class BaseDataset(ABC): @@ -561,12 +563,31 @@ async def async_request_video_sglang( return output +async def limited_request( + request_func: Callable[[RequestFuncInput, aiohttp.ClientSession, tqdm | None], Awaitable[RequestFuncOutput]], + req: RequestFuncInput, + session: aiohttp.ClientSession, + pbar: tqdm | None, + semaphore: asyncio.Semaphore | None, + arrival_time: float, +) -> RequestFuncOutput: + """Send `req` once a concurrency slot is free; `queue_wait` counts from `arrival_time`.""" + if semaphore is None: + output = await request_func(req, session, pbar) + else: + async with semaphore: + output = await request_func(req, session, pbar) + output.queue_wait = output.start_time - arrival_time + return output + + def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float): success_outputs = [o for o in outputs if o.success] error_outputs = [o for o in outputs if not o.success] num_success = len(success_outputs) latencies = [o.latency for o in success_outputs] + e2e_latencies = [o.queue_wait + o.latency for o in success_outputs] peak_memories = [o.peak_memory_mb for o in success_outputs if o.peak_memory_mb is not None and o.peak_memory_mb > 0] metrics = { @@ -578,6 +599,9 @@ def calculate_metrics(outputs: list[RequestFuncOutput], total_duration: float): "latency_median": np.median(latencies) if latencies else 0, "latency_p99": np.percentile(latencies, 99) if latencies else 0, "latency_p50": np.percentile(latencies, 50) if latencies else 0, + "e2e_latency_mean": np.mean(e2e_latencies) if e2e_latencies else 0, + "e2e_latency_median": np.median(e2e_latencies) if e2e_latencies else 0, + "e2e_latency_p99": np.percentile(e2e_latencies, 99) if e2e_latencies else 0, "peak_memory_mb_max": max(peak_memories) if peak_memories else 0, "peak_memory_mb_mean": np.mean(peak_memories) if peak_memories else 0, "peak_memory_mb_median": np.median(peak_memories) if peak_memories else 0, @@ -671,17 +695,6 @@ async def benchmark(args: argparse.Namespace) -> None: # Limit concurrency semaphore = asyncio.Semaphore(args.max_concurrency) if args.max_concurrency is not None else None - async def limited_request_func( - req: RequestFuncInput, - session: aiohttp.ClientSession, - pbar: tqdm | None, - ) -> RequestFuncOutput: - if semaphore: - async with semaphore: - return await request_func(req, session, pbar) - else: - return await request_func(req, session, pbar) - # Run benchmark pbar = tqdm(total=len(requests_list), disable=args.disable_tqdm) @@ -689,12 +702,14 @@ async def benchmark(args: argparse.Namespace) -> None: start_time = time.perf_counter() tasks = [] for req in requests_list: + arrival_time = start_time if args.request_rate != float("inf"): # Poisson process: inter-arrival times follow exponential distribution interval = np.random.exponential(1.0 / args.request_rate) await asyncio.sleep(interval) + arrival_time = time.perf_counter() - task = asyncio.create_task(limited_request_func(req, session, pbar)) + task = asyncio.create_task(limited_request(request_func, req, session, pbar, semaphore, arrival_time)) tasks.append(task) outputs = await asyncio.gather(*tasks) @@ -729,6 +744,9 @@ async def benchmark(args: argparse.Namespace) -> None: print("{:<40} {:<15.4f}".format("Latency Mean (s):", metrics["latency_mean"])) print("{:<40} {:<15.4f}".format("Latency Median (s):", metrics["latency_median"])) print("{:<40} {:<15.4f}".format("Latency P99 (s):", metrics["latency_p99"])) + print("{:<40} {:<15.4f}".format("E2E Latency Mean, incl. queue (s):", metrics["e2e_latency_mean"])) + print("{:<40} {:<15.4f}".format("E2E Latency Median, incl. queue (s):", metrics["e2e_latency_median"])) + print("{:<40} {:<15.4f}".format("E2E Latency P99, incl. queue (s):", metrics["e2e_latency_p99"])) if metrics["peak_memory_mb_max"] > 0: print(f"{'-' * 50}") diff --git a/fastvideo/tests/entrypoints/test_bench_serving_counts_queue_wait.py b/fastvideo/tests/entrypoints/test_bench_serving_counts_queue_wait.py new file mode 100644 index 000000000..b4c46842d --- /dev/null +++ b/fastvideo/tests/entrypoints/test_bench_serving_counts_queue_wait.py @@ -0,0 +1,75 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only contracts for how `fastvideo bench` measures request latency.""" + +from __future__ import annotations + +import asyncio +from types import SimpleNamespace + +import numpy as np +import pytest + +from fastvideo.entrypoints.cli.bench_serving import ( + RequestFuncInput, + RequestFuncOutput, + calculate_metrics, + limited_request, +) + + +def _output(latency: float, queue_wait: float, success: bool = True) -> RequestFuncOutput: + return RequestFuncOutput(success=success, latency=latency, queue_wait=queue_wait) + + +def test_e2e_latency_includes_the_wait_for_a_concurrency_slot(): + outputs = [_output(latency=1.0, queue_wait=float(k)) for k in range(5)] + outputs.append(_output(latency=0.1, queue_wait=9.0, success=False)) + + metrics = calculate_metrics(outputs, total_duration=5.0) + + assert metrics["latency_p99"] == pytest.approx(1.0) + assert metrics["e2e_latency_mean"] == pytest.approx(3.0) + assert metrics["e2e_latency_median"] == pytest.approx(3.0) + assert metrics["e2e_latency_p99"] == pytest.approx(np.percentile([1, 2, 3, 4, 5], 99)) + assert metrics["failed_requests"] == 1 + + +def test_e2e_latency_equals_latency_without_a_wait(): + outputs = [_output(latency=x, queue_wait=0.0) for x in (0.5, 1.5, 2.5)] + + metrics = calculate_metrics(outputs, total_duration=3.0) + + for stat in ("mean", "median", "p99"): + assert metrics[f"e2e_latency_{stat}"] == pytest.approx(metrics[f"latency_{stat}"]) + + +def test_limited_request_counts_the_wait_from_arrival(): + clock = SimpleNamespace(now=0.0) + started = {"a": asyncio.Event(), "b": asyncio.Event()} + done = {"a": asyncio.Event(), "b": asyncio.Event()} + + async def fake_request(req, session, pbar): + output = RequestFuncOutput(start_time=clock.now) + started[req.prompt].set() + await done[req.prompt].wait() + output.latency = clock.now - output.start_time + output.success = True + return output + + async def run() -> list: + semaphore = asyncio.Semaphore(1) + reqs = [RequestFuncInput(prompt=p, api_url="http://unused", model="m") for p in ("a", "b")] + tasks = [asyncio.create_task(limited_request(fake_request, r, None, None, semaphore, 0.0)) for r in reqs] + clock.now = 0.5 + await started["a"].wait() + clock.now = 3.0 + done["a"].set() + await started["b"].wait() + clock.now = 5.0 + done["b"].set() + return await asyncio.gather(*tasks) + + a, b = asyncio.run(run()) + + assert (a.queue_wait, a.latency) == (0.5, 2.5) + assert (b.queue_wait, b.latency) == (3.0, 2.0) From 9e1915ceb44a7da3dba2da5078730d26d7818605 Mon Sep 17 00:00:00 2001 From: Haoyu Cheng <149918604+CHYYYYYYYY@users.noreply.github.com> Date: Wed, 4 Nov 2026 15:25:14 -0800 Subject: [PATCH 5/7] [bugfix] Honor output_video_name when selecting output filenames (#2044) --- docs/getting_started/v1_api.md | 36 ++ fastvideo/entrypoints/video_generator.py | 70 ++-- .../tests/entrypoints/test_video_generator.py | 346 ++++++++++++++++++ 3 files changed, 426 insertions(+), 26 deletions(-) diff --git a/docs/getting_started/v1_api.md b/docs/getting_started/v1_api.md index 1104fb86c..09f6e8a69 100644 --- a/docs/getting_started/v1_api.md +++ b/docs/getting_started/v1_api.md @@ -26,3 +26,39 @@ to receive the decoded CPU tensor in `GenerationResult.samples` without building preview grids or writing media. Pixel samples retain the existing float `[0, 1]` contract. `return_frames=True` continues to return both samples and preview frames. The legacy sampling/keyword interface accepts the same `return_samples` option. + +### Choosing an output filename + +Set `request.output.output_video_name` to choose a filename when +`request.output.output_path` is a directory: + +```python +result = generator.generate({ + "prompt": "A waterfall in the mountains", + "output": {"output_path": "outputs/", "output_video_name": "my_clip"}, +}) +# Saves outputs/my_clip.mp4 when that file does not already exist. +``` + +An explicit file path such as `outputs/custom.mp4` takes precedence over +`output_video_name`. Without a name, the filename is derived from the prompt +(or the input video for prompt-less super-resolution pipelines). Names are +sanitized, and existing files are preserved by adding numeric suffixes such as +`my_clip_1.mp4`, including when processing multiple prompts. A matching extension +in the name is not duplicated. Image workloads accept PNG, JPEG, or WebP names +and default to PNG; audio workloads use WAV. + +For requests that save media, the sanitized filename is checked against the +destination filesystem's filename byte limit, including its extension and any +collision suffix. If the limit cannot be determined, 255 bytes is used. +Overlong names are rejected before model execution; use a shorter name rather +than relying on automatic truncation. Multibyte characters may use more than +one byte each. Prompt-file batches retain their existing per-prompt error +handling. + +Requests with `save_video=False` or `output_type="latent"` do not prepare an +output file: filename validation, directory creation, and collision checks are +skipped. This also applies to latent output when `save_video=True`. + +The legacy `SamplingParam` and `generate_video` keyword interfaces accept the +same `output_video_name` option. diff --git a/fastvideo/entrypoints/video_generator.py b/fastvideo/entrypoints/video_generator.py index 3ebac451b..74a06e0a2 100644 --- a/fastvideo/entrypoints/video_generator.py +++ b/fastvideo/entrypoints/video_generator.py @@ -637,6 +637,7 @@ class VideoGenerator: prompt_embeds = kwargs.pop("prompt_embeds", None) sampling_param.update(kwargs) kwargs["_extra_overrides"] = extra_overrides + save_to_disk = sampling_param.save_video and fastvideo_args.output_type != "latent" prompt_txt_path = sampling_param.prompt_path or fastvideo_args.prompt_txt if self.executor.uses_spmd_execution: @@ -667,7 +668,11 @@ class VideoGenerator: try: # Generate video for this prompt using the same logic below output_path = self._run_on_output_rank( - partial(self._prepare_output_path, sampling_param.output_path, batch_prompt), + partial(self._prepare_output_path, + sampling_param.output_path, + batch_prompt, + output_video_name=sampling_param.output_video_name, + save_to_disk=save_to_disk), description="output-path setup", ) kwargs["output_path"] = output_path @@ -712,8 +717,10 @@ class VideoGenerator: else: raise ValueError("Either prompt or prompt_txt must be provided") output_path = self._run_on_output_rank( - lambda: self._prepare_output_path(sampling_param.output_path, output_name_hint - if output_name_hint else prompt), + lambda: self._prepare_output_path(sampling_param.output_path, + output_name_hint if output_name_hint else prompt, + output_video_name=sampling_param.output_video_name, + save_to_disk=save_to_disk), description="output-path setup", ) kwargs["output_path"] = output_path @@ -744,26 +751,23 @@ class VideoGenerator: self, output_path: str, prompt: str, + output_video_name: str | None = None, + *, + save_to_disk: bool = True, ) -> str: - """Build a unique, sanitized output file path. + if not save_to_disk: + return output_path - Image workloads preserve explicit PNG, JPEG, or WebP extensions and - otherwise default to ``.png``. Video workloads use ``.mp4``. - - - If ``output_path`` already carries the correct extension, treat it - as a file path. - - Otherwise, treat ``output_path`` as a directory and derive the - filename from the prompt. - - Invalid filename characters are removed; if the name changes, a - warning is logged. - - If the target path already exists, a numeric suffix is appended. - """ base_path, extension = os.path.splitext(output_path) extension_lower = extension.lower() is_directory = os.path.isdir(output_path) if self._is_image_workload(): - target_ext = extension_lower if not is_directory and extension_lower in {".png", ".jpg", ".jpeg", ".webp" - } else ".png" + image_extensions = {".png", ".jpg", ".jpeg", ".webp"} + name_extension = os.path.splitext(output_video_name or "")[1].lower() + if not is_directory and extension_lower in image_extensions: + target_ext = extension_lower + else: + target_ext = name_extension if name_extension in image_extensions else ".png" elif self._is_audio_workload(): target_ext = ".wav" else: @@ -801,20 +805,34 @@ class VideoGenerator: target_ext, ) output_dir = output_path - prompt_component = _sanitize_filename_component(prompt[:100]) - out_name = f"{prompt_component}{target_ext}" + name_component = output_video_name or prompt[:100] + if output_video_name and name_component.lower().endswith(target_ext): + name_component = name_component[:-len(target_ext)] + sanitized_name = _sanitize_filename_component(name_component) + out_name = f"{sanitized_name}{target_ext}" if output_dir: os.makedirs(output_dir, exist_ok=True) - new_output_path = os.path.join(output_dir, out_name) - counter = 1 - while os.path.exists(new_output_path): - name_part, ext_part = os.path.splitext(out_name) - new_name = f"{name_part}_{counter}{ext_part}" - new_output_path = os.path.join(output_dir, new_name) + try: + name_max = os.pathconf(output_dir or ".", "PC_NAME_MAX") + except (AttributeError, OSError, ValueError): + name_max = 255 + if name_max <= 0: + name_max = 255 + + name_part, ext_part = os.path.splitext(out_name) + candidate_name = out_name + counter = 0 + while True: + if len(os.fsencode(candidate_name)) > name_max: + raise ValueError(f"Output filename exceeds the {name_max}-byte filesystem limit " + "(including extension and collision suffix). Use a shorter output name.") + new_output_path = os.path.join(output_dir, candidate_name) + if not os.path.exists(new_output_path): + return new_output_path counter += 1 - return new_output_path + candidate_name = f"{name_part}_{counter}{ext_part}" def _generate_single_video( self, diff --git a/fastvideo/tests/entrypoints/test_video_generator.py b/fastvideo/tests/entrypoints/test_video_generator.py index 8143305fe..92a59aa1e 100644 --- a/fastvideo/tests/entrypoints/test_video_generator.py +++ b/fastvideo/tests/entrypoints/test_video_generator.py @@ -41,6 +41,7 @@ def _new_runtime_video_generator() -> VideoGenerator: 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), ) @@ -265,6 +266,351 @@ def test_prepare_output_path_empty_prompt_fallback(tmp_path): 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(), From 229c43c8d31ff2d5bffadd13f784b10f07315f94 Mon Sep 17 00:00:00 2001 From: tianwei-mo Date: Wed, 4 Nov 2026 17:12:34 -0800 Subject: [PATCH 6/7] [feat] MiniMax-H3 Ref2VA SFT: train each clip at its own canvas and frame count (#2003) Co-authored-by: Claude Opus 5.5 --- docs/training/minimax_h3_ref2va_preprocess.md | 11 ++ .../preprocess_minimax_h3_overfit.py | 66 ++++++++-- .../preprocess_minimax_h3_ref2va.py | 30 ++++- .../test_minimax_h3_row_length_latents.py | 66 ++++++++++ ...t_minimax_h3_native_geometry_preprocess.py | 120 ++++++++++++++++++ .../train/models/minimax_h3/minimax_h3.py | 19 +++ 6 files changed, 292 insertions(+), 20 deletions(-) create mode 100644 fastvideo/tests/train/models/test_minimax_h3_row_length_latents.py create mode 100644 fastvideo/tests/workflow/test_minimax_h3_native_geometry_preprocess.py diff --git a/docs/training/minimax_h3_ref2va_preprocess.md b/docs/training/minimax_h3_ref2va_preprocess.md index aa5033fd4..f0fbdcc28 100644 --- a/docs/training/minimax_h3_ref2va_preprocess.md +++ b/docs/training/minimax_h3_ref2va_preprocess.md @@ -12,6 +12,17 @@ format, which stores already augmented `ref_visual_anchor` tensors. Keep using `preprocess_minimax_h3_ref2va` for that format. Do not rename or directly combine shards from these two schemas. +The SFT preprocessor keeps each target clip's own geometry: the canvas follows +the clip's aspect ratio at the inference short edge (`resolve_canvas_size`), and +the frame count is the longest `17 * n + 5` that fits the clip, capped at 362 +(about 15 seconds; a longer clip trains on its start). Clips under 124 frames at +24 fps fail. Target audio is cut to the samples covered by that frame count's +whole audio latents, so the record matches what inference generates. To train +on those lengths, leave `training.data.num_latent_t` and `num_frames` unset in +the training config. When they are set, every row is cropped to that length; +for example, the 124-frame overfit values crop a 10-second clip to its first +5.17 seconds while its caption still describes all 10. + ## Input Provide your own local MiniMax-H3 checkpoint with `vae`, `audio_vae`, `tokenizer`, diff --git a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py index 9ef03db80..cd786f313 100644 --- a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py +++ b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_overfit.py @@ -25,7 +25,15 @@ from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2va from fastvideo.fastvideo_args import FastVideoArgs from fastvideo.models.loader.component_loader import PipelineComponentLoader from fastvideo.pipelines import ForwardBatch -from fastvideo.pipelines.basic.minimax_h3.packing import MINIMAX_H3_FPS +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_AUDIO_LATENTS_PER_SECOND, + MINIMAX_H3_FPS, + MINIMAX_H3_FRAMES_PER_CHUNK, + MINIMAX_H3_LATENTS_PER_CHUNK, + MINIMAX_H3_MAX_ALIGNED_FRAMES, + MINIMAX_H3_MIN_ALIGNED_FRAMES, + audio_latent_num_frames, +) from fastvideo.pipelines.basic.minimax_h3.reference import ( decode_reference_video, prepare_reference_frames, @@ -60,20 +68,41 @@ def _init_single_process_distributed() -> None: maybe_init_distributed_environment_and_model_parallel(1, 1) -def load_training_media(path: Path) -> tuple[np.ndarray, torch.Tensor]: +def native_num_frames(available: int) -> int: + """Return the longest H3-native ``17 * n + 5`` frame count within ``available``. + + Counts are capped at the largest accepted bucket, so a long clip trains on + its first 15 seconds instead of failing. + """ + if available < MINIMAX_H3_MIN_ALIGNED_FRAMES: + raise ValueError(f"Training video has {available} frames at {MINIMAX_H3_FPS} FPS; " + f"at least {MINIMAX_H3_MIN_ALIGNED_FRAMES} are required") + num_frames = min(available, MINIMAX_H3_MAX_ALIGNED_FRAMES) + while num_frames % MINIMAX_H3_FRAMES_PER_CHUNK != MINIMAX_H3_LATENTS_PER_CHUNK: + num_frames -= 1 + return num_frames + + +def load_training_media(path: Path, num_frames: int | None = NUM_FRAMES) -> tuple[np.ndarray, torch.Tensor]: """Decode paired targets at the geometry and rates consumed by H3. - The video uses exactly 124 frames, and the audio is capped at the duration - of those frames before variational autoencoder (VAE) encoding. + With the default ``num_frames`` the video must resolve to the fixed + 124-frame 768x1344 overfit geometry. ``num_frames=None`` keeps the clip's + own canvas (``resolve_canvas_size``) and its longest native frame count. + The audio is capped at the duration of the selected frames, and at the + samples their ``audio_latent_num_frames`` cover, before VAE encoding. """ decoded_frames, source_fps, soundtrack = decode_reference_video(path) resampled_frames = resample_reference_frames(decoded_frames, source_fps) - if resampled_frames.shape[0] < NUM_FRAMES: + fixed_geometry = num_frames is not None + if num_frames is None: + num_frames = native_num_frames(resampled_frames.shape[0]) + if resampled_frames.shape[0] < num_frames: raise ValueError(f"Training video {path} has {resampled_frames.shape[0]} frames at {MINIMAX_H3_FPS} FPS; " - f"{NUM_FRAMES} are required") - frames = prepare_reference_frames(resampled_frames, NUM_FRAMES) - if frames.shape != (NUM_FRAMES, VIDEO_HEIGHT, VIDEO_WIDTH, 3): - raise ValueError(f"Training video must resolve to {(NUM_FRAMES, VIDEO_HEIGHT, VIDEO_WIDTH, 3)}, " + f"{num_frames} are required") + frames = prepare_reference_frames(resampled_frames, num_frames) + if fixed_geometry and frames.shape != (num_frames, VIDEO_HEIGHT, VIDEO_WIDTH, 3): + raise ValueError(f"Training video must resolve to {(num_frames, VIDEO_HEIGHT, VIDEO_WIDTH, 3)}, " f"got {tuple(frames.shape)}") if soundtrack is None: raise ValueError(f"Training video {path} requires a soundtrack") @@ -82,8 +111,14 @@ def load_training_media(path: Path) -> tuple[np.ndarray, torch.Tensor]: waveform, source_sample_rate, AUDIO_SAMPLE_RATE, - max_duration=NUM_FRAMES / MINIMAX_H3_FPS, + max_duration=num_frames / MINIMAX_H3_FPS, ) + # The audio VAE pads to whole latents, so audio past the last whole latent + # would add one more than inference generates for these frames (e.g. 264 + # instead of 263 at 158 frames). The packed Ref2VA preprocessor applies the + # same cap. + samples_per_latent = AUDIO_SAMPLE_RATE // MINIMAX_H3_AUDIO_LATENTS_PER_SECOND + waveform = waveform[..., :audio_latent_num_frames(num_frames) * samples_per_latent].contiguous() return frames, waveform @@ -234,6 +269,9 @@ def build_parquet_record( video_latents: torch.Tensor, audio_latents: torch.Tensor, text_embedding: torch.Tensor, + num_frames: int = NUM_FRAMES, + height: int = VIDEO_HEIGHT, + width: int = VIDEO_WIDTH, ) -> dict[str, Any]: """Serialize one synchronized H3 sample for the Parquet schema collator. @@ -255,10 +293,10 @@ def build_parquet_record( "file_name": file_name, "caption": caption, "media_type": "video_with_audio", - "width": VIDEO_WIDTH, - "height": VIDEO_HEIGHT, - "num_frames": NUM_FRAMES, - "duration_sec": NUM_FRAMES / MINIMAX_H3_FPS, + "width": width, + "height": height, + "num_frames": num_frames, + "duration_sec": num_frames / MINIMAX_H3_FPS, "fps": float(MINIMAX_H3_FPS), "audio_sample_rate": AUDIO_SAMPLE_RATE, }) diff --git a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_ref2va.py b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_ref2va.py index c45032d4f..9bab4a203 100644 --- a/fastvideo/pipelines/preprocess/preprocess_minimax_h3_ref2va.py +++ b/fastvideo/pipelines/preprocess/preprocess_minimax_h3_ref2va.py @@ -305,6 +305,9 @@ def build_ref2va_parquet_record( ref_visual_anchor: torch.Tensor, ref_audio_anchor: torch.Tensor, references: list[MiniMaxH3PreparedReference], + num_frames: int = NUM_FRAMES, + height: int = VIDEO_HEIGHT, + width: int = VIDEO_WIDTH, ) -> dict[str, Any]: record = build_parquet_record( file_name=file_name, @@ -312,6 +315,9 @@ def build_ref2va_parquet_record( video_latents=video_latents, audio_latents=audio_latents, text_embedding=text_embedding, + num_frames=num_frames, + height=height, + width=width, ) record["schema_version"] = MINIMAX_H3_REF2VA_SCHEMA_VERSION record["text_token_tags"] = text_token_tags.to(dtype=torch.long, device="cpu").tolist() @@ -326,8 +332,13 @@ def validate_record_contract(record: dict[str, Any]) -> None: if missing: raise ValueError(f"Ref2VA record is missing schema fields: {missing}") expected_target_shapes = { - "vae_latent_shape": [24, video_latent_num_frames(NUM_FRAMES), VIDEO_HEIGHT // 16, VIDEO_WIDTH // 16], - "audio_latent_shape": [2, 32, audio_latent_num_frames(NUM_FRAMES)], + "vae_latent_shape": [ + 24, + video_latent_num_frames(record["num_frames"]), + record["height"] // 16, + record["width"] // 16, + ], + "audio_latent_shape": [2, 32, audio_latent_num_frames(record["num_frames"])], } for name, expected in expected_target_shapes.items(): if record[name] != expected: @@ -337,9 +348,12 @@ def validate_record_contract(record: dict[str, Any]) -> None: collate_minimax_h3_ref2va_rows([record]) -def _prepare_references(references: tuple[MiniMaxH3RawReference, ...], ) -> list[MiniMaxH3PreparedReference]: +def _prepare_references( + references: tuple[MiniMaxH3RawReference, ...], + num_frames: int = NUM_FRAMES, +) -> list[MiniMaxH3PreparedReference]: raw_references = build_minimax_h3_references(references) - prepared = [prepare_reference(reference, NUM_FRAMES, AUDIO_SAMPLE_RATE) for reference in raw_references] + prepared = [prepare_reference(reference, num_frames, AUDIO_SAMPLE_RATE) for reference in raw_references] for index, (raw, resolved) in enumerate(zip(references, prepared, strict=True)): expected_audio = raw.media_type in {"audio", "video_audio"} if resolved.has_audio != expected_audio: @@ -377,8 +391,9 @@ def _build_record( fastvideo_args: FastVideoArgs, patch_size: tuple[int, int, int], ) -> dict[str, Any]: - target_frames, target_waveform = load_training_media(sample.target_video_path) - references = _prepare_references(sample.references) + target_frames, target_waveform = load_training_media(sample.target_video_path, num_frames=None) + num_frames, height, width = (int(value) for value in target_frames.shape[:3]) + references = _prepare_references(sample.references, num_frames) # The established single-sample helpers deliberately release each large # component before the next one is loaded, keeping this multi-sample path @@ -410,6 +425,9 @@ def _build_record( ref_visual_anchor=ref_visual_anchor, ref_audio_anchor=ref_audio_anchor, references=references, + num_frames=num_frames, + height=height, + width=width, ) record["id"] = sample.sample_id validate_record_contract(record) diff --git a/fastvideo/tests/train/models/test_minimax_h3_row_length_latents.py b/fastvideo/tests/train/models/test_minimax_h3_row_length_latents.py new file mode 100644 index 000000000..eba9b3b5d --- /dev/null +++ b/fastvideo/tests/train/models/test_minimax_h3_row_length_latents.py @@ -0,0 +1,66 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Without a fixed training shape, H3 rows train at their own length instead of being cropped.""" +from types import SimpleNamespace + +import pytest +import torch + +from fastvideo.pipelines.basic.minimax_h3.packing import audio_latent_num_frames, video_latent_num_frames +from fastvideo.train.models.minimax_h3.minimax_h3 import MiniMaxH3Model + + +def _model(num_latent_t: int, num_frames: int) -> MiniMaxH3Model: + model = MiniMaxH3Model.__new__(MiniMaxH3Model) + model.training_config = SimpleNamespace(data=SimpleNamespace( + native_shape_bucketing=False, num_latent_t=num_latent_t, num_frames=num_frames)) + return model + + +def _row(num_frames: int, audio_frames: int | None = None) -> dict[str, torch.Tensor]: + if audio_frames is None: + audio_frames = audio_latent_num_frames(num_frames) + return { + "vae_latent": torch.zeros(1, 24, video_latent_num_frames(num_frames), 4, 2), + "audio_latent": torch.zeros(1, 2, 32, audio_frames), + } + + +def _resolve(model: MiniMaxH3Model, raw_batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, torch.Tensor]: + return model._resolve_clean_latents(raw_batch, "data", torch.float32, torch.device("cpu")) + + +@pytest.mark.parametrize("num_frames", [124, 158, 243, 362]) +def test_unset_shape_keeps_each_row_length(num_frames: int) -> None: + video, audio = _resolve(_model(0, 0), _row(num_frames)) + + assert video.shape[2] == video_latent_num_frames(num_frames) + assert audio.shape[-1] == audio_latent_num_frames(num_frames) + + +def test_unset_shape_rejects_mismatched_latent_clocks() -> None: + with pytest.raises(ValueError, match="158 video frames need 263"): + _resolve(_model(0, 0), _row(158, audio_frames=264)) + + +def test_fixed_shape_still_crops_and_warns_once(monkeypatch: pytest.MonkeyPatch) -> None: + import fastvideo.train.models.minimax_h3.minimax_h3 as module + + warnings = [] + monkeypatch.setattr(module.logger, "warning_once", warnings.append) + + video, audio = _resolve(_model(37, 124), _row(243)) + + assert video.shape[2] == 37 + assert audio.shape[-1] == audio_latent_num_frames(124) + assert len(warnings) == 1 and "Cropping vae_latent from 72 to" in warnings[0] + + +def test_fixed_shape_at_its_own_length_does_not_warn(monkeypatch: pytest.MonkeyPatch) -> None: + import fastvideo.train.models.minimax_h3.minimax_h3 as module + + warnings = [] + monkeypatch.setattr(module.logger, "warning_once", warnings.append) + + _resolve(_model(37, 124), _row(124)) + + assert warnings == [] diff --git a/fastvideo/tests/workflow/test_minimax_h3_native_geometry_preprocess.py b/fastvideo/tests/workflow/test_minimax_h3_native_geometry_preprocess.py new file mode 100644 index 000000000..4e9210337 --- /dev/null +++ b/fastvideo/tests/workflow/test_minimax_h3_native_geometry_preprocess.py @@ -0,0 +1,120 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Ref2VA SFT preprocessing keeps each target clip's own canvas and native frame count.""" + +from pathlib import Path + +import numpy as np +import pytest +import torch + +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_TEXT_TAG, + audio_latent_num_frames, + resolve_canvas_size, + video_latent_num_frames, +) +from fastvideo.pipelines.basic.minimax_h3.ref2va_manifest import MiniMaxH3Ref2VARawSample +from fastvideo.pipelines.preprocess import preprocess_minimax_h3_overfit as overfit +from fastvideo.pipelines.preprocess import preprocess_minimax_h3_ref2va as ref2va + +_SAMPLES_PER_AUDIO_LATENT = 800 # 32 kHz / 40 latents per second + + +@pytest.mark.parametrize( + ("available", "expected"), + [(124, 124), (140, 124), (141, 141), (157, 141), (209, 209), (362, 362), (1000, 362)], +) +def test_native_num_frames_rounds_down_to_17n_plus_5(available: int, expected: int) -> None: + assert overfit.native_num_frames(available) == expected + + +def test_native_num_frames_rejects_clips_under_five_seconds() -> None: + with pytest.raises(ValueError, match="at least 124"): + overfit.native_num_frames(123) + + +def _stub_media(monkeypatch: pytest.MonkeyPatch, num_source_frames: int, seconds: float) -> None: + """Decode to tiny 30-fps frames and a 48 kHz stereo soundtrack; skip the canvas resize.""" + frames = np.zeros((num_source_frames, 8, 4, 3), dtype=np.uint8) + waveform = torch.zeros(2, int(seconds * 48000)) + monkeypatch.setattr(overfit, "decode_reference_video", lambda _path: (frames, 30.0, (waveform, 48000))) + monkeypatch.setattr(overfit, "prepare_reference_frames", lambda decoded, num_frames: decoded[:num_frames]) + + +def test_native_media_keeps_native_length_and_whole_audio_latents(monkeypatch: pytest.MonkeyPatch) -> None: + # 210 frames at 30 fps resample to 168 at 24 fps; the longest native count is 158. + _stub_media(monkeypatch, num_source_frames=210, seconds=7.0) + + frames, waveform = overfit.load_training_media(Path("clip.mp4"), num_frames=None) + + assert frames.shape[0] == 158 + # 158 frames last 210666 samples, past the 263rd whole latent; the VAE would pad a 264th. + assert audio_latent_num_frames(158) == 263 + assert waveform.shape == (2, 263 * _SAMPLES_PER_AUDIO_LATENT) + + +def test_fixed_overfit_media_audio_length_is_unchanged(monkeypatch: pytest.MonkeyPatch) -> None: + _stub_media(monkeypatch, num_source_frames=200, seconds=7.0) + monkeypatch.setattr( + overfit, + "prepare_reference_frames", + lambda decoded, num_frames: np.zeros((num_frames, overfit.VIDEO_HEIGHT, overfit.VIDEO_WIDTH, 3), np.uint8), + ) + + frames, waveform = overfit.load_training_media(Path("clip.mp4")) + + # 124 frames end inside the 207th audio latent, so the whole-latent cap cuts nothing. + uncapped = overfit.prepare_reference_waveform(torch.zeros(2, 7 * 48000), + 48000, + overfit.AUDIO_SAMPLE_RATE, + max_duration=overfit.NUM_FRAMES / 24) + assert frames.shape[0] == overfit.NUM_FRAMES + assert uncapped.shape[-1] < audio_latent_num_frames(overfit.NUM_FRAMES) * _SAMPLES_PER_AUDIO_LATENT + assert waveform.shape == uncapped.shape + + +def test_build_record_uses_the_clip_geometry(monkeypatch: pytest.MonkeyPatch) -> None: + """A portrait 158-frame clip yields a record the training collator accepts.""" + num_frames = 158 + height, width = resolve_canvas_size(9, 16) + target = np.zeros((num_frames, height, width, 3), dtype=np.uint8) + waveform = torch.zeros(2, audio_latent_num_frames(num_frames) * _SAMPLES_PER_AUDIO_LATENT) + monkeypatch.setattr(ref2va, "load_training_media", lambda _path, num_frames: (target, waveform)) + monkeypatch.setattr(ref2va, "_prepare_references", lambda _refs, _num_frames: []) + monkeypatch.setattr( + ref2va, + "encode_video_latents", + lambda frames, *_args: torch.zeros(24, video_latent_num_frames(frames.shape[0]), frames.shape[1] // 16, + frames.shape[2] // 16), + ) + monkeypatch.setattr( + ref2va, + "encode_audio_latents", + lambda audio, *_args: torch.zeros(2, 32, audio.shape[-1] // _SAMPLES_PER_AUDIO_LATENT), + ) + monkeypatch.setattr(ref2va, "encode_ref_visual_anchor", lambda *_args: torch.empty(0, 96)) + monkeypatch.setattr(ref2va, "encode_ref_audio_anchor", lambda *_args: torch.empty(0, 32)) + monkeypatch.setattr( + ref2va, + "encode_ref2va_conditioning", + lambda *_args: (torch.zeros(4, 5120), torch.full((4, ), MINIMAX_H3_TEXT_TAG)), + ) + sample = MiniMaxH3Ref2VARawSample( + sample_id="portrait", + target_file="portrait.mp4", + target_video_path=Path("portrait.mp4"), + caption="A person talks to the camera.", + references=(), + ) + + record = ref2va._build_record( + sample, + model_path=Path("model"), + model_index={}, + fastvideo_args=None, + patch_size=(1, 2, 2), + ) + + assert (record["num_frames"], record["height"], record["width"]) == (num_frames, height, width) + assert record["vae_latent_shape"] == [24, video_latent_num_frames(num_frames), height // 16, width // 16] + assert record["audio_latent_shape"] == [2, 32, audio_latent_num_frames(num_frames)] diff --git a/fastvideo/train/models/minimax_h3/minimax_h3.py b/fastvideo/train/models/minimax_h3/minimax_h3.py index 313dec6ac..a8c022202 100644 --- a/fastvideo/train/models/minimax_h3/minimax_h3.py +++ b/fastvideo/train/models/minimax_h3/minimax_h3.py @@ -19,6 +19,7 @@ from fastvideo.dataset.sequence_capacity import select_sequence_capacity, valida from fastvideo.distributed import get_sp_group from fastvideo.forward_context import set_forward_context from fastvideo.layers.pdd import shifted_noise_amount +from fastvideo.logger import init_logger from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler from fastvideo.pipelines import TrainingBatch from fastvideo.pipelines.basic.minimax_h3.packing import ( @@ -49,6 +50,8 @@ from fastvideo.train.utils.module_state import apply_trainable from fastvideo.train.utils.moduleloader import load_module_from_path from fastvideo.train.utils.lora import finalize_lora_training, LoraConfig +logger = init_logger(__name__) + if TYPE_CHECKING: from fastvideo.train.utils.training_config import TrainingConfig @@ -370,10 +373,26 @@ class MiniMaxH3Model(ModelBase): if latents_source == "data" and native_shapes: self._validate_native_latents(raw_batch, video_latents, audio_latents) + elif latents_source == "data" and data_config.num_latent_t <= 0 and data_config.num_frames <= 0: + # With no fixed shape configured, each row trains at its own + # length, as the Ref2VA SFT preprocessor records it. Its video and + # audio latents must still describe the same frame count. + num_video_latents = int(video_latents.shape[2]) + if (num_video_latents - 2) % 5: + raise ValueError(f"vae_latent has {num_video_latents} frames; MiniMax H3 encodes 5 * n + 2") + num_frames = (num_video_latents - 2) // 5 * 17 + 5 + if audio_latents.shape[-1] != audio_latent_num_frames(num_frames): + raise ValueError(f"audio_latent has {audio_latents.shape[-1]} frames; {num_frames} video frames " + f"need {audio_latent_num_frames(num_frames)}") elif not native_shapes: # Preserve the legacy fixed-shape contract for configs that have # not opted into exact-shape bucketing. if data_config.num_latent_t > 0: + if video_latents.shape[2] > data_config.num_latent_t: + logger.warning_once( + f"Cropping vae_latent from {video_latents.shape[2]} to training.data.num_latent_t=" + f"{data_config.num_latent_t} frames, and audio to match. Leave num_latent_t and " + "num_frames unset to train each row at its own length.") video_latents = video_latents[:, :, :data_config.num_latent_t] expected_audio_frames = audio_latent_num_frames(data_config.num_frames) audio_latents = audio_latents[:, :, :, :expected_audio_frames] From 36dc971dc6a87dd1d40a97204908e7b524d40310 Mon Sep 17 00:00:00 2001 From: tianwei-mo Date: Wed, 4 Nov 2026 17:24:11 -0800 Subject: [PATCH 7/7] [feat] MiniMax-H3 Ref2VA: continue an earlier video from its last frames (#2001) Co-authored-by: Claude Opus 5.5 --- .../inference_schema_parity_inventory.yaml | 1 + docs/inference/fasth3_omniref.md | 31 +++ fastvideo/entrypoints/openai/protocol.py | 3 + .../entrypoints/openai/request_adapter.py | 14 ++ .../pipelines/basic/minimax_h3/packing.py | 50 +++- .../pipelines/basic/minimax_h3/reference.py | 86 ++++++- .../stages/minimax_h3_input_preparation.py | 15 +- .../stages/minimax_h3_latent_preparation.py | 52 +++- .../entrypoints/test_openai_serving_engine.py | 35 +++ .../stages/test_minimax_h3_continuation.py | 225 ++++++++++++++++++ 10 files changed, 498 insertions(+), 14 deletions(-) create mode 100644 fastvideo/tests/stages/test_minimax_h3_continuation.py diff --git a/docs/design/inference_schema_parity_inventory.yaml b/docs/design/inference_schema_parity_inventory.yaml index 8348e5af9..d878de1e3 100644 --- a/docs/design/inference_schema_parity_inventory.yaml +++ b/docs/design/inference_schema_parity_inventory.yaml @@ -628,6 +628,7 @@ surfaces: audio_reference: request.inputs.references video_path: request.inputs.video_path video_url: request.inputs.video_path + continuation_video: request.inputs.references video_params: request.sampling.width,height,num_frames,fps size: target: request.sampling.width,height diff --git a/docs/inference/fasth3_omniref.md b/docs/inference/fasth3_omniref.md index b695ca250..a0dad96d4 100644 --- a/docs/inference/fasth3_omniref.md +++ b/docs/inference/fasth3_omniref.md @@ -148,6 +148,37 @@ Select the NVFP4 decoder per configuration with holds base MiniMax-H3's full 36-block VAE, so both the 36-block quality tier and the 8-block fast tier apply. +## Continuing a video + +A Ref2VA request can continue an earlier video. It takes the earlier video's last +5, 22 or 39 frames (and their sound, if any) as a continuation clip, and the +generated video starts by regenerating those frames at the same positions: the +clip's latent rows take the coordinates of the target's first latent frames and +audio latents, the reference context of +[ComfyUI-H3-Continuum](https://github.com/ukr8b3g-cmyk/ComfyUI-H3-Continuum). +Drop the clip's frames, and the matching audio, from the new video, then append +the rest to the earlier one. + +- Over HTTP, set `continuation_video` to the clip's URL or local path (under + `server.media_roots`) beside the usual references. In Python, append + `MiniMaxH3Reference(source=clip, media_type="continuation")` to `references`. +- The clip is resampled to 24 fps and must then have 5, 22 or 39 frames on the + target canvas. A request takes at most one clip and still needs at least one + reference; the clip does not count toward the reference limits and is not + shown to the text encoder. +- `num_frames` counts the clip: 22 frames plus five new seconds is 142, + aligned up to 158. +- Write the prompt from the clip's actual end state. When the prompt's opening + contradicts the clip, the model tends to cut to a new shot right after it. +- Each continuation drifts a little toward higher saturation and sharpness, so + fewer, longer segments (two 15-second ones rather than six 5-second ones) hold + up better. + +Measured on a FastH3 OmniRef PDD student at 1344×768, 22-frame clips re-encoded +from the earlier output: the jump at a join was about half that of conditioning +the next segment on the earlier last frame as an image reference, and the sound +continued instead of restarting from near silence. + ## Environment The profiles set these; set them yourself when driving the YAML recipes. diff --git a/fastvideo/entrypoints/openai/protocol.py b/fastvideo/entrypoints/openai/protocol.py index f23185ceb..da6b940b5 100644 --- a/fastvideo/entrypoints/openai/protocol.py +++ b/fastvideo/entrypoints/openai/protocol.py @@ -131,6 +131,9 @@ class VideoGenerationRequest(BaseModel): # SGLang's legacy direct-video spellings. video_path: str | None = None video_url: str | None = None + # MiniMax-H3 Ref2VA: the last 5, 22 or 39 frames of an earlier video (URL or local path) that this request + # continues; see MiniMaxH3Reference(media_type="continuation"). + continuation_video: str | None = None video_params: VideoParams | None = None user: str | None = None task: str | None = None diff --git a/fastvideo/entrypoints/openai/request_adapter.py b/fastvideo/entrypoints/openai/request_adapter.py index cd731e27e..356892d16 100644 --- a/fastvideo/entrypoints/openai/request_adapter.py +++ b/fastvideo/entrypoints/openai/request_adapter.py @@ -218,6 +218,15 @@ async def prepare_reference_media( policy=policy) if remote: track(request.video_url) + if request.continuation_video: + remote = is_remote(request.continuation_video) + request.continuation_video = await _materialize_video(request.continuation_video, + os.path.join(uploads_dir, + f"{request_id}_continuation"), + max_bytes=max_bytes, + policy=policy) + if remote: + track(request.continuation_video) except BaseException: # Includes cancellation, so a disconnected client leaves no downloaded inputs behind. discard_materialized_media(created) @@ -375,6 +384,8 @@ def _apply_reference_inputs( references = [MiniMaxH3Reference(source=source, media_type="image") for source in images] references.extend(MiniMaxH3Reference(source=source, media_type="video") for source in videos) references.extend(MiniMaxH3Reference(source=source, media_type="audio") for source in audios) + if request.continuation_video: + references.append(MiniMaxH3Reference(source=request.continuation_video, media_type="continuation")) if references: try: kwargs["references"] = validate_references(references) @@ -384,6 +395,9 @@ def _apply_reference_inputs( if request.task is not None and model_family != "minimax_h3": raise RequestAdaptationError("The task selector is only defined for MiniMax-H3 servers.") + if request.continuation_video: + raise RequestAdaptationError("continuation_video requires a MiniMax-H3 server configured with " + "override_pipeline_cls_name=MiniMaxH3Ref2VAModularPipeline.") if model_family == OPEN_H3_REGENERATE_2K_FAMILY: if len(videos) != 1 or images or audios: diff --git a/fastvideo/pipelines/basic/minimax_h3/packing.py b/fastvideo/pipelines/basic/minimax_h3/packing.py index 5c91a7207..e9dcf184c 100644 --- a/fastvideo/pipelines/basic/minimax_h3/packing.py +++ b/fastvideo/pipelines/basic/minimax_h3/packing.py @@ -93,7 +93,7 @@ class MiniMaxH3PackedLayout: # Per-reference conditioning spans in packed order (Ref2VA only; empty # elsewhere): (kind, row_count, latent_shape) covering exactly the rows # between the text tokens and the target audio. kind is "audio" / "image" - # / "video"; latent_shape is the reference's raw latent (frames, height, + # / "video" / "continuation"; latent_shape is the reference's raw latent (frames, height, # width) for visual kinds and (0, 0, 0) for audio rows. A video reference # with sound contributes two entries — its audio rows precede its video # rows. Retained so VSA-H3 can 3-D-tile (and sparsify) reference-video @@ -167,6 +167,15 @@ class MiniMaxH3PackedBatchLayout: raise ValueError("H3 packed-batch token_tags must cover the complete sequence") +@dataclass(frozen=True) +class MiniMaxH3ContinuationLayout: + """Latent geometry of a continuation clip (``reference.MiniMaxH3PreparedContinuation``).""" + + num_latent_frames: int + # 0 for a clip without a soundtrack. + num_audio_latents: int + + def _validate_packed_rows(layout: MiniMaxH3PackedLayout | MiniMaxH3PackedBatchLayout) -> None: if type(layout.pad_tokens) is not int or layout.pad_tokens < 0: raise ValueError("pad_tokens must be a non-negative integer") @@ -510,8 +519,15 @@ def build_ref2va_packed_sequence( latent_width: int, num_audio_latents: int, patch_size: tuple[int, int, int], + continuation: MiniMaxH3ContinuationLayout | None = None, ) -> MiniMaxH3PackedLayout: - """Build `[text | ordered references | target audio | target video]`.""" + """Build `[text | ordered references | continuation clip | target audio | target video]`. + + A continuation clip follows the references and advances the rotary clock as a video reference with its + soundtrack would; its video rows then take the coordinates of the target's first latent frames and its audio rows + those of the target's first audio latents, so the clip conditions the start of the target in place (the + reference context of ComfyUI-H3-Continuum). VSA-H3 keeps its rows dense. + """ if patch_size != (1, 2, 2): raise ValueError(f"MiniMax-H3 Ref2VA requires patch_size=(1, 2, 2), got {patch_size}.") if text_token_tags.ndim != 1: @@ -544,6 +560,13 @@ def build_ref2va_packed_sequence( num_reference_video_rows = sum(visual_row_counts) num_reference_audio_rows = sum(reference.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS for reference in references if reference.has_audio) + rows_per_target_frame = (latent_height // patch_h) * (latent_width // patch_w) + if continuation is not None: + if not 0 < continuation.num_latent_frames < num_latent_frames or continuation.num_audio_latents < 0: + raise ValueError(f"Continuation clip {continuation} does not fit a target of {num_latent_frames} " + "latent frames.") + num_reference_video_rows += continuation.num_latent_frames * rows_per_target_frame + num_reference_audio_rows += continuation.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS sequence_length = (num_text_tokens + num_reference_video_rows + num_reference_audio_rows + num_target_audio_rows + num_target_video_rows) @@ -617,6 +640,29 @@ def build_ref2va_packed_sequence( else: raise ValueError(f"Unsupported prepared reference type: {reference.media_type!r}.") + if continuation is not None: + audio_count = continuation.num_audio_latents * MINIMAX_H3_AUDIO_CHANNELS + visual_row_count = continuation.num_latent_frames * rows_per_target_frame + if audio_count: + reference_segments.append(("audio", audio_count, (0, 0, 0))) + # Not "video": _h3_vsa_ref2va_segments keeps every other span dense. + reference_segments.append( + ("continuation", visual_row_count, (continuation.num_latent_frames, latent_height, latent_width))) + audio_rows = slice(cursor, cursor + audio_count) + video_rows = slice(audio_rows.stop, audio_rows.stop + visual_row_count) + cursor = video_rows.stop + if audio_count: + audio_indices.append(torch.arange(audio_rows.start, audio_rows.stop)) + video_indices.append(torch.arange(video_rows.start, video_rows.stop)) + rotary_time += max(float(continuation.num_audio_latents), + _reference_temporal_span(continuation.num_latent_frames)) + frame_time = temporal_position_grid(continuation.num_latent_frames, rotary_time) + position_ids[video_rows, 0] = frame_time.repeat_interleave(rows_per_target_frame) + position_ids[video_rows, 1:] = target_frame_grid.repeat(continuation.num_latent_frames // patch_t, 1) + if audio_count: + _fill_audio_positions(position_ids, audio_rows, continuation.num_audio_latents, rotary_time, + target_width_grid) + audio_start = cursor video_start = audio_start + num_target_audio_rows _fill_audio_positions(position_ids, slice(audio_start, video_start), num_audio_latents, rotary_time, diff --git a/fastvideo/pipelines/basic/minimax_h3/reference.py b/fastvideo/pipelines/basic/minimax_h3/reference.py index 965d3327d..76c83442c 100644 --- a/fastvideo/pipelines/basic/minimax_h3/reference.py +++ b/fastvideo/pipelines/basic/minimax_h3/reference.py @@ -34,14 +34,24 @@ MINIMAX_H3_MAX_REFERENCE_IMAGES = 9 MINIMAX_H3_MAX_REFERENCE_VIDEOS = 3 MINIMAX_H3_MAX_REFERENCE_AUDIOS = 3 MINIMAX_H3_MAX_REFERENCES = 12 +# A continuation clip ends on a causal-VAE chunk boundary: 2, 7 or 12 latent frames. +MINIMAX_H3_CONTINUATION_FRAMES = (5, 22, 39) +# ForwardBatch.extra key for the decoded continuation clip, set by input preparation. +MINIMAX_H3_PREPARED_CONTINUATION_KEY = "minimax_h3_prepared_continuation" @dataclass(frozen=True) class MiniMaxH3Reference: - """One ordered Ref2VA input; decoding is deferred until request preparation.""" + """One ordered Ref2VA input; decoding is deferred until request preparation. + + ``media_type="continuation"`` is not a reference: it is the end of an earlier video (5, 22 or 39 frames, on the + target canvas) that the request continues. The generated video starts by regenerating those frames, at the same + positions, so a caller joins the two by dropping them. At most one per request; it is not presented to the text + encoder and does not count toward the reference limits. + """ source: Any - media_type: Literal["image", "video", "audio"] = "image" + media_type: Literal["image", "video", "audio", "continuation"] = "image" soundtrack: Any | None = None fps: float | None = None sample_rate: int | None = None @@ -49,13 +59,13 @@ class MiniMaxH3Reference: def __post_init__(self) -> None: if self.source is None: raise ValueError("A MiniMax-H3 reference requires a media source.") - if self.media_type not in ("image", "video", "audio"): + if self.media_type not in ("image", "video", "audio", "continuation"): raise ValueError(f"Unsupported MiniMax-H3 reference type: {self.media_type!r}.") - if self.soundtrack is not None and self.media_type != "video": - raise ValueError("Only a video reference may carry a separate soundtrack.") - if self.fps is not None and (self.media_type != "video" or isinstance(self.fps, bool) + if self.soundtrack is not None and self.media_type not in ("video", "continuation"): + raise ValueError("Only a video reference or a continuation clip may carry a separate soundtrack.") + if self.fps is not None and (self.media_type not in ("video", "continuation") or isinstance(self.fps, bool) or not isinstance(self.fps, Real) or self.fps <= 0): - raise ValueError("Reference `fps` must be positive and is only valid for video.") + raise ValueError("Reference `fps` must be positive and is only valid for video and continuation clips.") if self.sample_rate is not None and (self.media_type == "image" or isinstance(self.sample_rate, bool) or not isinstance(self.sample_rate, Integral) or self.sample_rate <= 0): raise ValueError("Reference `sample_rate` must be positive and is only valid for audio-bearing media.") @@ -99,8 +109,13 @@ def validate_references(references: list[Any]) -> list[MiniMaxH3Reference]: raise ValueError("MiniMax-H3 Ref2VA requires at least one reference.") typed = [coerce_reference(reference) for reference in references] + if sum(reference.media_type == "continuation" for reference in typed) > 1: + raise ValueError("MiniMax-H3 accepts at most one continuation clip.") + typed_references = [reference for reference in typed if reference.media_type != "continuation"] + if not typed_references: + raise ValueError("A MiniMax-H3 continuation needs at least one reference besides its clip.") counts = { - media_type: sum(reference.media_type == media_type for reference in typed) + media_type: sum(reference.media_type == media_type for reference in typed_references) for media_type in ("image", "video", "audio") } limits = { @@ -111,9 +126,9 @@ def validate_references(references: list[Any]) -> list[MiniMaxH3Reference]: for media_type, limit in limits.items(): if counts[media_type] > limit: raise ValueError(f"MiniMax-H3 accepts at most {limit} {media_type} references.") - if len(typed) > MINIMAX_H3_MAX_REFERENCES: + if len(typed_references) > MINIMAX_H3_MAX_REFERENCES: raise ValueError(f"MiniMax-H3 accepts at most {MINIMAX_H3_MAX_REFERENCES} references.") - if counts["audio"] == len(typed): + if counts["audio"] == len(typed_references): raise ValueError("Audio references must be paired with at least one image or video reference.") return typed @@ -350,6 +365,8 @@ def prepare_reference( edge and reference-video canvas rule. MiniMax-H3's released values are the defaults; Pika H3 Flex passes its 480p/768p tier values. """ + if reference.media_type == "continuation": + raise ValueError("A continuation clip is decoded by prepare_continuation, not as a reference.") prepared = MiniMaxH3PreparedReference(media_type=reference.media_type) if reference.media_type == "image": source = reference.source @@ -411,19 +428,68 @@ def prepare_reference( return prepared +@dataclass +class MiniMaxH3PreparedContinuation: + """A decoded continuation clip: 24-fps RGB frames on the target canvas, and its soundtrack if it has one.""" + + frames: np.ndarray + waveform: torch.Tensor | None = None + + +def prepare_continuation( + reference: MiniMaxH3Reference, + height: int, + width: int, + target_sample_rate: int, +) -> MiniMaxH3PreparedContinuation: + """Decode a continuation clip, resampled to 24 fps; its frames and canvas must match the contract. + + Like a video reference, the clip is a path (its own soundtrack, if any, is used) or in-memory frames, and + ``soundtrack``/``fps``/``sample_rate`` on the reference override what the file carries. + """ + sound: torch.Tensor | None = None + sound_rate: int | None = None + if isinstance(reference.source, str | os.PathLike): + frames, fps, decoded = decode_reference_video(reference.source) + if decoded is not None: + sound, sound_rate = decoded + else: + frames, fps = reference_media_to_uint8(reference.source), float(MINIMAX_H3_FPS) + if reference.soundtrack is not None: + sound, sound_rate = _resolve_audio_source(reference.soundtrack) + frames = resample_reference_frames(frames, float(reference.fps if reference.fps is not None else fps)) + if frames.shape[0] not in MINIMAX_H3_CONTINUATION_FRAMES: + raise ValueError(f"A continuation clip must have {', '.join(map(str, MINIMAX_H3_CONTINUATION_FRAMES))} frames " + f"at {MINIMAX_H3_FPS} fps, got {frames.shape[0]}.") + if frames.shape[1:3] != (height, width): + raise ValueError(f"A continuation clip must be on the target canvas {width}x{height}, " + f"got {frames.shape[2]}x{frames.shape[1]}.") + waveform = None + if sound is not None: + waveform = prepare_reference_waveform(sound, + int(reference.sample_rate or sound_rate or target_sample_rate), + target_sample_rate, + max_duration=frames.shape[0] / MINIMAX_H3_FPS) + return MiniMaxH3PreparedContinuation(frames=frames, waveform=waveform) + + __all__ = [ "MINIMAX_H3_MAX_REFERENCES", "MINIMAX_H3_MAX_REFERENCE_AUDIOS", "MINIMAX_H3_MAX_REFERENCE_IMAGES", "MINIMAX_H3_MAX_REFERENCE_VIDEOS", + "MINIMAX_H3_CONTINUATION_FRAMES", + "MINIMAX_H3_PREPARED_CONTINUATION_KEY", "MINIMAX_H3_QWEN_TEMPORAL_PATCH", "MINIMAX_H3_QWEN_VIDEO_SAMPLE_FPS", "MINIMAX_H3_REFERENCE_IMAGE_SHORT_EDGE", "coerce_reference", + "MiniMaxH3PreparedContinuation", "MiniMaxH3PreparedReference", "MiniMaxH3Reference", "decode_reference_audio", "decode_reference_video", + "prepare_continuation", "prepare_reference", "prepare_reference_frames", "prepare_reference_image", diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py index 4bccc2284..9b769372d 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_input_preparation.py @@ -25,6 +25,8 @@ from fastvideo.pipelines.basic.minimax_h3.packing import ( from fastvideo.pipelines.basic.minimax_h3.reference import ( MiniMaxH3PreparedReference, prepare_reference, + MINIMAX_H3_PREPARED_CONTINUATION_KEY, + prepare_continuation, validate_references, ) from fastvideo.pipelines.pipeline_batch_info import ForwardBatch @@ -202,8 +204,19 @@ class MiniMaxH3InputPreparationStage(PipelineStage): height, width, ratio = resolve_target_canvas(batch, self.vae, (16, 9)) num_frames = resolve_target_num_frames(batch.num_frames) target_sample_rate = int(self.audio_vae.sampling_rate) - batch.references = [prepare_reference(reference, num_frames, target_sample_rate) for reference in references] + clips = [reference for reference in references if reference.media_type == "continuation"] + batch.references = [ + prepare_reference(reference, num_frames, target_sample_rate) for reference in references + if reference.media_type != "continuation" + ] self._write_target_geometry(batch, height, width, ratio, num_frames) + batch.extra.pop(MINIMAX_H3_PREPARED_CONTINUATION_KEY, None) + if clips: + continuation = prepare_continuation(clips[0], height, width, target_sample_rate) + if continuation.frames.shape[0] >= num_frames: + raise ValueError(f"A {continuation.frames.shape[0]}-frame continuation clip leaves no new frames in " + f"{num_frames}.") + batch.extra[MINIMAX_H3_PREPARED_CONTINUATION_KEY] = continuation batch.extra[MINIMAX_H3_KEYFRAMES_KEY] = [] batch.extra[MINIMAX_H3_KEYFRAME_ANCHORS_KEY] = () diff --git a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py index f684f01dd..d5d55d838 100644 --- a/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py +++ b/fastvideo/pipelines/basic/minimax_h3/stages/minimax_h3_latent_preparation.py @@ -19,6 +19,7 @@ from fastvideo.pipelines.basic.minimax_h3.packing import ( MINIMAX_H3_AUDIO_CHANNELS, MINIMAX_H3_KEYFRAME_ENCODE_SEED, MINIMAX_H3_KEYFRAME_NOISE_AUG, + MiniMaxH3ContinuationLayout, MiniMaxH3PackedLayout, audio_latent_num_frames, build_packed_sequence, @@ -29,6 +30,8 @@ from fastvideo.pipelines.basic.minimax_h3.packing import ( patchify_video_latents, ) from fastvideo.pipelines.basic.minimax_h3.reference import ( + MINIMAX_H3_PREPARED_CONTINUATION_KEY, + MiniMaxH3PreparedContinuation, MiniMaxH3PreparedReference, trim_reference_num_frames, ) @@ -45,6 +48,7 @@ from fastvideo.pipelines.stages.validators import VerificationResult logger = init_logger(__name__) MINIMAX_H3_LAYOUT_KEY = "minimax_h3_layout" +_CONTINUATION_LAYOUT_KEY = "minimax_h3_continuation_layout" def _video_geometry(batch: ForwardBatch) -> tuple[int, int, int, int]: @@ -169,6 +173,19 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): rows.append(patchify_video_latents(latents, patch_size)) return rows + def _encode_continuation_video(self, continuation: MiniMaxH3PreparedContinuation) -> torch.Tensor: + """The clip's normalized video latents, sampled as a reference video's are.""" + pixels = torch.from_numpy(np.ascontiguousarray(continuation.frames)).permute(3, 0, 1, 2)[None] + posterior = self.vae.encode_pixels(pixels).latent_dist + return self.vae.normalize_latents(_sample_visual_posterior(posterior).to(torch.float16).float()).cpu() + + def _encode_continuation_audio(self, continuation: MiniMaxH3PreparedContinuation, + device: torch.device) -> torch.Tensor: + """The clip soundtrack's normalized latents as channel-major rows, encoded as a reference track is.""" + posterior = self.audio_vae.encode(continuation.waveform.to(device)[:, None]).latent_dist + latents = self.audio_vae.normalize_latents(posterior.mode().float()).cpu().transpose(1, 2) + return latents.reshape(-1, self.audio_vae.latent_channels) + def _encode_audio_rows( self, references: list[MiniMaxH3PreparedReference], @@ -243,20 +260,27 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): if not references or not all(isinstance(item, MiniMaxH3PreparedReference) for item in references): raise TypeError("MiniMax-H3 Ref2VA latent preparation requires prepared references.") + continuation = batch.extra.pop(MINIMAX_H3_PREPARED_CONTINUATION_KEY, None) + continuation_video = continuation_audio = None vae_device = get_local_torch_device() pinned_offload.load(self.vae, vae_device, pin=fastvideo_args.pin_cpu_memory) try: video_rows = self._encode_visual_rows(references, vae_device, fastvideo_args) + if continuation is not None: + continuation_video = self._encode_continuation_video(continuation) finally: if fastvideo_args.vae_cpu_offload: pinned_offload.unload(self.vae) audio_rows: list[torch.Tensor] = [] - if any(reference.has_audio for reference in references): + has_continuation_audio = continuation is not None and continuation.waveform is not None + if any(reference.has_audio for reference in references) or has_continuation_audio: audio_device = get_local_torch_device() pinned_offload.load(self.audio_vae, audio_device, pin=fastvideo_args.pin_cpu_memory) try: audio_rows = self._encode_audio_rows(references, audio_device) + if has_continuation_audio: + continuation_audio = self._encode_continuation_audio(continuation, audio_device) finally: if fastvideo_args.vae_cpu_offload: pinned_offload.unload(self.audio_vae) @@ -277,6 +301,31 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): MINIMAX_H3_KEYFRAME_NOISE_AUG, noise, ) + batch.extra.pop(_CONTINUATION_LAYOUT_KEY, None) + if continuation_video is not None: + # The clip's rows follow the references' and carry the same noise augmentation, drawn from a generator of + # its own so the target noise does not depend on whether a clip is present. + patch_size = h3_dit_patch_size(fastvideo_args) + clip_noise = keyframe_condition_noise( + (tuple(int(size) for size in continuation_video.shape[2:]), ), + patch_size, + _video_latent_channels(fastvideo_args), + generator=torch.Generator("cpu").manual_seed(MINIMAX_H3_KEYFRAME_ENCODE_SEED + 1), + device=device, + ) + clip_rows = self.scheduler.scale_noise( + patchify_video_latents(continuation_video, patch_size).to(device), + MINIMAX_H3_KEYFRAME_NOISE_AUG, + clip_noise, + ) + video_conditions = torch.cat((video_conditions, clip_rows)) + if continuation_audio is not None: + audio_rows.append(continuation_audio) + clip_audio_latents = (0 if continuation_audio is None else continuation_audio.shape[0] // + MINIMAX_H3_AUDIO_CHANNELS) + batch.extra[_CONTINUATION_LAYOUT_KEY] = MiniMaxH3ContinuationLayout(num_latent_frames=int( + continuation_video.shape[2]), + num_audio_latents=clip_audio_latents) audio_conditions = torch.cat(audio_rows).to(device) if audio_rows else None for reference in references: @@ -303,6 +352,7 @@ class MiniMaxH3LatentPreparationStage(PipelineStage): width, num_audio_latents, h3_dit_patch_size(fastvideo_args), + continuation=batch.extra.pop(_CONTINUATION_LAYOUT_KEY, None), ) anchors = batch.extra.get(MINIMAX_H3_KEYFRAME_ANCHORS_KEY, ()) if not isinstance(anchors, tuple): diff --git a/fastvideo/tests/entrypoints/test_openai_serving_engine.py b/fastvideo/tests/entrypoints/test_openai_serving_engine.py index ace7dd80b..3f91b5ae5 100644 --- a/fastvideo/tests/entrypoints/test_openai_serving_engine.py +++ b/fastvideo/tests/entrypoints/test_openai_serving_engine.py @@ -411,6 +411,41 @@ def test_fasth3_ref2va_limits_references_at_admission(tmp_path: Path) -> None: ) +def test_fasth3_ref2va_continuation_video_becomes_a_continuation_reference(tmp_path: Path) -> None: + request = VideoGenerationRequest( + prompt="a fox", + task="ref2va", + aspect_ratio="16:9", + num_frames=141, + image_reference=[{"image_url": "/tmp/reference.png"}], + continuation_video="/tmp/clip.mkv", + ) + adapted = build_generation_request( + "video_gen_test", + request, + _args("MiniMaxAI/MiniMax-H3", override_pipeline_cls_name="MiniMaxH3Ref2VAModularPipeline"), + served_model_name="minimax-h3", + output_dir=str(tmp_path), + ) + references = adapted.inputs.references + assert [(reference.media_type, reference.source) for reference in references] == [ + ("image", "/tmp/reference.png"), + ("continuation", "/tmp/clip.mkv"), + ] + + +def test_continuation_video_needs_a_ref2va_server(tmp_path: Path) -> None: + request = VideoGenerationRequest(prompt="a fox", aspect_ratio="16:9", num_frames=141, continuation_video="/tmp/clip.mkv") + with pytest.raises(RequestAdaptationError, match="continuation_video requires"): + build_generation_request( + "video_gen_test", + request, + _args("FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2"), + served_model_name="fasth3", + output_dir=str(tmp_path), + ) + + def test_aspect_ratio_overrides_operator_dimensions(tmp_path: Path) -> None: default_request = parse_config( GenerationRequest, diff --git a/fastvideo/tests/stages/test_minimax_h3_continuation.py b/fastvideo/tests/stages/test_minimax_h3_continuation.py new file mode 100644 index 000000000..9bcbdcaf4 --- /dev/null +++ b/fastvideo/tests/stages/test_minimax_h3_continuation.py @@ -0,0 +1,225 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MiniMax-H3 Ref2VA continuation clips: validation, decoding and their place in the packed sequence.""" + +from types import SimpleNamespace +from unittest.mock import patch + +import numpy as np +from PIL import Image +import pytest +import torch + +from fastvideo.pipelines.basic.minimax_h3 import MiniMaxH3Reference +from fastvideo.pipelines.basic.minimax_h3.packing import ( + MINIMAX_H3_AUDIO_CHANNELS, + MiniMaxH3ContinuationLayout, + _reference_temporal_span, + audio_latent_num_frames, + build_ref2va_packed_sequence, + video_latent_num_frames, +) +from fastvideo.models import pinned_offload +from fastvideo.models.schedulers.scheduling_minimax_h3 import MiniMaxH3Scheduler +from fastvideo.pipelines.basic.minimax_h3.reference import ( + MINIMAX_H3_PREPARED_CONTINUATION_KEY, + prepare_continuation, + validate_references, +) +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_conditioning import MINIMAX_H3_TEXT_TOKEN_TAGS_KEY +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_denoising import _h3_vsa_ref2va_segments +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_input_preparation import MiniMaxH3InputPreparationStage +from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_latent_preparation import MiniMaxH3LatentPreparationStage +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + +TEXT_TOKENS = 137 +LATENT_HEIGHT, LATENT_WIDTH = 48, 84 +ROWS_PER_FRAME = (LATENT_HEIGHT // 2) * (LATENT_WIDTH // 2) + + +def _image_reference(): + return SimpleNamespace(media_type="image", + num_latent_frames=1, + latent_height=64, + latent_width=40, + has_audio=False, + num_audio_latents=0) + + +def _layout(clip_audio_latents: int, num_frames: int = 141, clip_frames: int = 22): + clip = MiniMaxH3ContinuationLayout(video_latent_num_frames(clip_frames), clip_audio_latents) + packed = build_ref2va_packed_sequence( + torch.ones(TEXT_TOKENS, dtype=torch.long), + [_image_reference()], + video_latent_num_frames(num_frames), + LATENT_HEIGHT, + LATENT_WIDTH, + audio_latent_num_frames(num_frames), + (1, 2, 2), + continuation=clip, + ) + return packed, clip + + +@pytest.mark.parametrize("clip_audio_latents", [37, 0]) +def test_clip_rows_take_the_coordinates_of_the_target_start(clip_audio_latents): + packed, clip = _layout(clip_audio_latents) + positions = packed.position_ids + clip_rows = clip.num_latent_frames * ROWS_PER_FRAME + conditions = packed.num_condition_video_rows + clip_video = packed.video_indices[conditions - clip_rows:conditions] + target_start = packed.video_indices[conditions:conditions + clip_rows] + assert torch.equal(positions[clip_video], positions[target_start]) + + image_rows = (64 // 2) * (40 // 2) + assert conditions == image_rows + clip_rows + origin = TEXT_TOKENS + 1.0 + max(float(clip_audio_latents), _reference_temporal_span(clip.num_latent_frames)) + assert positions[target_start[0], 0].item() == pytest.approx(origin) + + audio_conditions = packed.num_condition_audio_rows + assert audio_conditions == clip_audio_latents * MINIMAX_H3_AUDIO_CHANNELS + clip_audio = packed.audio_indices[:audio_conditions] + target_audio = packed.audio_indices[audio_conditions:] + for channel in range(MINIMAX_H3_AUDIO_CHANNELS if clip_audio_latents else 0): + clip_channel = clip_audio[channel * clip_audio_latents:(channel + 1) * clip_audio_latents] + start = channel * packed.num_audio_latents + assert torch.equal(positions[clip_channel], positions[target_audio[start:start + clip_audio_latents]]) + + +@pytest.mark.parametrize("clip_audio_latents", [37, 0]) +def test_vsa_keeps_the_clip_dense(clip_audio_latents): + packed, clip = _layout(clip_audio_latents) + segments = _h3_vsa_ref2va_segments(packed, (1, 2, 2)) + assert [segment for segment in segments if isinstance(segment, tuple)] == [ + (packed.num_video_latent_frames, LATENT_HEIGHT, LATENT_WIDTH) + ] + assert clip.num_latent_frames * ROWS_PER_FRAME in segments + + +def test_a_clip_as_long_as_the_target_is_rejected(): + with pytest.raises(ValueError, match="does not fit"): + _layout(37, num_frames=22, clip_frames=22) + + +def test_a_continuation_is_not_a_reference(): + image = MiniMaxH3Reference(source="picture.png") + clip = MiniMaxH3Reference(source="clip.mkv", media_type="continuation") + assert len(validate_references([image] * 9 + [MiniMaxH3Reference("a.wav", media_type="audio")] * 3 + [clip])) == 13 + with pytest.raises(ValueError, match="at most one continuation"): + validate_references([image, clip, clip]) + with pytest.raises(ValueError, match="besides its clip"): + validate_references([clip]) + + +def test_continuation_clips_are_checked_against_the_contract(): + height, width = 768, 1344 + frames = np.zeros((22, height, width, 3), dtype=np.uint8) + prepared = prepare_continuation(MiniMaxH3Reference(frames, media_type="continuation"), height, width, 32000) + assert prepared.frames.shape == (22, height, width, 3) and prepared.waveform is None + with pytest.raises(ValueError, match="5, 22, 39 frames"): + prepare_continuation(MiniMaxH3Reference(frames[:20], media_type="continuation"), height, width, 32000) + with pytest.raises(ValueError, match="target canvas"): + prepare_continuation(MiniMaxH3Reference(frames, media_type="continuation"), width, height, 32000) + + +class _Posterior: + + def __init__(self, latents: torch.Tensor) -> None: + self.latents = latents + + def sample(self, generator=None) -> torch.Tensor: + return self.latents + + def mode(self) -> torch.Tensor: + return self.latents + + +class _VideoVAE: + """Encodes each clip to latent frames of 2.0 and each still image to 1.0, on the H3 latent geometry.""" + + latent_channels = 4 + spatial_compression_ratio = 16 + + def normalize_pixels(self, pixels): + return pixels + + def normalize_latents(self, latents): + return latents + + def _latents(self, pixels: torch.Tensor, frames: int, value: float): + height, width = pixels.shape[-2] // 16, pixels.shape[-1] // 16 + return SimpleNamespace(latent_dist=_Posterior(torch.full((1, 4, frames, height, width), value))) + + def encode_keyframe(self, pixels): + return self._latents(pixels, 1, 1.0) + + def encode_pixels(self, pixels): + return self._latents(pixels, video_latent_num_frames(pixels.shape[2]), 2.0) + + +class _AudioVAE: + """Encodes each stereo channel to one latent per 800 samples whose values number the channel and step.""" + + latent_channels = 4 + sampling_rate = 32000 + + def normalize_latents(self, latents): + return latents + + def encode(self, waveform): + channels, steps = waveform.shape[0], waveform.shape[-1] // 800 + index = torch.arange(channels * steps, dtype=torch.float32).reshape(channels, 1, steps) + return SimpleNamespace(latent_dist=_Posterior(index.repeat(1, self.latent_channels, 1))) + + +def _args(): + return SimpleNamespace( + pin_cpu_memory=False, + vae_cpu_offload=False, + vae_parallel_encode=False, + pipeline_config=SimpleNamespace( + dit_config=SimpleNamespace(patch_size=(1, 2, 2)), + vae_config=SimpleNamespace(arch_config=SimpleNamespace(latent_channels=4)), + ), + ) + + +def test_a_clip_with_sound_reaches_the_packed_sequence(): + height = width = 64 + clip_frames = np.zeros((22, height, width, 3), dtype=np.uint8) + clip_samples = round(22 / 24 * 32000) + clip = MiniMaxH3Reference(clip_frames, + media_type="continuation", + soundtrack=torch.linspace(-0.5, 0.5, clip_samples).repeat(2, 1), + sample_rate=32000) + image = MiniMaxH3Reference(Image.new("RGB", (64, 64), "red")) + batch = ForwardBatch(data_type="video", + prompt="x", + height=height, + width=width, + num_frames=141, + references=[image, clip], + generator=torch.Generator("cpu").manual_seed(0)) + video, audio = _VideoVAE(), _AudioVAE() + MiniMaxH3InputPreparationStage(video, audio, ref2va=True)._prepare_ref2va(batch) + assert [reference.media_type for reference in batch.references] == ["image"] + prepared = batch.extra[MINIMAX_H3_PREPARED_CONTINUATION_KEY] + assert prepared.frames.shape == (22, height, width, 3) and prepared.waveform.shape == (2, clip_samples) + + stage = MiniMaxH3LatentPreparationStage(video, audio, MiniMaxH3Scheduler(), ref2va=True) + with patch.object(pinned_offload, "load"), patch.object(pinned_offload, "unload"): + video_rows, audio_rows = stage._encode_ref2va_conditions(batch, _args(), torch.device("cpu")) + image_rows = (batch.references[0].latent_height // 2) * (batch.references[0].latent_width // 2) + clip_rows = video_latent_num_frames(22) * (height // 16 // 2) * (width // 16 // 2) + assert video_rows.shape[0] == image_rows + clip_rows + clip_audio_latents = clip_samples // 800 + expected_audio = torch.arange(2 * clip_audio_latents, dtype=torch.float32)[:, None].repeat(1, 4) + assert torch.equal(audio_rows, expected_audio) + + batch.extra[MINIMAX_H3_TEXT_TOKEN_TAGS_KEY] = torch.ones(5, dtype=torch.long) + layout = stage._build_layout(batch, _args()) + assert layout.num_condition_video_rows == image_rows + clip_rows + assert layout.num_condition_audio_rows == 2 * clip_audio_latents + assert layout.reference_segments[-2:] == ( + ("audio", 2 * clip_audio_latents, (0, 0, 0)), + ("continuation", clip_rows, (video_latent_num_frames(22), height // 16, width // 16)), + )