[bugfix] MiniMax H3: drop the audio/video decode overlap; tests use envs overrides; regenerate env docs
This commit is contained in:
Aryan Kumar
@@ -268,7 +268,7 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
|
||||
| `FASTVIDEO_H3_FP8_ATTENTION` | bool | `0` | performance | With NVFP4 layer_profile h3_dit_ffn, run MiniMax-H3 attention projections in FP8. |
|
||||
| `FASTVIDEO_H3_FP8_GRANULARITY` | one of tensor, channel | `tensor` | performance | FP8 scaling granularity for FASTVIDEO_H3_FP8_ATTENTION. |
|
||||
| `FASTVIDEO_LORA_SIDE_BRANCH` | bool | `0` | performance | Apply inference LoRA adapters as a bf16 side branch next to every target linear instead of merging them, so the adapter strength can change at runtime. Packed NVFP4 or FP8 linears without their original dense weight always use the side branch. |
|
||||
| `FASTVIDEO_NVFP4_MM_BACKEND` | str | `auto` | performance | FlashInfer mm_fp4 backend for NVFP4 linears, e.g. auto or cutlass. |
|
||||
| `FASTVIDEO_NVFP4_MM_BACKEND` | str | `auto` | performance | FlashInfer mm_fp4 backend for NVFP4 linears: auto (cudnn for MiniMax-H3 DiT profiles on sm_100), cutlass or cudnn. |
|
||||
| `FASTVIDEO_NVFP4_MM_ROW_SPLIT` | one of off, table, auto | `off` | performance | Run large cuDNN NVFP4 GEMMs over row chunks where its plan heuristic is slow: off, table (fixed per-shape chunk counts) or auto (time 1/2/4 chunks once per shape; not deterministic across runs). |
|
||||
| `FASTVIDEO_NVFP4_ACT_AMAX` | path | unset | performance | JSON of calibrated NVFP4 input amax per linear, keyed b<block>.<sub> or full prefix; sets a static activation scale. |
|
||||
| `FASTVIDEO_NVFP4_DYNAMIC_ACT` | str | `""` | performance | NVFP4 linears that derive the activation scale per call: all, or comma-separated layer-name suffixes such as ff.fc_out. |
|
||||
@@ -288,6 +288,9 @@ longer exists also fails the test, so the fixing pull request deletes its entry.
|
||||
| `FASTVIDEO_H3_VAE_INT8_OVERLAY` | bool | `1` | performance | Apply the INT8 ConvRot decoder overlay when the MiniMax-H3 VAE folder ships one; 0 decodes with the dense weights. |
|
||||
| `FASTVIDEO_H3_VAE_INT8_SHARED_QKV` | bool | `0` | performance | Share the INT8 activation rotation and quantization across the MiniMax-H3 VAE Q/K/V projections. |
|
||||
| `FASTVIDEO_H3_VAE_INT8_TRANSPOSE_VIEW` | bool | `0` | performance | Use transposed weight views in the MiniMax-H3 VAE INT8 projections. |
|
||||
| `FASTVIDEO_H3_VAE_COMPILE_BLOCKS` | bool | `0` | performance | Run the MiniMax-H3 video VAE decoder transformer blocks through one regional torch.compile graph: Inductor fuses the fp32 residual scale/add pairs, the RMSNorm+cast and the SwiGLU silu\*mul. |
|
||||
| `FASTVIDEO_H3_VAE_UINT8_PIXELS` | bool | `1` | performance | Quantize single-GPU MiniMax-H3 video VAE output to uint8 on the GPU per chunk and stream uint8 into the host buffer; frames are bitwise identical to the fp32 path. 0 streams fp32. |
|
||||
| `FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH` | int | `1` | performance | Spatial tiles per MiniMax-H3 video VAE encoder call (reference and keyframe encodes); 1 encodes per tile. Batched tiles differ from per-tile encoding only in reduction order. |
|
||||
| `FASTVIDEO_H3_VAE_INT8_FUSED_DEQUANT` | bool | `0` | performance | Fused dequantization epilogue for the MiniMax-H3 VAE INT8 projections. |
|
||||
| `FASTVIDEO_H3_PINNED_SWAP` | bool | `1` | performance | Swap offloaded MiniMax-H3 modules through exact-size pinned host arenas. |
|
||||
| `FASTVIDEO_H3_PARK_MODULES` | str | unset | performance | Comma-separated MiniMax-H3 denoise modules parked on the host while the text encoder runs, e.g. vae,audio_vae. |
|
||||
|
||||
@@ -941,9 +941,6 @@ class VideoGenerator:
|
||||
# *do* ask for the latent samples via ``return_frames=True``.
|
||||
# ``skip_pixel_prealloc`` also gates the slow-path warning.
|
||||
needs_samples_out = batch.return_frames or batch.return_samples
|
||||
if (needs_samples_out and not batch.return_samples and not is_latent_output
|
||||
and envs.FASTVIDEO_H3_SKIP_FP32_SAMPLES_MIRROR.get()):
|
||||
needs_samples_out = False
|
||||
# ``output_shape_from_input`` pipelines (Kandinsky6 SR) decide the output geometry from the input video, so the
|
||||
# request's height / width / num_frames say nothing about the buffer to pre-allocate.
|
||||
output_shape_from_input = bool(getattr(fastvideo_args.pipeline_config, "output_shape_from_input", False))
|
||||
|
||||
@@ -766,16 +766,6 @@ FASTVIDEO_H3_VAE_COMPILE_BLOCKS = EnvBool(False,
|
||||
doc="Run the MiniMax-H3 video VAE decoder transformer blocks through one "
|
||||
"regional torch.compile graph: Inductor fuses the fp32 residual scale/add "
|
||||
"pairs, the RMSNorm+cast and the SwiGLU silu*mul.")
|
||||
FASTVIDEO_H3_SKIP_FP32_SAMPLES_MIRROR = EnvBool(
|
||||
False,
|
||||
category="performance",
|
||||
doc="Do not allocate or fill the fp32 [0, 1] ``samples`` mirror when a request asks for frames but not "
|
||||
"samples and the decode already returned uint8 pixels (``result.samples`` is then None).")
|
||||
FASTVIDEO_H3_OVERLAP_AUDIO_DECODE = EnvBool(
|
||||
True,
|
||||
category="performance",
|
||||
doc="Decode the MiniMax-H3 audio latents on a second CUDA stream while the video VAE decodes, instead of "
|
||||
"serially after it; the waveform is bitwise identical. 0 decodes serially.")
|
||||
FASTVIDEO_H3_VAE_UINT8_PIXELS = EnvBool(True,
|
||||
category="performance",
|
||||
doc="Quantize single-GPU MiniMax-H3 video VAE output to uint8 on the GPU per "
|
||||
|
||||
@@ -3,11 +3,9 @@
|
||||
|
||||
from __future__ import annotations
|
||||
|
||||
import copy
|
||||
import gc
|
||||
import json
|
||||
import math
|
||||
import threading
|
||||
from dataclasses import dataclass
|
||||
from pathlib import Path
|
||||
from typing import Any
|
||||
@@ -655,16 +653,14 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
self._restore_denoise_modules(fastvideo_args)
|
||||
if not self._denoise_stages_ready:
|
||||
self._add_denoise_stages(ref2va=self._ref2va)
|
||||
return self._run_stage_sequence(
|
||||
(
|
||||
for name in (
|
||||
"latent_preparation_stage",
|
||||
"denoising_stage",
|
||||
"video_decoding_stage",
|
||||
"audio_decoding_stage",
|
||||
),
|
||||
batch,
|
||||
fastvideo_args,
|
||||
)
|
||||
):
|
||||
batch = self._stage_name_mapping[name](batch, fastvideo_args)
|
||||
return batch
|
||||
|
||||
def _input_video_geometry(self, fastvideo_args: FastVideoArgs) -> Any:
|
||||
"""Read canvas scalars from checkpoint JSON, not a live VAE proxy."""
|
||||
@@ -764,53 +760,6 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
if self._denoise_modules_loaded():
|
||||
self._add_denoise_stages(ref2va=ref2va)
|
||||
|
||||
def _run_stage_sequence(self, names: tuple[str, ...], batch: ForwardBatch,
|
||||
fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
"""Run named stages in order; with ``FASTVIDEO_H3_OVERLAP_AUDIO_DECODE`` on CUDA, decode audio on a side
|
||||
stream while the video VAE decodes."""
|
||||
stages = [self._stage_name_mapping[name] for name in names]
|
||||
audio_stage = next((stage for stage in stages if isinstance(stage, MiniMaxH3AudioDecodingStage)), None)
|
||||
video_stage = next((stage for stage in stages if isinstance(stage, MiniMaxH3VideoDecodingStage)), None)
|
||||
if (not envs.FASTVIDEO_H3_OVERLAP_AUDIO_DECODE.get() or not torch.cuda.is_available() or audio_stage is None
|
||||
or video_stage is None or stages.index(audio_stage) < stages.index(video_stage)):
|
||||
for stage in stages:
|
||||
batch = stage(batch, fastvideo_args)
|
||||
return batch
|
||||
|
||||
box: dict[str, Any] = {}
|
||||
stream = torch.cuda.Stream(device=torch.cuda.current_device())
|
||||
|
||||
def run_audio() -> None:
|
||||
local = copy.copy(batch)
|
||||
local.extra = dict(batch.extra)
|
||||
try:
|
||||
with torch.cuda.stream(stream):
|
||||
local = audio_stage.forward(local, fastvideo_args)
|
||||
stream.synchronize()
|
||||
box["audio"] = local.extra.get("audio")
|
||||
box["audio_sample_rate"] = local.extra.get("audio_sample_rate")
|
||||
except BaseException as error: # noqa: BLE001 - re-raised on the main thread
|
||||
box["error"] = error
|
||||
|
||||
thread: threading.Thread | None = None
|
||||
for stage in stages:
|
||||
if stage is audio_stage:
|
||||
continue
|
||||
if stage is video_stage:
|
||||
thread = threading.Thread(target=run_audio, name="minimax-h3-audio-decode", daemon=True)
|
||||
thread.start()
|
||||
batch = stage(batch, fastvideo_args)
|
||||
if thread is not None:
|
||||
thread.join()
|
||||
if "error" in box:
|
||||
raise box["error"]
|
||||
batch.extra["audio"] = box.get("audio")
|
||||
batch.extra["audio_sample_rate"] = box.get("audio_sample_rate")
|
||||
# forward() bypassed PipelineStage.__call__, which normally releases these.
|
||||
audio_stage._release_deferred_modules(f"{getattr(audio_stage, '_pipeline_stage_name', 'audio')}|"
|
||||
f"{type(audio_stage).__name__}")
|
||||
return batch
|
||||
|
||||
def _forward_encoder_split(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
encoder_worker = h3_is_encoder_worker(fastvideo_args)
|
||||
batch = self._stage_name_mapping["input_preparation_stage"](batch, fastvideo_args)
|
||||
@@ -834,16 +783,14 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
# gigabytes for a response the executor never reads.
|
||||
return ForwardBatch(data_type=batch.data_type)
|
||||
h3_receive_condition(batch, fastvideo_args)
|
||||
return self._run_stage_sequence(
|
||||
(
|
||||
for name in (
|
||||
"latent_preparation_stage",
|
||||
"denoising_stage",
|
||||
"video_decoding_stage",
|
||||
"audio_decoding_stage",
|
||||
),
|
||||
batch,
|
||||
fastvideo_args,
|
||||
)
|
||||
):
|
||||
batch = self._stage_name_mapping[name](batch, fastvideo_args)
|
||||
return batch
|
||||
|
||||
def forward(self, batch: ForwardBatch, fastvideo_args: FastVideoArgs) -> ForwardBatch:
|
||||
if not self.post_init_called:
|
||||
@@ -873,13 +820,7 @@ class MiniMaxH3BasePipeline(LoRAPipeline, ComposedPipelineBase):
|
||||
except BaseException:
|
||||
self._release_all_lazy_modules()
|
||||
raise
|
||||
if not envs.FASTVIDEO_H3_OVERLAP_AUDIO_DECODE.get():
|
||||
return super().forward(batch, fastvideo_args)
|
||||
try:
|
||||
return self._run_stage_sequence(tuple(self._stage_name_mapping), batch, fastvideo_args)
|
||||
except BaseException:
|
||||
self._release_all_lazy_modules()
|
||||
raise
|
||||
return super().forward(batch, fastvideo_args)
|
||||
|
||||
|
||||
class MiniMaxH3Pipeline(MiniMaxH3BasePipeline):
|
||||
|
||||
@@ -1,4 +1,5 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import fastvideo.envs as envs
|
||||
import pytest
|
||||
|
||||
from fastvideo.layers.quantization import nvfp4_config
|
||||
@@ -12,13 +13,13 @@ from fastvideo.layers.quantization import nvfp4_config
|
||||
(None, True, "auto"),
|
||||
("h3_dit", False, "auto"),
|
||||
])
|
||||
def test_auto_backend_resolution(monkeypatch, profile, sm100, expected) -> None:
|
||||
monkeypatch.delenv("FASTVIDEO_NVFP4_MM_BACKEND", raising=False)
|
||||
def test_auto_backend_resolution(monkeypatch, profile, sm100, expected, env_overrides) -> None:
|
||||
env_overrides.enter_context(envs.FASTVIDEO_NVFP4_MM_BACKEND.override(None))
|
||||
monkeypatch.setattr(nvfp4_config, "_is_sm100", lambda: sm100)
|
||||
assert nvfp4_config._mm_fp4_backend(profile) == expected
|
||||
|
||||
|
||||
def test_explicit_backend_wins(monkeypatch) -> None:
|
||||
monkeypatch.setenv("FASTVIDEO_NVFP4_MM_BACKEND", "cutlass")
|
||||
def test_explicit_backend_wins(monkeypatch, env_overrides) -> None:
|
||||
env_overrides.enter_context(envs.FASTVIDEO_NVFP4_MM_BACKEND.override("cutlass"))
|
||||
monkeypatch.setattr(nvfp4_config, "_is_sm100", lambda: True)
|
||||
assert nvfp4_config._mm_fp4_backend("h3_dit") == "cutlass"
|
||||
|
||||
@@ -1,148 +0,0 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
"""FASTVIDEO_H3_OVERLAP_AUDIO_DECODE: audio decode beside the video decode."""
|
||||
|
||||
import time
|
||||
|
||||
import pytest
|
||||
import torch
|
||||
|
||||
from fastvideo.pipelines.basic.minimax_h3.minimax_h3_pipeline import MiniMaxH3Pipeline
|
||||
from fastvideo.pipelines.basic.minimax_h3.stages.minimax_h3_decoding import (
|
||||
MiniMaxH3AudioDecodingStage,
|
||||
MiniMaxH3VideoDecodingStage,
|
||||
)
|
||||
from fastvideo.pipelines.pipeline_batch_info import ForwardBatch
|
||||
|
||||
|
||||
class _FakeStream:
|
||||
|
||||
device = torch.device("cpu")
|
||||
|
||||
def __enter__(self):
|
||||
return self
|
||||
|
||||
def __exit__(self, *exc_info):
|
||||
return False
|
||||
|
||||
def synchronize(self) -> None:
|
||||
pass
|
||||
|
||||
|
||||
@pytest.fixture
|
||||
def fake_cuda(monkeypatch):
|
||||
"""The overlap path only needs a stream-like context manager and a device index."""
|
||||
monkeypatch.setattr(torch.cuda, "Stream", lambda *args, **kwargs: _FakeStream())
|
||||
monkeypatch.setattr(torch.cuda, "stream", lambda *args, **kwargs: _FakeStream())
|
||||
monkeypatch.setattr(torch.cuda, "current_device", lambda: 0)
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: True)
|
||||
|
||||
|
||||
def _pipeline(video: MiniMaxH3VideoDecodingStage, audio: MiniMaxH3AudioDecodingStage) -> MiniMaxH3Pipeline:
|
||||
pipeline = object.__new__(MiniMaxH3Pipeline)
|
||||
pipeline._stage_name_mapping = {"video_decoding_stage": video, "audio_decoding_stage": audio}
|
||||
return pipeline
|
||||
|
||||
|
||||
def _fake_video(order: list[str], hold: float):
|
||||
video = MiniMaxH3VideoDecodingStage(None)
|
||||
|
||||
def forward(batch, fastvideo_args):
|
||||
order.append("video_start")
|
||||
batch.extra["video"] = torch.ones(1)
|
||||
time.sleep(hold)
|
||||
order.append("video_end")
|
||||
return batch
|
||||
|
||||
video.forward = forward
|
||||
return video
|
||||
|
||||
|
||||
def _fake_audio(order: list[str], hold: float, *, fail: bool = False):
|
||||
audio = MiniMaxH3AudioDecodingStage(None)
|
||||
|
||||
def forward(batch, fastvideo_args):
|
||||
order.append("audio_start")
|
||||
time.sleep(hold)
|
||||
if fail:
|
||||
raise RuntimeError("audio decode failed")
|
||||
batch.extra["audio"] = torch.full((2, 3), 0.5)
|
||||
batch.extra["audio_sample_rate"] = 16000
|
||||
# The real stage clears the runtime tensors it consumed.
|
||||
batch.latents = None
|
||||
batch.extra.pop("minimax_h3_layout", None)
|
||||
order.append("audio_end")
|
||||
return batch
|
||||
|
||||
audio.forward = forward
|
||||
return audio
|
||||
|
||||
|
||||
def test_audio_decode_overlaps_the_video_decode(monkeypatch, fake_cuda) -> None:
|
||||
monkeypatch.setenv("FASTVIDEO_H3_OVERLAP_AUDIO_DECODE", "1")
|
||||
order: list[str] = []
|
||||
video = _fake_video(order, hold=0.20)
|
||||
audio = _fake_audio(order, hold=0.10)
|
||||
batch = ForwardBatch(data_type="video", latents=torch.zeros(1), extra={"minimax_h3_layout": object()})
|
||||
|
||||
result = _pipeline(video, audio)._run_stage_sequence(("video_decoding_stage", "audio_decoding_stage"), batch,
|
||||
None)
|
||||
|
||||
# The audio decode started before the video decode finished: they overlap.
|
||||
assert order.index("audio_start") < order.index("video_end")
|
||||
assert order.index("audio_end") < order.index("video_end")
|
||||
# The results merge back onto the caller's batch; the audio thread's copy did
|
||||
# not clear the caller's runtime tensors.
|
||||
assert torch.equal(result.extra["audio"], torch.full((2, 3), 0.5))
|
||||
assert result.extra["audio_sample_rate"] == 16000
|
||||
assert result.extra["video"] is not None
|
||||
assert result.latents is not None
|
||||
assert "minimax_h3_layout" in result.extra
|
||||
|
||||
|
||||
def test_audio_decode_stays_serial_when_disabled(monkeypatch, fake_cuda) -> None:
|
||||
monkeypatch.setenv("FASTVIDEO_H3_OVERLAP_AUDIO_DECODE", "0")
|
||||
order: list[str] = []
|
||||
video = _fake_video(order, hold=0.05)
|
||||
audio = _fake_audio(order, hold=0.05)
|
||||
batch = ForwardBatch(data_type="video", latents=torch.zeros(1), extra={})
|
||||
|
||||
_pipeline(video, audio)._run_stage_sequence(("video_decoding_stage", "audio_decoding_stage"), batch, None)
|
||||
|
||||
assert order == ["video_start", "video_end", "audio_start", "audio_end"]
|
||||
|
||||
|
||||
def test_audio_decode_error_is_raised_on_the_main_thread(monkeypatch, fake_cuda) -> None:
|
||||
monkeypatch.setenv("FASTVIDEO_H3_OVERLAP_AUDIO_DECODE", "1")
|
||||
order: list[str] = []
|
||||
video = _fake_video(order, hold=0.05)
|
||||
audio = _fake_audio(order, hold=0.01, fail=True)
|
||||
batch = ForwardBatch(data_type="video", latents=torch.zeros(1), extra={})
|
||||
|
||||
with pytest.raises(RuntimeError, match="audio decode failed"):
|
||||
_pipeline(video, audio)._run_stage_sequence(("video_decoding_stage", "audio_decoding_stage"), batch, None)
|
||||
|
||||
|
||||
def test_audio_decode_runs_after_video_when_listed_first(monkeypatch, fake_cuda) -> None:
|
||||
"""An audio stage that precedes the video stage in the order cannot overlap it."""
|
||||
monkeypatch.setenv("FASTVIDEO_H3_OVERLAP_AUDIO_DECODE", "1")
|
||||
order: list[str] = []
|
||||
video = _fake_video(order, hold=0.05)
|
||||
audio = _fake_audio(order, hold=0.05)
|
||||
batch = ForwardBatch(data_type="video", latents=torch.zeros(1), extra={})
|
||||
|
||||
_pipeline(video, audio)._run_stage_sequence(("audio_decoding_stage", "video_decoding_stage"), batch, None)
|
||||
|
||||
assert order == ["audio_start", "audio_end", "video_start", "video_end"]
|
||||
|
||||
|
||||
def test_audio_decode_stays_serial_without_cuda(monkeypatch) -> None:
|
||||
monkeypatch.setenv("FASTVIDEO_H3_OVERLAP_AUDIO_DECODE", "1")
|
||||
monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
|
||||
order: list[str] = []
|
||||
video = _fake_video(order, hold=0.05)
|
||||
audio = _fake_audio(order, hold=0.05)
|
||||
batch = ForwardBatch(data_type="video", latents=torch.zeros(1), extra={})
|
||||
|
||||
_pipeline(video, audio)._run_stage_sequence(("video_decoding_stage", "audio_decoding_stage"), batch, None)
|
||||
|
||||
assert order == ["video_start", "video_end", "audio_start", "audio_end"]
|
||||
@@ -4,6 +4,7 @@ from types import SimpleNamespace
|
||||
|
||||
import numpy as np
|
||||
import pytest
|
||||
import fastvideo.envs as envs
|
||||
import torch
|
||||
|
||||
from fastvideo.pipelines.basic.minimax_h3 import encoder_split
|
||||
@@ -74,8 +75,8 @@ def test_reference_video_encode_keeps_pixels_on_cpu() -> None:
|
||||
assert rows[0].shape == (7 * 4 * 4, 4)
|
||||
|
||||
|
||||
def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None:
|
||||
monkeypatch.setenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", "0")
|
||||
def test_decode_stage_uses_cpu_output_buffer(monkeypatch, env_overrides) -> None:
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.override(False))
|
||||
latent_shape = (1, 4, 2, 4, 4)
|
||||
latents = torch.randn(latent_shape)
|
||||
rows = patchify_video_latents(latents, (1, 1, 1))
|
||||
@@ -129,8 +130,8 @@ def test_decode_stage_uses_cpu_output_buffer(monkeypatch) -> None:
|
||||
assert torch.equal(result.output, expected)
|
||||
|
||||
|
||||
def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None:
|
||||
monkeypatch.setenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", "0")
|
||||
def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch, env_overrides) -> None:
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.override(False))
|
||||
latent_shape = (1, 4, 2, 4, 4)
|
||||
latents = torch.randn(latent_shape)
|
||||
rows = patchify_video_latents(latents, (1, 1, 1))
|
||||
@@ -179,7 +180,7 @@ def test_decode_stage_requests_pin_fallback_output_buffer(monkeypatch) -> None:
|
||||
}
|
||||
|
||||
|
||||
def test_decode_stage_streams_uint8_by_default(monkeypatch) -> None:
|
||||
def test_decode_stage_streams_uint8_by_default(monkeypatch, env_overrides) -> None:
|
||||
latent_shape = (1, 4, 2, 4, 4)
|
||||
rows = patchify_video_latents(torch.randn(latent_shape), (1, 1, 1))
|
||||
batch = ForwardBatch(data_type="video", latents=rows, raw_latent_shape=latent_shape)
|
||||
@@ -201,7 +202,7 @@ def test_decode_stage_streams_uint8_by_default(monkeypatch) -> None:
|
||||
observed["dtype"] = output.dtype
|
||||
output.fill_(63)
|
||||
|
||||
monkeypatch.delenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", raising=False)
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.override(None))
|
||||
monkeypatch.setattr(minimax_h3_decoding, "get_local_torch_device", lambda: torch.device("cpu"))
|
||||
result = MiniMaxH3VideoDecodingStage(VAE()).forward(
|
||||
batch,
|
||||
@@ -221,10 +222,10 @@ def test_decode_stage_streams_uint8_by_default(monkeypatch) -> None:
|
||||
@pytest.mark.parametrize("uint8_pixels", [False, True])
|
||||
@pytest.mark.parametrize("unified_memory", [True, False])
|
||||
def test_decode_stage_reuses_its_pixel_buffer_only_on_unified_memory(monkeypatch, unified_memory,
|
||||
uint8_pixels) -> None:
|
||||
uint8_pixels, env_overrides) -> None:
|
||||
"""On unified memory, same-geometry requests reuse one pre-touched CPU buffer; returned
|
||||
frames never alias it. Discrete GPUs allocate per request and hold nothing between them."""
|
||||
monkeypatch.setenv("FASTVIDEO_H3_VAE_UINT8_PIXELS", "1" if uint8_pixels else "0")
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_UINT8_PIXELS.override(uint8_pixels))
|
||||
latent_shape = (1, 4, 2, 4, 4)
|
||||
rows = patchify_video_latents(torch.randn(latent_shape), (1, 1, 1))
|
||||
allocations = []
|
||||
|
||||
@@ -11,6 +11,7 @@ from typing import Any
|
||||
from unittest.mock import Mock, patch
|
||||
|
||||
import pytest
|
||||
import fastvideo.envs as envs
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
@@ -551,9 +552,9 @@ def test_decoder_forward_uses_the_compiled_block_region() -> None:
|
||||
assert torch.equal(actual, expected)
|
||||
|
||||
|
||||
def test_regional_block_compile_flag_installs_on_construction(monkeypatch) -> None:
|
||||
def test_regional_block_compile_flag_installs_on_construction(monkeypatch, env_overrides) -> None:
|
||||
"""The env flag is read when the decoder is built, like the other VAE load-time knobs."""
|
||||
monkeypatch.setenv("FASTVIDEO_H3_VAE_COMPILE_BLOCKS", "1")
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_COMPILE_BLOCKS.override(True))
|
||||
compiled_block = Mock(name="compiled_block")
|
||||
with patch("fastvideo.models.vaes.minimax_h3_video.torch.compile", return_value=compiled_block) as compile_mock:
|
||||
decoder = _tiny_h3_decoder()
|
||||
|
||||
@@ -1,5 +1,6 @@
|
||||
# SPDX-License-Identifier: Apache-2.0
|
||||
import pytest
|
||||
import fastvideo.envs as envs
|
||||
import torch
|
||||
from torch.testing import assert_close
|
||||
|
||||
@@ -124,7 +125,7 @@ def test_decode_to_pixels_uint8_matches_fp32_quantization(device: str) -> None:
|
||||
|
||||
|
||||
@torch.inference_mode()
|
||||
def test_batched_tile_encode_matches_per_tile(monkeypatch) -> None:
|
||||
def test_batched_tile_encode_matches_per_tile(monkeypatch, env_overrides) -> None:
|
||||
torch.manual_seed(20260810)
|
||||
vae = _tiny_vae()
|
||||
vae.enable_tiling(tile_sample_min_height=16,
|
||||
@@ -133,9 +134,9 @@ def test_batched_tile_encode_matches_per_tile(monkeypatch) -> None:
|
||||
tile_sample_min_overlap_width=8)
|
||||
pixels = torch.rand(1, 3, 1, 32, 32)
|
||||
|
||||
monkeypatch.setenv("FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH", "1")
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH.override(1))
|
||||
per_tile = vae.encode(vae.normalize_pixels(pixels)).latent_dist.parameters
|
||||
monkeypatch.setenv("FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH", "8")
|
||||
env_overrides.enter_context(envs.FASTVIDEO_H3_VAE_ENCODE_TILE_BATCH.override(8))
|
||||
batched = vae.encode(vae.normalize_pixels(pixels)).latent_dist.parameters
|
||||
|
||||
assert_close(batched, per_tile, atol=1e-5, rtol=1e-5)
|
||||
|
||||
Reference in New Issue
Block a user