diff --git a/.github/workflows/ci-macos-mlx.yml b/.github/workflows/ci-macos-mlx.yml index a560bb8aa5..62054c8393 100644 --- a/.github/workflows/ci-macos-mlx.yml +++ b/.github/workflows/ci-macos-mlx.yml @@ -86,6 +86,7 @@ jobs: fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py \ + fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py \ fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \ fastvideo/tests/mlx/test_taehv_decode.py \ fastvideo/tests/mlx/test_frame_upsample.py \ @@ -149,6 +150,7 @@ jobs: fastvideo/tests/mlx/test_mlx_minimax_h3_vsa.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_vsa_regressions.py \ fastvideo/tests/mlx/test_mlx_minimax_h3_fast_mode.py \ + fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py \ fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \ fastvideo/tests/mlx/test_taehv_decode.py \ fastvideo/tests/mlx/test_frame_upsample.py \ diff --git a/docs/assets/cookbook-recipes.json b/docs/assets/cookbook-recipes.json index 69e87c9ed6..577cb9a609 100644 --- a/docs/assets/cookbook-recipes.json +++ b/docs/assets/cookbook-recipes.json @@ -476,7 +476,7 @@ }, "evidence": "Verified", "expected_artifact": "MP4 with H.264 video and stereo AAC audio at outputs/fasth3_int6.mp4", - "limitations": ["The MLX path supports T2VA and optional temporal --fast mode. FL2VA, Ref2VA, spatial fast mode, two-pass refinement, and VSA are not wired."] + "limitations": ["The MLX path supports T2VA and optional temporal --fast and spatial --fast-spatial modes. FL2VA, Ref2VA, two-pass refinement, and VSA are not wired."] }, { "id": "minimax-h3-fl2va", diff --git a/docs/cookbook/minimax-h3.md b/docs/cookbook/minimax-h3.md index 9e836102e8..f465208ead 100644 --- a/docs/cookbook/minimax-h3.md +++ b/docs/cookbook/minimax-h3.md @@ -135,7 +135,7 @@ for the download, conversion, and storage requirements. - The full CUDA H3 examples request four GPUs by default. Their sources do not claim a GPU model or memory minimum. - The FastH3 CUDA performance profile was measured on four GB200 GPUs. Use its strict profile when exact operation order matters more than the measured performance configuration. -- The MLX source runtime is limited to T2VA. FL2VA, Ref2VA, VSA, spatial fast mode, and two-pass refinement are not wired on MLX. +- The MLX source runtime is limited to T2VA. FL2VA, Ref2VA, VSA, and two-pass refinement are not wired on MLX. - Gated or missing checkpoints: run `huggingface-cli login` and confirm you accepted the model's license on Hugging Face. ## Evidence status diff --git a/docs/getting_started/installation/mps.md b/docs/getting_started/installation/mps.md index 573da07a73..3a830a8c01 100644 --- a/docs/getting_started/installation/mps.md +++ b/docs/getting_started/installation/mps.md @@ -192,6 +192,25 @@ python examples/inference/basic/mlx_fasth3.py \ --output-path ./outputs/fasth3_int6_fast_720p.mp4 ``` +Add `--fast-spatial` for spatial fast mode, `--fast`'s spatial twin. It +denoises and decodes on the smallest 32px-aligned canvas covering the +requested size divided by `--fast-spatial-scale` (a 480x832 request runs on a +256x416 canvas), then resamples the decoded frames up to the requested size +in pixel space. It composes with `--fast`. This is a speed/quality trade-off +and stays off by default: the output carries the reduced canvas's detail +budget, so it reads softer than a native-resolution render, with the unsharp +pass countering some but not all of the difference: + +```bash +python examples/inference/basic/mlx_fasth3.py \ + --model-root ./FastH3-Preview-v0.2 \ + --mlx-checkpoint ./FastH3-MLX/int6 \ + --prompt "(S1) A presenter says [English] Fast H3 is fastest." \ + --height 480 --width 832 --num-frames 124 --seed 2028 \ + --fast --fast-spatial \ + --output-path ./outputs/fasth3_int6_fast_spatial.mp4 +``` + VSA is off by default. A dense-only checkpoint (no `--include-vsa`) keeps the existing fused-SDPA path. After converting with `--include-vsa`, enable the sparse path explicitly: @@ -212,8 +231,8 @@ falls back to reference on unsupported shapes. It is not the default. `--vsa-impl reference` is the same as `auto`. !!! note "Current MLX scope" - This source runtime supports T2VA, temporal `--fast`, and opt-in VSA. - FL2VA, Ref2VA, spatial fast mode, two-pass refinement, and + This source runtime supports T2VA, temporal `--fast`, spatial + `--fast-spatial`, and opt-in VSA. FL2VA, Ref2VA, two-pass refinement, and `VideoGenerator` registry dispatch are not wired yet. INT8/INT6/INT4 quantization is **weight-only**; VSA attention Q/K/V stay BF16. Old dense MLX checkpoints remain valid for dense inference and raise a reconvert diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 060c05a3bc..4457d6209a 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -185,7 +185,7 @@ optimizations: absence means **untested**, not incompatible. | MLX FastMetal T2V 1.3B | [`FastVideo/FastMetal-1.3B-QAD`](https://huggingface.co/FastVideo/FastMetal-1.3B-QAD) | 480x832, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 16 GB+ unified memory | Released | | MLX FastMetal TI2V 5B | [`FastVideo/FastMetal-5B-QAD`](https://huggingface.co/FastVideo/FastMetal-5B-QAD) | 480p / 720p, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 16 GB+ unified memory | Released | | MLX FastMetal T2V 14B | [`FastVideo/FastMetal-14B-QAD`](https://huggingface.co/FastVideo/FastMetal-14B-QAD) | 480p / 720p, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 36 GB+ unified memory | Released | -| MLX FastH3 Preview T2VA | [`FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2`](https://huggingface.co/FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2) + locally converted DiT | 480p / 720p, 124 frames, 4-step DMD2, INT8/INT6/INT4 **weight-only** DiT, native video + audio VAE; optional temporal RIFE fast mode; optional VSA (tile 64/256, exempt/compete) on `--include-vsa` checkpoints | Apple M4 Max, 36 GB unified memory | Source runtime; T2VA only | +| MLX FastH3 Preview T2VA | [`FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2`](https://huggingface.co/FastVideo/FastVideo-Minimax-FastH3-Preview-v0.2) + locally converted DiT | 480p / 720p, 124 frames, 4-step DMD2, INT8/INT6/INT4 **weight-only** DiT, native video + audio VAE; optional temporal RIFE fast mode; optional spatial fast mode; optional VSA (tile 64/256, exempt/compete) on `--include-vsa` checkpoints | Apple M4 Max, 36 GB unified memory | Source runtime; T2VA only | Apple Silicon uses the native MLX runtime. FastMetal-QAD is the packaged Wan release, while FastH3 Preview currently uses a source checkout and local DiT diff --git a/examples/inference/basic/README.md b/examples/inference/basic/README.md index 9b69e43741..0bad54cbe1 100644 --- a/examples/inference/basic/README.md +++ b/examples/inference/basic/README.md @@ -51,11 +51,12 @@ python examples/inference/basic/mlx_fasth3.py \ --output-path ./outputs/fasth3_int6.mp4 ``` -Pass `--fast` for temporal RIFE fast mode. VSA is opt-in: convert with -`--include-vsa` and pass `--vsa` (see the +Pass `--fast` for temporal RIFE fast mode and `--fast-spatial` for spatial +fast mode (reduced-canvas denoise + pixel-space upsample); the two compose. +VSA is opt-in: convert with `--include-vsa` and pass `--vsa` (see the [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/)). -This MLX entrypoint currently supports T2VA only; FL2VA, Ref2VA, spatial fast -mode, and two-pass refinement remain follow-up work. INT6/INT8/INT4 are +This MLX entrypoint currently supports T2VA only; FL2VA, Ref2VA, and +two-pass refinement remain follow-up work. INT6/INT8/INT4 are weight-only; VSA attention activations stay BF16. Dense-only checkpoints keep working for dense inference. diff --git a/examples/inference/basic/mlx_fasth3.py b/examples/inference/basic/mlx_fasth3.py index ee5493cfc4..d13bdcd03e 100644 --- a/examples/inference/basic/mlx_fasth3.py +++ b/examples/inference/basic/mlx_fasth3.py @@ -19,8 +19,16 @@ requested frame count. A 1280x720 request runs on H3's 1280x736 grid and is center-cropped after decode. +``--fast-spatial`` is spatial fast mode, ``--fast``'s spatial twin. It +denoises and decodes on the smallest 32px-aligned canvas covering +height/width divided by ``--fast-spatial-scale``, then resamples the decoded +frames up to the requested size in pixel space. The two modes compose. +This trades fine detail for speed: the output carries the reduced canvas's +detail budget and reads softer than a native-resolution render, so it stays +off by default. + This entrypoint currently supports text-to-video-with-audio only. It does not -yet wire FL2VA, Ref2VA, spatial fast mode, or two-pass refinement. +yet wire FL2VA, Ref2VA, or two-pass refinement. VSA is off by default; existing dense MLX checkpoints remain supported. H3 uses fused MLX RMSNorm, which can change BF16 rounding compared with the @@ -71,6 +79,22 @@ def parse_args() -> argparse.Namespace: help="temporal reduction target for --fast (default: 2)") parser.add_argument("--fast-sharpen", type=float, default=0.6, help="unsharp strength after RIFE interpolation (0 disables)") + parser.add_argument( + "--fast-spatial", + action=argparse.BooleanOptionalAction, + default=False, + help="denoise and decode at height/width // fast-spatial-scale on H3's 32px grid, " + "then resample the decoded frames up to the requested size; composes with --fast. " + "Trades fine detail for speed", + ) + parser.add_argument("--fast-spatial-scale", type=int, default=2, + help="spatial reduction factor for --fast-spatial (default: 2)") + parser.add_argument("--fast-spatial-upsample-mode", + choices=("lanczos", "cubic", "bilinear", "nearest"), + default="lanczos", + help="pixel interpolation kernel for the post-decode upsample") + parser.add_argument("--fast-spatial-sharpen", type=float, default=0.4, + help="unsharp strength after the upsample (0 disables)") parser.add_argument("--rife-weights-dir", type=Path, default=None, help="optional local mlx-community/RIFE-4.25 snapshot") parser.add_argument("--vae-dtype", choices=("fp32", "fp16", "bf16"), default="fp32") @@ -160,6 +184,10 @@ def main() -> None: fast_factor=args.fast_factor, fast_sharpen=args.fast_sharpen, rife_weights_dir=args.rife_weights_dir, + fast_spatial=args.fast_spatial, + fast_spatial_scale=args.fast_spatial_scale, + fast_spatial_upsample_mode=args.fast_spatial_upsample_mode, + fast_spatial_sharpen=args.fast_spatial_sharpen, vsa=args.vsa, vsa_sparsity=args.vsa_sparsity, vsa_tile_size=args.vsa_tile_size, diff --git a/fastvideo/mlx_runtime/minimax_h3_pipeline.py b/fastvideo/mlx_runtime/minimax_h3_pipeline.py index 12a3af4b6c..4310ee8c3b 100644 --- a/fastvideo/mlx_runtime/minimax_h3_pipeline.py +++ b/fastvideo/mlx_runtime/minimax_h3_pipeline.py @@ -33,6 +33,11 @@ import numpy as np from fastvideo.logger import init_logger +from fastvideo.mlx_runtime.frame_upsample import ( + DEFAULT_PIXEL_UPSAMPLE_MODE, + PIXEL_UPSAMPLE_MODES, + upsample_frames, +) from fastvideo.mlx_runtime.minimax_h3 import ( H3_MANIFEST_FILENAME, MINIMAX_H3_AUDIO_SHIFT, @@ -99,6 +104,72 @@ def plan_fast_temporal(target_frames: int, factor: int = 2) -> FastTemporalPlan: ) +# Resampling from a smaller decode softens output the same way on every +# runtime; 0.4 matches the tuned Wan default without the halos that show +# up by ~0.8. +DEFAULT_FAST_SPATIAL_SHARPEN = 0.4 + + +@dataclass(frozen=True) +class FastSpatialPlan: + """Reduced-canvas geometry for spatial fast mode (RIFE's spatial twin).""" + + target_height: int + target_width: int + stage1_height: int + stage1_width: int + canvas_height: int + canvas_width: int + scale: int + upsample_mode: str + sharpen: float + + +def plan_fast_spatial( + height: int, + width: int, + *, + scale: int = 2, + upsample_mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE, + sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN, +) -> FastSpatialPlan: + """Choose the smallest H3-valid canvas that covers ``target / scale``. + + H3 geometry rounds *up* to the 32px model grid and center-crops after + decode — the same convention plain 720p generation uses via + ``_model_canvas_size`` — so no size the full-resolution path accepts is + rejected here. The return trip to the target size runs in pixel space + after the VAE decode, never on latents; see + :mod:`fastvideo.mlx_runtime.frame_upsample` for why. + """ + if scale < 2: + raise ValueError(f"fast-spatial scale must be at least 2, got {scale}.") + if upsample_mode not in PIXEL_UPSAMPLE_MODES: + raise ValueError(f"Unsupported upsample mode: {upsample_mode!r} " + f"(expected one of {', '.join(PIXEL_UPSAMPLE_MODES)})") + if sharpen < 0: + raise ValueError(f"fast_spatial_sharpen must be non-negative, got {sharpen}.") + target_canvas_height, target_canvas_width = _model_canvas_size(height, width) + stage1_height = math.ceil(height / scale) + stage1_width = math.ceil(width / scale) + canvas_height, canvas_width = _model_canvas_size(stage1_height, stage1_width) + if canvas_height * canvas_width >= target_canvas_height * target_canvas_width: + raise ValueError( + f"fast-spatial scale {scale} does not reduce the H3 canvas for {height}x{width} " + f"(stage-1 canvas {canvas_width}x{canvas_height} vs {target_canvas_width}x{target_canvas_height}).") + return FastSpatialPlan( + target_height=height, + target_width=width, + stage1_height=stage1_height, + stage1_width=stage1_width, + canvas_height=canvas_height, + canvas_width=canvas_width, + scale=scale, + upsample_mode=upsample_mode, + sharpen=sharpen, + ) + + def _model_canvas_size(height: int, width: int) -> tuple[int, int]: """Round an exact output size up to H3's 32-pixel model grid.""" if height <= 0 or width <= 0: @@ -202,10 +273,16 @@ def _validate_checkpoint_step_ladder(checkpoint_dir: str | Path, num_steps: int) f"{num_steps}. Use the step count used during conversion (normally 4), or re-export the checkpoint.") -def _preflight_media_dependencies(*, fast: bool, fast_sharpen: float, rife_weights_dir: str | Path | None) -> None: +def _preflight_media_dependencies(*, + fast: bool, + fast_sharpen: float, + rife_weights_dir: str | Path | None, + fast_spatial: bool = False) -> None: """Fail before conditioning when required output dependencies are unavailable.""" if shutil.which("ffmpeg") is None: raise RuntimeError("ffmpeg is required for MP4 muxing; install it before generation.") + if fast_spatial and importlib.util.find_spec("cv2") is None: + raise RuntimeError("OpenCV is required for --fast-spatial resampling.") if not fast: return if fast_sharpen > 0 and importlib.util.find_spec("cv2") is None: @@ -629,6 +706,10 @@ def generate( fast_factor: int = 2, fast_sharpen: float = 0.6, rife_weights_dir: str | Path | None = None, + fast_spatial: bool = False, + fast_spatial_scale: int = 2, + fast_spatial_upsample_mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE, + fast_spatial_sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN, vsa: bool = False, vsa_sparsity: float = 0.9, vsa_tile_size: int = 64, @@ -654,12 +735,23 @@ def generate( ) if vsa else MiniMaxH3VSAConfig() if vsa_config.enabled and not mlx_h3_checkpoint_vsa_capable(self.dit_checkpoint): raise dense_only_vsa_error(self.dit_checkpoint) + spatial_plan = plan_fast_spatial( + height, + width, + scale=fast_spatial_scale, + upsample_mode=fast_spatial_upsample_mode, + sharpen=fast_spatial_sharpen, + ) if fast_spatial else None _preflight_media_dependencies( fast=fast, fast_sharpen=fast_sharpen, rife_weights_dir=rife_weights_dir, + fast_spatial=fast_spatial, ) - canvas_height, canvas_width = _model_canvas_size(height, width) + if spatial_plan is not None: + canvas_height, canvas_width = spatial_plan.canvas_height, spatial_plan.canvas_width + else: + canvas_height, canvas_width = _model_canvas_size(height, width) target_geometry = self.resolve_geometry(canvas_height, canvas_width, num_frames) fast_plan = plan_fast_temporal(target_geometry["num_frames"], fast_factor) if fast else None video_num_frames = fast_plan.source_frames if fast_plan is not None else target_geometry["num_frames"] @@ -671,7 +763,7 @@ def generate( enforce_duration=fast_plan is None, ) logger.info( - "Geometry: output=%dx%dx%d model=%dx%dx%d audio_frames=%d fast=%s", + "Geometry: output=%dx%dx%d model=%dx%dx%d audio_frames=%d fast=%s fast_spatial=%s", width, height, target_geometry["num_frames"], @@ -680,6 +772,7 @@ def generate( video_geometry["num_frames"], target_geometry["num_frames"], fast_plan, + spatial_plan, ) _reset_peak_memory() @@ -718,7 +811,10 @@ def generate( num_frames=video_geometry["num_frames"], tiled=tiled_video_decode, ) - frames = _center_crop_frames(frames, height, width) + if spatial_plan is not None: + frames = _center_crop_frames(frames, spatial_plan.stage1_height, spatial_plan.stage1_width) + else: + frames = _center_crop_frames(frames, height, width) timings["video_decode_s"] = time.perf_counter() - started peaks["video_decode_gib"] = _peak_memory_gib() _cleanup_mlx() @@ -735,7 +831,8 @@ def generate( target_geometry["num_frames"], model=model, ) - interpolated = _sharpen_frames(interpolated, fast_sharpen) + if spatial_plan is None: + interpolated = _sharpen_frames(interpolated, fast_sharpen) frames = np.stack(interpolated) if frames.shape[0] != target_geometry["num_frames"]: raise RuntimeError( @@ -748,6 +845,22 @@ def generate( del model _cleanup_mlx() + if spatial_plan is not None: + started = time.perf_counter() + # One sharpen pass, at full resolution: RIFE and the resample soften + # for the same reason, so the stronger requested amount is applied + # once instead of stacking two unsharp masks. + sharpen = spatial_plan.sharpen if fast_plan is None else max(spatial_plan.sharpen, fast_sharpen) + frames = np.stack( + upsample_frames( + frames, + width=spatial_plan.target_width, + height=spatial_plan.target_height, + mode=spatial_plan.upsample_mode, + sharpen=sharpen, + )) + timings["spatial_upsample_s"] = time.perf_counter() - started + _reset_peak_memory() started = time.perf_counter() waveform = self.decode_audio(audio_rows, num_frames=target_geometry["num_frames"]) @@ -759,8 +872,8 @@ def generate( video_path = self.mux(frames, waveform, output_path) timings["mux_s"] = time.perf_counter() - started timings["generate_s"] = sum( - timings.get(key, 0.0) - for key in ("condition_s", "denoise_s", "video_decode_s", "rife_s", "audio_decode_s", "mux_s")) + timings.get(key, 0.0) for key in ("condition_s", "denoise_s", "video_decode_s", "rife_s", + "spatial_upsample_s", "audio_decode_s", "mux_s")) result = GenerationResult( video_path=str(video_path), diff --git a/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py new file mode 100644 index 0000000000..66964da533 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_minimax_h3_fast_spatial.py @@ -0,0 +1,251 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Contracts for MiniMax-H3 spatial fast mode on Apple Silicon.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +pytest.importorskip("mlx.core", reason="MLX is required for MiniMax H3 fast-spatial tests") + +from fastvideo.mlx_runtime import rife_interp # noqa: E402 +from fastvideo.mlx_runtime.frame_upsample import upsample_frames # noqa: E402 +from fastvideo.mlx_runtime.minimax_h3 import ( # noqa: E402 + build_packed_layout, + video_latent_num_frames, +) +from fastvideo.mlx_runtime.minimax_h3_pipeline import ( # noqa: E402 + DEFAULT_FAST_SPATIAL_SHARPEN, + MiniMaxH3MLXPipeline, + _center_crop_frames, + _preflight_media_dependencies, + plan_fast_spatial, +) + + +def test_spatial_plan_rounds_stage1_canvas_up_to_model_grid() -> None: + plan = plan_fast_spatial(480, 832) + + assert (plan.target_height, plan.target_width) == (480, 832) + assert (plan.stage1_height, plan.stage1_width) == (240, 416) + assert (plan.canvas_height, plan.canvas_width) == (256, 416) + assert plan.scale == 2 + + +def test_spatial_plan_720p_lands_on_384x640_canvas() -> None: + plan = plan_fast_spatial(720, 1280) + + assert (plan.stage1_height, plan.stage1_width) == (360, 640) + assert (plan.canvas_height, plan.canvas_width) == (384, 640) + + +def test_spatial_plan_rejects_scale_below_two() -> None: + with pytest.raises(ValueError, match="at least 2"): + plan_fast_spatial(480, 832, scale=1) + + +def test_spatial_plan_rejects_non_reducing_scale() -> None: + with pytest.raises(ValueError, match="does not reduce"): + plan_fast_spatial(32, 32, scale=2) + + +def test_spatial_plan_rejects_unknown_upsample_mode() -> None: + with pytest.raises(ValueError, match="Unsupported upsample mode"): + plan_fast_spatial(480, 832, upsample_mode="metalfx") + + +def test_spatial_plan_rejects_negative_sharpen() -> None: + with pytest.raises(ValueError, match="non-negative"): + plan_fast_spatial(480, 832, sharpen=-0.1) + + +def test_spatial_preflight_requires_opencv(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline.shutil.which", lambda _name: "/opt/ffmpeg") + monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline.importlib.util.find_spec", lambda _name: None) + + with pytest.raises(RuntimeError, match="OpenCV is required"): + _preflight_media_dependencies(fast=False, fast_sharpen=0.0, rife_weights_dir=None, fast_spatial=True) + + +def test_spatial_crop_then_upsample_restores_exact_target_size() -> None: + pytest.importorskip("cv2", reason="OpenCV backs the pixel-space resample") + plan = plan_fast_spatial(480, 832) + frames = np.zeros((2, plan.canvas_height, plan.canvas_width, 3), dtype=np.uint8) + + cropped = _center_crop_frames(frames, plan.stage1_height, plan.stage1_width) + assert cropped.shape == (2, 240, 416, 3) + + upsampled = upsample_frames(cropped, width=plan.target_width, height=plan.target_height, + mode="bilinear", sharpen=0.0) + assert np.stack(upsampled).shape == (2, 480, 832, 3) + + +# -- generation-level orchestration contracts (heavyweight phases mocked) ---- + + +def _generate_with_mocked_phases(monkeypatch, tmp_path, **generate_kwargs): + """Run real ``generate()`` orchestration with condition/denoise/decode/mux stubbed.""" + events: list[str] = [] + calls: dict = {} + + pipeline = MiniMaxH3MLXPipeline.__new__(MiniMaxH3MLXPipeline) + pipeline.dit_checkpoint = tmp_path + + monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline._validate_checkpoint_step_ladder", + lambda _checkpoint, _steps: None) + monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline._preflight_media_dependencies", + lambda **_kwargs: None) + monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline.mlx_h3_checkpoint_vsa_capable", + lambda _checkpoint: False) + + def fake_encode_prompt(_prompt): + return np.zeros((8, 8), dtype=np.float32), np.zeros(8, dtype=np.int64) + + def fake_denoise(_text_rows, _token_tags, **kwargs): + events.append("denoise") + calls["denoise"] = kwargs + return np.zeros((4, 4), dtype=np.float32), np.zeros((4, 4), dtype=np.float32) + + def fake_decode_video(_rows, *, height, width, num_frames, tiled): + events.append("decode_video") + calls["decode_video"] = {"height": height, "width": width, "num_frames": num_frames} + return np.zeros((num_frames, height, width, 3), dtype=np.uint8) + + def fake_decode_audio(_rows, *, num_frames): + events.append("decode_audio") + calls["decode_audio"] = {"num_frames": num_frames} + return np.zeros((2, 64), dtype=np.float32) + + pipeline.encode_prompt = fake_encode_prompt + pipeline.denoise = fake_denoise + pipeline.decode_video = fake_decode_video + pipeline.decode_audio = fake_decode_audio + pipeline.mux = lambda _frames, _waveform, output_path: Path(output_path) + + def fake_interpolate(frames, target, *, model): + events.append("rife") + calls["rife"] = {"target": target} + return [np.array(frames[0]) for _ in range(target)] + + def fake_load_model(weights_dir=None): + return object() + + fake_load_model.cache_clear = lambda: None + monkeypatch.setattr(rife_interp, "interpolate_to_frame_count", fake_interpolate) + monkeypatch.setattr(rife_interp, "load_model", fake_load_model) + + def fake_sharpen(frames, amount): + events.append("sharpen") + calls.setdefault("sharpen", []).append(amount) + return list(frames) + + def fake_upsample(frames, *, width, height, mode, sharpen): + events.append("upsample") + calls["upsample"] = {"width": width, "height": height, "mode": mode, "sharpen": sharpen} + return [np.zeros((height, width, 3), dtype=np.uint8) for _ in frames] + + monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline._sharpen_frames", fake_sharpen) + monkeypatch.setattr("fastvideo.mlx_runtime.minimax_h3_pipeline.upsample_frames", fake_upsample) + + result = pipeline.generate( + "(S1) test prompt", + output_path=tmp_path / "out.mp4", + height=480, + width=832, + num_frames=124, + save_frames=True, + **generate_kwargs, + ) + return events, calls, result + + +def test_generate_spatial_only_denoises_reduced_canvas_with_full_audio( + monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + events, calls, result = _generate_with_mocked_phases(monkeypatch, tmp_path, fast_spatial=True) + + assert calls["denoise"]["height"] == 256 + assert calls["denoise"]["width"] == 416 + assert calls["denoise"]["num_frames"] == 124 + assert calls["denoise"]["audio_num_frames"] is None + assert calls["denoise"]["video_temporal_scale"] == 1.0 + assert calls["decode_video"] == {"height": 256, "width": 416, "num_frames": 124} + assert calls["decode_audio"] == {"num_frames": 124} + assert calls["upsample"] == { + "width": 832, "height": 480, "mode": "lanczos", + "sharpen": pytest.approx(DEFAULT_FAST_SPATIAL_SHARPEN), + } + assert "rife" not in events + assert "sharpen" not in events + assert result.frames.shape == (124, 480, 832, 3) + assert "spatial_upsample_s" in result.timings + + +def test_generate_temporal_only_control_keeps_full_canvas(monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + events, calls, result = _generate_with_mocked_phases(monkeypatch, tmp_path, fast=True) + + assert calls["denoise"]["height"] == 480 + assert calls["denoise"]["width"] == 832 + assert calls["denoise"]["num_frames"] == 73 + assert calls["denoise"]["audio_num_frames"] == 124 + assert calls["denoise"]["video_temporal_scale"] > 1.0 + assert calls["decode_video"] == {"height": 480, "width": 832, "num_frames": 73} + assert calls["rife"] == {"target": 124} + assert calls["decode_audio"] == {"num_frames": 124} + assert calls["sharpen"] == [pytest.approx(0.6)] + assert "upsample" not in events + assert result.frames.shape == (124, 480, 832, 3) + + +def test_generate_stacked_runs_rife_before_upsample_with_one_sharpen( + monkeypatch: pytest.MonkeyPatch, tmp_path) -> None: + events, calls, result = _generate_with_mocked_phases(monkeypatch, tmp_path, fast=True, fast_spatial=True) + + assert calls["denoise"]["height"] == 256 + assert calls["denoise"]["width"] == 416 + assert calls["denoise"]["num_frames"] == 73 + assert calls["denoise"]["audio_num_frames"] == 124 + assert calls["decode_video"] == {"height": 256, "width": 416, "num_frames": 73} + assert calls["rife"] == {"target": 124} + assert events.index("rife") < events.index("upsample") + assert "sharpen" not in events + assert calls["upsample"]["sharpen"] == pytest.approx(0.6) + assert calls["decode_audio"] == {"num_frames": 124} + assert result.frames.shape == (124, 480, 832, 3) + + +# -- reduced VSA layout contracts -------------------------------------------- + + +def test_reduced_layout_has_8x13_video_grid_and_unchanged_audio_prefix() -> None: + reduced = build_packed_layout(8, 37, 16, 26, 207) + full = build_packed_layout(8, 37, 30, 52, 207) + + assert reduced.video_indices.shape[0] == 37 * 8 * 13 + assert full.video_indices.shape[0] == 37 * 15 * 26 + np.testing.assert_array_equal(reduced.audio_indices, full.audio_indices) + # A/V sync rides on the audio rows' temporal positions (column 0), which + # must not move with the canvas. Column 2 is excluded on purpose: audio + # rows borrow the video width grid's edge coordinates for their spatial + # tag, so it tracks the canvas the same way any native resolution change + # does. + np.testing.assert_array_equal( + reduced.position_ids[reduced.audio_indices, 0], + full.position_ids[full.audio_indices, 0], + ) + np.testing.assert_array_equal(reduced.position_ids[reduced.audio_indices, 1], + np.zeros(reduced.audio_indices.shape[0])) + + +def test_reduced_layout_stacked_uses_22_latent_frames() -> None: + assert video_latent_num_frames(73) == 22 + stacked = build_packed_layout(8, 22, 16, 26, 207, video_temporal_scale=1.7) + baseline = build_packed_layout(8, 22, 16, 26, 207) + + assert stacked.video_indices.shape[0] == 22 * 8 * 13 + np.testing.assert_array_equal(stacked.audio_indices, baseline.audio_indices) + np.testing.assert_array_equal( + stacked.position_ids[stacked.audio_indices], + baseline.position_ids[baseline.audio_indices], + )