From 8537dcd6ded64052327c4c8e2d7acdfa903c845e Mon Sep 17 00:00:00 2001 From: Aryan Kumar <156166681+aryan5v@users.noreply.github.com> Date: Tue, 18 Aug 2026 15:08:00 -0700 Subject: [PATCH] =?UTF-8?q?[feat]:=20Apple=20Silicon=20MLX=20runtime=20?= =?UTF-8?q?=E2=80=94=20INT8=20Wan2.1=20and=20Wan2.2=20inference=20(#1638)?= MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Co-authored-by: Aryan Kumar --- .github/workflows/ci-macos-mlx.yml | 149 +++ README.md | 5 + docs/design/apple_silicon_fast_mode.md | 128 ++ docs/getting_started/installation/mps.md | 10 +- docs/inference/support_matrix.md | 11 + .../inference/basic/mlx_rife_smoke_test.py | 46 + .../inference/basic/mlx_wan22_generate.py | 510 ++++++++ .../basic/mlx_wan_decode_benchmark.py | 130 ++ .../basic/mlx_wan_prompt_to_video.py | 1114 +++++++++++++++++ .../basic/mlx_wan_quant_benchmark.py | 365 ++++++ .../inference/basic/mlx_wan_video_quality.py | 104 ++ fastvideo/benchmarks/__init__.py | 1 + fastvideo/benchmarks/eval_metalfx_rife.py | 375 ++++++ fastvideo/benchmarks/mlx_fastwan_bench.py | 826 ++++++++++++ fastvideo/mlx_runtime/__init__.py | 115 ++ fastvideo/mlx_runtime/checkpoint.py | 273 ++++ fastvideo/mlx_runtime/fast_spatial.py | 228 ++++ fastvideo/mlx_runtime/fastwan.py | 980 +++++++++++++++ fastvideo/mlx_runtime/frame_upsample.py | 154 +++ fastvideo/mlx_runtime/memory.py | 243 ++++ fastvideo/mlx_runtime/prompt_cache.py | 129 ++ fastvideo/mlx_runtime/prompt_enhance.py | 454 +++++++ fastvideo/mlx_runtime/quant_backends.py | 284 +++++ fastvideo/mlx_runtime/refine.py | 690 ++++++++++ fastvideo/mlx_runtime/rife_interp.py | 137 ++ fastvideo/mlx_runtime/sampling.py | 139 ++ fastvideo/mlx_runtime/taehv_decode.py | 191 +++ fastvideo/mlx_runtime/wan22.py | 286 +++++ fastvideo/mlx_runtime/wan22_sample.py | 113 ++ fastvideo/mlx_runtime/wan_vae.py | 557 +++++++++ fastvideo/mlx_runtime/windowed_attention.py | 223 ++++ fastvideo/platforms/mps.py | 5 +- fastvideo/tests/mlx/test_dmd_sampling.py | 124 ++ fastvideo/tests/mlx/test_frame_upsample.py | 99 ++ fastvideo/tests/mlx/test_import_laziness.py | 30 + fastvideo/tests/mlx/test_memory_limits.py | 100 ++ fastvideo/tests/mlx/test_mlx_checkpoint.py | 123 ++ .../tests/mlx/test_mlx_compile_parity.py | 134 ++ .../test_mlx_diffusers_compile_dispatch.py | 77 ++ fastvideo/tests/mlx/test_mlx_dit_parity.py | 93 ++ fastvideo/tests/mlx/test_mlx_fast_spatial.py | 134 ++ .../tests/mlx/test_mlx_fastwan_benchmark.py | 141 +++ .../tests/mlx/test_mlx_prompt_enhance.py | 86 ++ .../mlx/test_mlx_prompt_to_video_decode.py | 85 ++ fastvideo/tests/mlx/test_mlx_refine.py | 248 ++++ .../tests/mlx/test_mlx_rife_interpolation.py | 63 + fastvideo/tests/mlx/test_mlx_taehv_parity.py | 40 + .../mlx/test_mlx_wan22_memory_handoff.py | 116 ++ fastvideo/tests/mlx/test_mlx_wan22_parity.py | 142 +++ ...test_mlx_wan22_prompt_cache_fingerprint.py | 318 +++++ .../mlx/test_mlx_wan22_real_weights_parity.py | 222 ++++ .../mlx/test_mlx_wan_vae_denormalization.py | 60 + fastvideo/tests/mlx/test_quant_backends.py | 126 ++ fastvideo/tests/mlx/test_quant_capability.py | 72 ++ fastvideo/tests/mlx/test_taehv_decode.py | 34 + fastvideo/tests/mlx/test_wan22_sample.py | 32 + .../tests/mlx/test_windowed_attention.py | 118 ++ fastvideo/tests/mlx/tiny_wan.py | 290 +++++ .../tests/platforms/test_mps_vsa_error.py | 19 + fastvideo/third_party/rife_mlx/LICENSE | 26 + fastvideo/third_party/rife_mlx/__init__.py | 17 + fastvideo/third_party/rife_mlx/config.py | 45 + .../third_party/rife_mlx/model/IFNet_HDv3.py | 142 +++ .../third_party/rife_mlx/model/RIFE_HDv3.py | 37 + .../third_party/rife_mlx/model/__init__.py | 5 + .../third_party/rife_mlx/model/warplayer.py | 37 + .../third_party/rife_mlx/ops/__init__.py | 6 + .../third_party/rife_mlx/ops/grid_sample.py | 60 + .../third_party/rife_mlx/ops/interpolate.py | 51 + fastvideo/third_party/rife_mlx/ops/spatial.py | 18 + .../third_party/rife_mlx/pipeline_mlx.py | 80 ++ .../third_party/rife_mlx/utils/__init__.py | 1 + .../third_party/rife_mlx/utils/weights.py | 44 + fastvideo/third_party/taehv/LICENSE | 21 + fastvideo/third_party/taehv/__init__.py | 10 + fastvideo/third_party/taehv/taehv.py | 475 +++++++ mkdocs.yml | 1 + pyproject.toml | 8 + 78 files changed, 13180 insertions(+), 5 deletions(-) create mode 100644 .github/workflows/ci-macos-mlx.yml create mode 100644 docs/design/apple_silicon_fast_mode.md create mode 100644 examples/inference/basic/mlx_rife_smoke_test.py create mode 100644 examples/inference/basic/mlx_wan22_generate.py create mode 100644 examples/inference/basic/mlx_wan_decode_benchmark.py create mode 100644 examples/inference/basic/mlx_wan_prompt_to_video.py create mode 100644 examples/inference/basic/mlx_wan_quant_benchmark.py create mode 100644 examples/inference/basic/mlx_wan_video_quality.py create mode 100644 fastvideo/benchmarks/__init__.py create mode 100644 fastvideo/benchmarks/eval_metalfx_rife.py create mode 100644 fastvideo/benchmarks/mlx_fastwan_bench.py create mode 100644 fastvideo/mlx_runtime/__init__.py create mode 100644 fastvideo/mlx_runtime/checkpoint.py create mode 100644 fastvideo/mlx_runtime/fast_spatial.py create mode 100644 fastvideo/mlx_runtime/fastwan.py create mode 100644 fastvideo/mlx_runtime/frame_upsample.py create mode 100644 fastvideo/mlx_runtime/memory.py create mode 100644 fastvideo/mlx_runtime/prompt_cache.py create mode 100644 fastvideo/mlx_runtime/prompt_enhance.py create mode 100644 fastvideo/mlx_runtime/quant_backends.py create mode 100644 fastvideo/mlx_runtime/refine.py create mode 100644 fastvideo/mlx_runtime/rife_interp.py create mode 100644 fastvideo/mlx_runtime/sampling.py create mode 100644 fastvideo/mlx_runtime/taehv_decode.py create mode 100644 fastvideo/mlx_runtime/wan22.py create mode 100644 fastvideo/mlx_runtime/wan22_sample.py create mode 100644 fastvideo/mlx_runtime/wan_vae.py create mode 100644 fastvideo/mlx_runtime/windowed_attention.py create mode 100644 fastvideo/tests/mlx/test_dmd_sampling.py create mode 100644 fastvideo/tests/mlx/test_frame_upsample.py create mode 100644 fastvideo/tests/mlx/test_import_laziness.py create mode 100644 fastvideo/tests/mlx/test_memory_limits.py create mode 100644 fastvideo/tests/mlx/test_mlx_checkpoint.py create mode 100644 fastvideo/tests/mlx/test_mlx_compile_parity.py create mode 100644 fastvideo/tests/mlx/test_mlx_diffusers_compile_dispatch.py create mode 100644 fastvideo/tests/mlx/test_mlx_dit_parity.py create mode 100644 fastvideo/tests/mlx/test_mlx_fast_spatial.py create mode 100644 fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py create mode 100644 fastvideo/tests/mlx/test_mlx_prompt_enhance.py create mode 100644 fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py create mode 100644 fastvideo/tests/mlx/test_mlx_refine.py create mode 100644 fastvideo/tests/mlx/test_mlx_rife_interpolation.py create mode 100644 fastvideo/tests/mlx/test_mlx_taehv_parity.py create mode 100644 fastvideo/tests/mlx/test_mlx_wan22_memory_handoff.py create mode 100644 fastvideo/tests/mlx/test_mlx_wan22_parity.py create mode 100644 fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py create mode 100644 fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py create mode 100644 fastvideo/tests/mlx/test_mlx_wan_vae_denormalization.py create mode 100644 fastvideo/tests/mlx/test_quant_backends.py create mode 100644 fastvideo/tests/mlx/test_quant_capability.py create mode 100644 fastvideo/tests/mlx/test_taehv_decode.py create mode 100644 fastvideo/tests/mlx/test_wan22_sample.py create mode 100644 fastvideo/tests/mlx/test_windowed_attention.py create mode 100644 fastvideo/tests/mlx/tiny_wan.py create mode 100644 fastvideo/tests/platforms/test_mps_vsa_error.py create mode 100644 fastvideo/third_party/rife_mlx/LICENSE create mode 100644 fastvideo/third_party/rife_mlx/__init__.py create mode 100644 fastvideo/third_party/rife_mlx/config.py create mode 100644 fastvideo/third_party/rife_mlx/model/IFNet_HDv3.py create mode 100644 fastvideo/third_party/rife_mlx/model/RIFE_HDv3.py create mode 100644 fastvideo/third_party/rife_mlx/model/__init__.py create mode 100644 fastvideo/third_party/rife_mlx/model/warplayer.py create mode 100644 fastvideo/third_party/rife_mlx/ops/__init__.py create mode 100644 fastvideo/third_party/rife_mlx/ops/grid_sample.py create mode 100644 fastvideo/third_party/rife_mlx/ops/interpolate.py create mode 100644 fastvideo/third_party/rife_mlx/ops/spatial.py create mode 100644 fastvideo/third_party/rife_mlx/pipeline_mlx.py create mode 100644 fastvideo/third_party/rife_mlx/utils/__init__.py create mode 100644 fastvideo/third_party/rife_mlx/utils/weights.py create mode 100644 fastvideo/third_party/taehv/LICENSE create mode 100644 fastvideo/third_party/taehv/__init__.py create mode 100644 fastvideo/third_party/taehv/taehv.py diff --git a/.github/workflows/ci-macos-mlx.yml b/.github/workflows/ci-macos-mlx.yml new file mode 100644 index 000000000..0737de6e2 --- /dev/null +++ b/.github/workflows/ci-macos-mlx.yml @@ -0,0 +1,149 @@ +name: macOS MLX Smoke + +on: + pull_request: + branches: [main] + paths: + - ".github/workflows/ci-macos-mlx.yml" + - "fastvideo/mlx_runtime/**" + - "fastvideo/tests/mlx/**" + - "fastvideo/tests/platforms/test_mps_vsa_error.py" + - "fastvideo/platforms/mps.py" + - "fastvideo/platforms/__init__.py" + - "fastvideo/__init__.py" + - "examples/inference/basic/mlx_*.py" + - "fastvideo/benchmarks/mlx_*.py" + - "pyproject.toml" + workflow_dispatch: + +permissions: + contents: read + +concurrency: + group: macos-mlx-${{ github.ref }} + cancel-in-progress: true + +jobs: + mlx-smoke: + if: github.event_name == 'workflow_dispatch' || github.event.pull_request.draft != true + runs-on: macos-15 + timeout-minutes: 25 + env: + FASTVIDEO_ATTENTION_BACKEND: TORCH_SDPA + TOKENIZERS_PARALLELISM: "false" + MASTER_ADDR: localhost + MASTER_PORT: "29513" + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + + - uses: astral-sh/setup-uv@v3 + + - name: Install lightweight MLX smoke dependencies + run: | + uv pip install --system \ + --index-url https://download.pytorch.org/whl/cpu \ + torch==2.11.0 torchvision torchaudio + uv pip install --system \ + pytest numpy scipy pillow imageio einops cloudpickle filelock \ + PyYAML diffusers huggingface_hub remote-pdb safetensors loguru mlx \ + "ftfy>=6.3.1" "opencv-python>=4.10.0.84" psutil "transformers>=5.0.0" + + - name: Show Apple runtime + run: | + python - <<'PY' + import platform + import mlx.core as mx + import torch + + print("machine:", platform.machine()) + print("processor:", platform.processor()) + print("mlx default device:", mx.default_device()) + memory_size = mx.metal.device_info().get("memory_size") if mx.metal.is_available() else "metal unavailable" + print("mlx memory_size:", memory_size) + print("torch:", torch.__version__) + print("torch mps available:", torch.backends.mps.is_available()) + PY + + - name: Run MLX smoke tests + run: | + python -m pytest \ + fastvideo/tests/mlx/test_dmd_sampling.py \ + fastvideo/tests/mlx/test_memory_limits.py \ + fastvideo/tests/mlx/test_quant_capability.py \ + fastvideo/tests/mlx/test_mlx_dit_parity.py \ + fastvideo/tests/mlx/test_mlx_compile_parity.py \ + fastvideo/tests/mlx/test_mlx_checkpoint.py \ + fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \ + fastvideo/tests/mlx/test_taehv_decode.py \ + fastvideo/tests/mlx/test_frame_upsample.py \ + fastvideo/tests/mlx/test_mlx_fast_spatial.py \ + fastvideo/tests/mlx/test_mlx_refine.py \ + fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \ + fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \ + fastvideo/tests/mlx/test_wan22_sample.py \ + fastvideo/tests/mlx/test_windowed_attention.py \ + fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_download_unavailable_has_specific_error \ + fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_backend_regression_is_not_skip_eligible \ + fastvideo/tests/platforms/test_mps_vsa_error.py \ + -q + + # Same tests on MLX's CPU backend. Hosted macOS runners are scarce and + # slower to schedule; this Linux job gives fast PR signal on the identical + # graph (the parity tests were designed to be backend-agnostic), while the + # macOS job above stays the source of truth for Metal behavior. + mlx-smoke-linux-cpu: + if: github.event_name == 'workflow_dispatch' || github.event.pull_request.draft != true + runs-on: ubuntu-latest + timeout-minutes: 20 + env: + FASTVIDEO_ATTENTION_BACKEND: TORCH_SDPA + TOKENIZERS_PARALLELISM: "false" + MASTER_ADDR: localhost + MASTER_PORT: "29513" + steps: + - uses: actions/checkout@v4 + + - uses: actions/setup-python@v5 + with: + python-version: "3.12" + cache: pip + + - uses: astral-sh/setup-uv@v3 + + - name: Install lightweight MLX smoke dependencies (CPU backend) + run: | + uv pip install --system \ + --index-url https://download.pytorch.org/whl/cpu \ + torch==2.11.0 torchvision torchaudio + uv pip install --system \ + pytest numpy scipy pillow imageio einops cloudpickle filelock \ + PyYAML diffusers huggingface_hub remote-pdb safetensors loguru "mlx[cpu]" \ + "ftfy>=6.3.1" "opencv-python>=4.10.0.84" psutil "transformers>=5.0.0" + + - name: Run MLX smoke tests (CPU backend) + run: | + python -m pytest \ + fastvideo/tests/mlx/test_dmd_sampling.py \ + fastvideo/tests/mlx/test_memory_limits.py \ + fastvideo/tests/mlx/test_quant_capability.py \ + fastvideo/tests/mlx/test_mlx_dit_parity.py \ + fastvideo/tests/mlx/test_mlx_compile_parity.py \ + fastvideo/tests/mlx/test_mlx_checkpoint.py \ + fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \ + fastvideo/tests/mlx/test_taehv_decode.py \ + fastvideo/tests/mlx/test_frame_upsample.py \ + fastvideo/tests/mlx/test_mlx_fast_spatial.py \ + fastvideo/tests/mlx/test_mlx_refine.py \ + fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \ + fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \ + fastvideo/tests/mlx/test_wan22_sample.py \ + fastvideo/tests/mlx/test_windowed_attention.py \ + fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_download_unavailable_has_specific_error \ + fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_backend_regression_is_not_skip_eligible \ + fastvideo/tests/platforms/test_mps_vsa_error.py \ + -q diff --git a/README.md b/README.md index fb293374a..373db6e89 100644 --- a/README.md +++ b/README.md @@ -62,6 +62,11 @@ UV_TORCH_BACKEND=cu126 uv pip install fastvideo Use `UV_TORCH_BACKEND=cu130` on CUDA 13. Apple silicon users should follow the [MPS installation guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/). +> **On an Apple Silicon Mac?** FastVideo runs FastWan text-to-video natively +> through an MLX runtime — a 5-second 480p clip generated locally, no cloud, +> no discrete GPU. Install with `uv pip install -e '.[mlx]'` and follow the +> [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/). + Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions. > **On an NVIDIA DGX Spark (GB10 / ARM64 + CUDA 13)?** There's no prebuilt ARM wheel for the FastVideo CUDA kernel, so it's an editable from-source install (`UV_TORCH_BACKEND=cu130 uv pip install -e .`, which compiles that kernel for you) rather than `UV_TORCH_BACKEND=cu130 uv pip install fastvideo`. A compatible prebuilt ARM64 FlashAttention wheel is available separately. Follow the [DGX Spark install guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/spark/). diff --git a/docs/design/apple_silicon_fast_mode.md b/docs/design/apple_silicon_fast_mode.md new file mode 100644 index 000000000..5ab5017a2 --- /dev/null +++ b/docs/design/apple_silicon_fast_mode.md @@ -0,0 +1,128 @@ +# Fast mode (RIFE) — Apple Silicon + +`--fast` makes local generation ~2.7× faster by **generating fewer frames and +interpolating the rest** with an Apple-Silicon-native RIFE model, instead of +denoising every frame. Video-diffusion denoise is dominated by self-attention, +which is O(tokens²); halving the frames cuts the token count ~2× and the denoise +compute ~3.7×, so the wall-clock drops far more than 2×. RIFE (which estimates +its own optical flow — no motion vectors needed) fills the dropped frames back +in for ~1.4 s, and a light unsharp pass counters its softening. + +Measured on the 1.3B INT8 QAD model (fox, 480×832×81, M4): generate 41 + RIFE→81 +runs in ~35 s of denoise vs ~90 s full, at reconstruction MS-SSIM **0.97**. +Reproduce with `python -m fastvideo.benchmarks.eval_metalfx_rife --mode int8`. + +> **Note:** Apple's *MetalFX* frame interpolation is **not** usable here — it +> requires game-engine motion vectors + depth, which diffusion output lacks. We +> use the video-native **`rife-mlx`** model instead (Metal-backed, torch-free). + +## Install + +```bash +uv pip install -e ".[mlx]" # RIFE ships vendored; this only needs MLX +``` + +## Use + +```bash +python examples/inference/basic/mlx_wan_prompt_to_video.py \ + --mlx-checkpoint \ + --prompt "A red fox trotting through a snowy pine forest at golden hour, cinematic" \ + --num-frames 81 --fast \ + --output-path video_samples/fox_fast.mp4 +``` + +`--num-frames` stays the *target* length; fast mode generates the smallest +VAE-aligned keyframe count that RIFE can interpolate to that target. + +| Flag | Default | Meaning | +|---|---|---| +| `--fast` / `--no-fast` | off | enable fast mode | +| `--fast-factor` | 2 | generate 1/factor of the frames (2 = half) | +| `--fast-sharpen` | 0.6 | light unsharp strength to counter RIFE softness (0 disables) | + +Fast mode composes with everything else (`--mlx-quantization int8`, +`--mlx-compile`, TAEHV vs `--decode-backend wan-vae`). Keep `--fast-factor` at 2 +for quality — larger temporal gaps are where RIFE starts inventing motion. + +## Spatial fast mode (`--fast-spatial`) + +The spatial twin of `--fast`: instead of dropping frames, drop pixels. Denoise +*and decode* at `height/width // fast-spatial-scale`, then resample the decoded +frames up to the requested size. Self-attention is O(tokens²), so halving each +spatial axis cuts the token count 4× and the denoise time far more than that — +measured on the 1.3B INT8 QAD model at 480×832×81, M4 Max: **86.1 s → 10.3 s** +of denoise. It composes with `--fast`; both together run the same clip in +**4.5 s** of denoise. + +```bash +python examples/inference/basic/mlx_wan_prompt_to_video.py \ + --prompt "A red fox trotting through a snowy pine forest at golden hour, cinematic" \ + --height 480 --width 832 --num-frames 81 --fast-spatial \ + --output-path video_samples/fox_fast_spatial.mp4 +``` + +| Flag | Default | Meaning | +|---|---|---| +| `--fast-spatial` / `--no-fast-spatial` | off | enable spatial fast mode | +| `--fast-spatial-scale` | 2 | denoise at 1/scale of each spatial axis | +| `--fast-spatial-upsample-mode` | `lanczos` | pixel interpolation kernel (`lanczos`, `cubic`, `bilinear`, `nearest`) | +| `--fast-spatial-sharpen` | 0.4 | light unsharp strength to counter resampling softness (0 disables) | + +### The upsample must happen in pixel space + +This is the one thing to get right. The obvious implementation — bilinearly +upsample the finished latents and decode at the target size — **does not work**, +and produces a distinctive failure: correct composition and silhouette under a +smeared, hazy veil, with ringing along strong edges. + +A Wan latent cell is a *learned code* for an 8×8 (Wan2.1) or 16×16 (Wan2.2) +pixel block, not a low-pass sample of the image. The average of two adjacent +codes is not the code of the averaged blocks; it is a vector the decoder was +never trained on. Measured on Wan2.1-1.3B at 480×832, a 2× bilinear latent +upsample destroys **62%** of the latent's high-frequency energy while leaving +its overall magnitude intact — exactly the signature of that veil. At Wan2.2-5B +the same operation degrades to black or noise. + +Decoded RGB frames have no such problem: an image *is* a sampled 2-D signal, so +Lanczos interpolation is the operation it was defined for. The result is soft — +it carries stage-1's real detail budget and no more — but clean and coherent. + +`--refine` gets away with a latent-space upsample only because a second DMD pass +re-denoises the hand-off; spatial fast mode passes the latent straight to the +decoder, so it cannot. + +## Refine (`--refine`) stage-2 timesteps + +`--refine` hands stage 1 to stage 2 as `(1 - sigma) * upsampled + sigma * noise`, +where `sigma` comes from the *first* stage-2 timestep. FastWan's DMD grid opens +at `t=1000`, which is `sigma == 1` exactly — so a stage-2 grid that starts there +weights the stage-1 result at zero and refine silently degrades into a plain +full-resolution run at twice the cost. + +Left unset, `--refine-dmd-denoising-steps` now derives the stage-2 grid from the +stage-1 one with leading full-noise steps dropped (`1000,757,522` → `757,522`). +That keeps the pass on timesteps the distilled student was trained on while +letting stage-1 structure through: hand-off `sigma = 0.757`, stage-1 weight +`0.243`. Passing a grid that starts at full noise is now an error rather than a +silently wasted pass. + +The run prints the resolved hand-off so it is visible: + +``` +[refine] stage-2 hand-off sigma=0.7568 (stage-1 weight 0.2432) +``` + +There is a trade-off in choosing that grid. Later start = more of the draft +survives, but fewer stage-2 steps. On Wan2.1 the default `757,522` gives weight +0.243 with two steps; `--refine-dmd-denoising-steps 522` gives weight 0.478 with +one. `--refine-sigma` decouples the noise level from the timestep entirely — it +logs a warning, because the DiT is then told a timestep that does not match the +noise it receives. + +**Wan2.2-5B has a lower ceiling.** Its warped schedule maps `1000,757,522` to +sigmas `1.000, 0.940, 0.845`, so the best available stage-1 weight is **0.060** +(vs 0.243 at 1.3B). Un-warped (`--no-warp`) the same grid gives `1.000, 0.757, +0.522` and a weight of 0.243 — but warping is what matches the FastVideo +sampling schedule, so turning it off changes the timesteps the distilled student +sees. Which is better at 5B is unresolved and needs a run on real 5B weights. diff --git a/docs/getting_started/installation/mps.md b/docs/getting_started/installation/mps.md index 852f71234..2446ff88f 100644 --- a/docs/getting_started/installation/mps.md +++ b/docs/getting_started/installation/mps.md @@ -49,10 +49,12 @@ brew install ffmpeg ### Installation +FastWan's native Apple Silicon runtime requires the `mlx` extra. + #### With uv (recommended) ```bash -uv pip install fastvideo +uv pip install "fastvideo[mlx]" ``` #### With Conda environment (alternative) @@ -60,7 +62,7 @@ uv pip install fastvideo `uv` works inside an active conda env too, so prefer `uv pip` for the actual install: ```bash -uv pip install fastvideo +uv pip install "fastvideo[mlx]" ``` ### Installation from Source @@ -76,13 +78,13 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo Basic installation: ```bash -uv pip install -e . +uv pip install -e ".[mlx]" ``` Alternative with Conda environment: ```bash -uv pip install -e . +uv pip install -e ".[mlx]" ``` ## Development Environment Setup diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index 3de50def2..cee5598f6 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -178,6 +178,17 @@ optimizations: absence means **untested**, not incompatible. | Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 | ⭕ | ⭕ | ⭕ | ⭕ | ⭕ | | GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p | ❌ | ❌ | ❌ | ⭕ | ⭕ | +## Apple Silicon native runtime + +| Release path | Model | Mode | Validated hardware | Status | +| --- | --- | --- | --- | --- | +| MLX FastWan T2V | FastWan-QAD-INT8-1.3B `[release model ID pending]` | 480x832, 81 frames, 3-step DMD, INT8 DiT + TAEHV decode | Apple M4 Max, 36 GB unified-memory class, MLX 0.31.2 | Release candidate; requires release-owner visual sign-off | + +This is a text-to-video-only source-install release. It is validated on the +hardware listed above; MLX allocator caps are not evidence of support for a +physical 16 GB Mac. See [Apple Silicon FastWan](../getting_started/installation/mps.md) +for the supported command and release gates. + **Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue. ***Lucy Edit Dev uses a non-commercial model license. FastVideo support is diff --git a/examples/inference/basic/mlx_rife_smoke_test.py b/examples/inference/basic/mlx_rife_smoke_test.py new file mode 100644 index 000000000..f5d239fc9 --- /dev/null +++ b/examples/inference/basic/mlx_rife_smoke_test.py @@ -0,0 +1,46 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tiny MLX RIFE frame-interpolation smoke test.""" + +from __future__ import annotations + +import argparse +import time + +import numpy as np + +from fastvideo.mlx_runtime.rife_interp import interpolate, load_model + + +def main() -> None: + parser = argparse.ArgumentParser( + description="MLX RIFE 4.25 frame interpolation smoke test." + ) + parser.add_argument( + "--self-test", + action="store_true", + help="Run a tiny two-frame interpolation test.", + ) + args = parser.parse_args() + if not args.self_test: + raise SystemExit("Nothing to do; pass --self-test") + + frame0 = np.zeros((64, 96, 3), dtype=np.uint8) + frame1 = np.zeros((64, 96, 3), dtype=np.uint8) + frame1[:, :, 0] = 255 + start = time.perf_counter() + model = load_model() + load_s = time.perf_counter() - start + start = time.perf_counter() + frames = interpolate([frame0, frame1], factor=2, model=model) + interp_s = time.perf_counter() - start + assert len(frames) == 3 + assert frames[1].shape == frame0.shape + assert frames[1].dtype == np.uint8 + print( + "MLX RIFE self-test passed: " + f"load_s={load_s:.3f} interp_s={interp_s:.3f} shape={frames[1].shape}" + ) + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/mlx_wan22_generate.py b/examples/inference/basic/mlx_wan22_generate.py new file mode 100644 index 000000000..30e247537 --- /dev/null +++ b/examples/inference/basic/mlx_wan22_generate.py @@ -0,0 +1,510 @@ +# SPDX-License-Identifier: Apache-2.0 +"""End-to-end Wan2.2-TI2V-5B generation on Apple Silicon (MLX DiT + MLX TAEHV). + +Pipeline: torch/MPS UMT5 encode (shared with 1.3B) → MLXWan22DiT 3-step DMD +(warped schedule, flow_shift=5) → MLX TAEHV decode (taew2_2.pth). Fully MLX +on the heavy DiT + decode path. + + PYTHONPATH=$PWD python examples/inference/basic/mlx_wan22_generate.py \ + --prompt "A red fox trotting through a snowy pine forest at golden hour" \ + --output-path video_samples/demo_5b/fox_5b_mlx.mp4 + +Decoder backends: ``taehv`` (default, MLX, ~seconds), ``taehv-torch`` (parity), +``wan-vae`` (full AutoencoderKLWan on MPS, slow). +""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np + +from fastvideo.mlx_runtime.fast_spatial import DEFAULT_FAST_SPATIAL_SHARPEN +from fastvideo.mlx_runtime.frame_upsample import DEFAULT_PIXEL_UPSAMPLE_MODE, PIXEL_UPSAMPLE_MODES +from fastvideo.mlx_runtime.memory import cleanup_mlx +from fastvideo.mlx_runtime.prompt_cache import ( + fingerprint_digest, + load_prompt_cache, + save_prompt_cache, + text_encoder_fingerprint, +) +from fastvideo.mlx_runtime.rife_interp import aligned_keyframe_count + +FASTWAN21_MODEL_ID = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers" +FASTWAN22_MODEL_ID = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers" +DEFAULT_HEIGHT = 448 +DEFAULT_WIDTH = 832 +DEFAULT_NUM_FRAMES = 121 + +def _resolve_model_paths( + *, + text_encoder_root: Path | None, + dit_checkpoint: Path | None, + dit_config: Path | None, + vae_root: Path | None, + mlx_checkpoint: Path | None, + decode_backend: str, +) -> tuple[Path, Path | None, Path | None, Path | None]: + """Download only the missing assets required by the selected Wan2.2 path.""" + from huggingface_hub import snapshot_download + + if text_encoder_root is None: + text_encoder_root = Path(snapshot_download( + FASTWAN21_MODEL_ID, + allow_patterns=["tokenizer/*", "text_encoder/*"], + )) + if mlx_checkpoint is None and (dit_checkpoint is None or dit_config is None): + patterns = [] + if dit_checkpoint is None: + patterns.append("transformer/diffusion_pytorch_model.safetensors") + if dit_config is None: + patterns.append("transformer/config.json") + model_root = Path(snapshot_download(FASTWAN22_MODEL_ID, allow_patterns=patterns)) + dit_checkpoint = dit_checkpoint or model_root / "transformer/diffusion_pytorch_model.safetensors" + dit_config = dit_config or model_root / "transformer/config.json" + if decode_backend == "wan-vae" and vae_root is None: + model_root = Path(snapshot_download(FASTWAN22_MODEL_ID, allow_patterns=["vae/*"])) + vae_root = model_root / "vae" + return text_encoder_root, dit_checkpoint, dit_config, vae_root + + +def _prompt_cache_fingerprint( + *, + prompt: str, + prompt_used: str, + enhance_prompt: bool, + enhance_prompt_backend: str, + text_encoder_root: Path, + max_sequence_length: int, + dtype: str, +) -> dict[str, object]: + return { + "prompt": prompt, + "prompt_used": prompt_used, + "enhance_prompt": enhance_prompt, + "enhance_prompt_backend": enhance_prompt_backend, + "text_encoder": text_encoder_fingerprint(text_encoder_root), + "max_sequence_length": max_sequence_length, + "dtype": dtype, + } + + +def _default_prompt_cache_path(fingerprint: dict[str, object]) -> Path: + """Content-addressed default cache file for a prompt fingerprint. + + The Wan2.1 entrypoint caches prompt embeddings by default; this one only + did so when handed an explicit ``--prompt-embeds-cache`` path, so every 5B + run paid a full UMT5 encode (~45s on an M4 Max) even for a repeat prompt. + The fingerprint already covers everything that changes the embedding, so + hash it for the filename. + """ + digest = fingerprint_digest(fingerprint)[:32] + return Path.home() / ".cache" / "fastvideo" / "prompt_embeds" / f"wan22_{digest}.npy" + + +def main() -> None: + parser = argparse.ArgumentParser( + description="MLX Wan2.2-5B T2V (encode → DiT DMD → TAEHV/VAE decode)" + ) + parser.add_argument( + "--prompt", + default="A red fox trotting through a snowy pine forest at golden hour, cinematic", + ) + parser.add_argument( + "--output-path", + type=Path, + default=Path("video_samples/demo_5b/fox_5b_mlx.mp4"), + ) + parser.add_argument( + "--text-encoder-root", + type=Path, + default=None, + help="Root with text_encoder/ + tokenizer/", + ) + parser.add_argument( + "--prompt-embeds-cache", + type=Path, + default=None, + help="Explicit .npy UMT5 embedding cache file. Overrides the automatic " + "content-addressed cache (--prompt-cache).", + ) + parser.add_argument( + "--prompt-cache", + action=argparse.BooleanOptionalAction, + default=True, + help="Cache prompt embeddings under ~/.cache/fastvideo/prompt_embeds so " + "repeat runs skip the text encoder entirely. Default: on.", + ) + parser.add_argument( + "--text-encoder-device", + choices=("auto", "cpu", "mps"), + default="cpu", + help="Device for UMT5 encoding. CPU is safest beside the 5B MLX DiT.", + ) + parser.add_argument( + "--enhance-prompt", + action="store_true", + help="Apply deterministic local cinematic prompt enrichment before UMT5.", + ) + parser.add_argument( + "--enhance-prompt-backend", + choices=("template",), + default="template", + help="Prompt enrichment backend.", + ) + parser.add_argument( + "--dit-checkpoint", + type=Path, + default=None, + ) + parser.add_argument("--dit-config", type=Path, default=None) + parser.add_argument( + "--mlx-checkpoint", + type=Path, + default=None, + help="Pre-quantized MLX DiT checkpoint directory. Rewrapped with Wan2.2 per-token conditioning.", + ) + parser.add_argument("--vae-root", type=Path, default=None) + parser.add_argument("--height", type=int, default=DEFAULT_HEIGHT) + parser.add_argument("--width", type=int, default=DEFAULT_WIDTH) + parser.add_argument( + "--num-frames", + type=int, + default=DEFAULT_NUM_FRAMES, + help="Pixel frames (121 at 24fps = 5.04 seconds)", + ) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--renoise-seed", type=int, default=0) + parser.add_argument("--fps", type=int, default=24) + parser.add_argument("--flow-shift", type=float, default=5.0) + parser.add_argument("--dmd-denoising-steps", default="1000,757,522") + parser.add_argument( + "--no-warp", + action="store_true", + help="Disable schedule warping (debug only).", + ) + parser.add_argument( + "--fast", + action="store_true", + help="Generate fewer frames then RIFE-interpolate to --num-frames.", + ) + parser.add_argument("--fast-factor", type=int, default=2) + parser.add_argument("--fast-sharpen", type=float, default=0.6) + parser.add_argument( + "--fast-spatial", + action="store_true", + help="Denoise and decode at reduced spatial resolution, then resample " + "the decoded frames up to the target size.", + ) + parser.add_argument("--fast-spatial-scale", type=int, default=2) + parser.add_argument( + "--fast-spatial-upsample-mode", + choices=PIXEL_UPSAMPLE_MODES, + default=DEFAULT_PIXEL_UPSAMPLE_MODE, + ) + parser.add_argument("--fast-spatial-sharpen", type=float, default=DEFAULT_FAST_SPATIAL_SHARPEN) + parser.add_argument( + "--refine", + action="store_true", + help="Two-pass DMD: coarse denoise, upsample/re-noise, full-res denoise.", + ) + parser.add_argument("--refine-scale", type=int, default=2) + parser.add_argument( + "--refine-upsample-mode", + choices=("bilinear", "nearest"), + default="bilinear", + ) + parser.add_argument("--no-refine-add-noise", action="store_true") + parser.add_argument( + "--decode-backend", + choices=("taehv", "taehv-torch", "wan-vae"), + default="taehv", + ) + parser.add_argument("--save-latents", type=Path, default=None) + parser.add_argument("--metrics-json", type=Path, default=None, + help="Write measured run metadata as JSON for reports or galleries.") + parser.add_argument( + "--compile", + action="store_true", + help="Compile the DiT forward with mx.compile; fallback to eager on failure.", + ) + args = parser.parse_args() + + if args.fast_factor < 2: + parser.error("--fast-factor must be at least 2") + # --fast-spatial used to be rejected here because it upsampled the completed + # 48-channel latent, which is out of distribution for the decoder and gave + # black or noisy video. The upsample now runs on decoded frames, so the + # latent never leaves the grid it was denoised on and the mode is usable. + if args.refine and args.fast_spatial: + print("[wan22] --refine takes precedence over --fast-spatial") + args.text_encoder_root, args.dit_checkpoint, args.dit_config, args.vae_root = _resolve_model_paths( + text_encoder_root=args.text_encoder_root, + dit_checkpoint=args.dit_checkpoint, + dit_config=args.dit_config, + vae_root=args.vae_root, + mlx_checkpoint=args.mlx_checkpoint, + decode_backend=args.decode_backend, + ) + target_frames = args.num_frames + if args.fast: + args.num_frames = aligned_keyframe_count(target_frames, args.fast_factor) + print( + f"[wan22 fast] generating {args.num_frames} frames, " + f"RIFE {args.fast_factor}x -> {target_frames}" + ) + + import mlx.core as mx + import torch + + from examples.inference.basic.mlx_wan_prompt_to_video import ( + _postprocess_video, + encode_prompt, + make_rotary_embeddings, + ) + from fastvideo.mlx_runtime.fast_spatial import plan_fast_spatial + from fastvideo.mlx_runtime.refine import ( + default_refine_timesteps, + plan_refine_resolutions, + prepare_refine_latents, + ) + from fastvideo.mlx_runtime.wan22 import ( + mlx_wan22_dit_from_diffusers_safetensors, + mlx_wan22_dit_from_mlx_checkpoint, + ) + from fastvideo.mlx_runtime.wan22_sample import build_wan22_dmd_schedule, sample_wan22_dmd + from fastvideo.mlx_runtime.wan_vae import decode_latents_to_video + + if args.mlx_checkpoint is not None: + config = json.loads((args.mlx_checkpoint / "mlx_dit.json").read_text())["config"] + else: + config = json.loads(args.dit_config.read_text()) + patch_size = tuple(config.get("patch_size", (1, 2, 2))) + if args.refine: + active_plan = plan_refine_resolutions( + height=args.height, width=args.width, num_frames=args.num_frames, + spatial_scale=args.refine_scale, vae_spatial_compression=16, + vae_temporal_compression=4, patch_size=patch_size, enabled=True, + ) + spatial_mode = "refine" + elif args.fast_spatial: + fast_spatial_plan = plan_fast_spatial( + height=args.height, width=args.width, num_frames=args.num_frames, + spatial_scale=args.fast_spatial_scale, vae_spatial_compression=16, + vae_temporal_compression=4, patch_size=patch_size, + upsample_mode=args.fast_spatial_upsample_mode, + sharpen=args.fast_spatial_sharpen, enabled=True, + ) + active_plan = fast_spatial_plan.plan + spatial_mode = "fast_spatial" + else: + active_plan = plan_refine_resolutions( + height=args.height, width=args.width, num_frames=args.num_frames, + spatial_scale=1, vae_spatial_compression=16, vae_temporal_compression=4, + patch_size=patch_size, enabled=False, + ) + spatial_mode = "off" + lat_h, lat_w = active_plan.stage1_latent_height, active_plan.stage1_latent_width + lat_t = active_plan.latent_frames + in_ch = int(config["in_channels"]) + print(f"[5B] latent {in_ch}x{lat_t}x{lat_h}x{lat_w}", flush=True) + + total_start = time.perf_counter() + prompt_for_encode = args.prompt + enhance_backend = None + enhance_elapsed_s = 0.0 + if args.enhance_prompt: + from fastvideo.mlx_runtime.prompt_enhance import enhance_prompt + + enhancement = enhance_prompt(args.prompt, backend=args.enhance_prompt_backend) + prompt_for_encode = enhancement.enhanced + enhance_backend = enhancement.backend + enhance_elapsed_s = enhancement.elapsed_s + print(f"[enhance] backend={enhance_backend} in {enhance_elapsed_s:.2f}s", flush=True) + print(f"[enhance] prompt: {prompt_for_encode}", flush=True) + + t0 = time.perf_counter() + prompt_cache_fingerprint = _prompt_cache_fingerprint( + prompt=args.prompt, + prompt_used=prompt_for_encode, + enhance_prompt=args.enhance_prompt, + enhance_prompt_backend=args.enhance_prompt_backend, + text_encoder_root=args.text_encoder_root, + max_sequence_length=512, + dtype="fp16", + ) + prompt_cache_path = args.prompt_embeds_cache + if prompt_cache_path is None and args.prompt_cache: + prompt_cache_path = _default_prompt_cache_path(prompt_cache_fingerprint) + cached_embeds = load_prompt_cache( + prompt_cache_path, + prompt_cache_fingerprint, + ) + if cached_embeds is not None: + embeds = torch.from_numpy(cached_embeds).contiguous() + else: + embeds = encode_prompt( + model_root=args.text_encoder_root, + prompt=prompt_for_encode, + max_sequence_length=512, + device_arg=args.text_encoder_device, + dtype_arg="fp16", + ) + save_prompt_cache( + prompt_cache_path, + embeds.cpu().numpy(), + prompt_cache_fingerprint, + ) + ehs = mx.array(embeds.numpy()).astype(mx.float16) + prompt_encode_s = time.perf_counter() - t0 + print(f"[5B] prompt encoded {tuple(ehs.shape)} in {prompt_encode_s:.1f}s", flush=True) + + t1 = time.perf_counter() + if args.mlx_checkpoint is not None: + dit = mlx_wan22_dit_from_mlx_checkpoint( + args.mlx_checkpoint, + compile=args.compile, + ) + else: + dit = mlx_wan22_dit_from_diffusers_safetensors( + args.dit_checkpoint, + args.dit_config, + dtype="fp16", + compile=args.compile, + ) + dit_load_s = time.perf_counter() - t1 + print(f"[5B] DiT loaded in {dit_load_s:.1f}s", flush=True) + + freqs = make_rotary_embeddings(config, latent_frames=lat_t, latent_height=lat_h, latent_width=lat_w) + gen = torch.Generator().manual_seed(args.seed) + noise = mx.array( + torch.randn(1, in_ch, lat_t, lat_h, lat_w, generator=gen, dtype=torch.float32).numpy()).astype(mx.float16) + + steps = [int(s) for s in args.dmd_denoising_steps.split(",") if s.strip()] + t2 = time.perf_counter() + mx.reset_peak_memory() + latents = sample_wan22_dmd( + dit, + ehs, + noise, + freqs, + dmd_denoising_steps=steps, + flow_shift=args.flow_shift, + warp_denoising_step=not args.no_warp, + seed=args.renoise_seed, + ) + if spatial_mode == "refine": + schedule, warped_steps = build_wan22_dmd_schedule( + steps, flow_shift=args.flow_shift, warp_denoising_step=not args.no_warp, + ) + # The grid opens at sigma == 1, where the hand-off + # `(1 - sigma) * upsampled + sigma * noise` weights stage 1 at zero and + # refine silently becomes a plain full-res run. Drop the leading + # full-noise steps so stage 1 actually reaches stage 2. + stage2_warped = default_refine_timesteps(schedule, warped_steps) + stage2_steps = steps[len(warped_steps) - len(stage2_warped):] + sigma = schedule.sigma_for(stage2_warped[0]) + print(f"[5B refine] stage-2 steps={stage2_steps} sigma={sigma:.4f} " + f"(stage-1 weight {1.0 - sigma:.4f})", flush=True) + latents = prepare_refine_latents( + latents, scale=args.refine_scale, sigma=sigma, + add_noise_flag=not args.no_refine_add_noise, + upsample_mode=args.refine_upsample_mode, seed=args.renoise_seed + 1, + ) + freqs_stage2 = make_rotary_embeddings( + config, latent_frames=lat_t, + latent_height=active_plan.stage2_latent_height, + latent_width=active_plan.stage2_latent_width, + ) + latents = sample_wan22_dmd( + dit, ehs, latents, freqs_stage2, dmd_denoising_steps=stage2_steps, + flow_shift=args.flow_shift, warp_denoising_step=not args.no_warp, + seed=args.renoise_seed + 2, + ) + # spatial_mode == "fast_spatial" leaves the latents on the stage-1 grid; + # the resample happens after decode, in _postprocess_video. + denoise_s = time.perf_counter() - t2 + peak = mx.get_peak_memory() / (1024**3) + print(f"[5B] denoise {len(steps)} steps in {denoise_s:.1f}s, peak {peak:.2f} GiB", flush=True) + + latents_np = np.array(latents.astype(mx.float32)) + if args.save_latents is not None: + args.save_latents.parent.mkdir(parents=True, exist_ok=True) + np.savez(args.save_latents, latents=latents_np, prompt=args.prompt, seed=args.seed) + print(f"[5B] wrote latents {args.save_latents}", flush=True) + + if spatial_mode == "refine": + del freqs_stage2 + del dit, latents, ehs, noise, freqs + cleanup_mlx() + + metrics = decode_latents_to_video( + latents_np, + args.output_path, + fps=args.fps, + backend=args.decode_backend, + vae_dir=args.vae_root if args.decode_backend == "wan-vae" else None, + z_dim=in_ch, + ) + # One h264 round-trip for both post-decode passes (see _postprocess_video). + rife_s = 0.0 + rife_request = ({ + "factor": args.fast_factor, + "target_frames": target_frames, + "sharpen": args.fast_sharpen, + } if args.fast else None) + spatial_request = fast_spatial_plan if spatial_mode == "fast_spatial" else None + if rife_request is not None or spatial_request is not None: + rife_start = time.perf_counter() + _postprocess_video( + video_path=args.output_path, fps=args.fps, + rife=rife_request, spatial=spatial_request, + ) + rife_s = time.perf_counter() - rife_start + print(f"[5B] decoded via {metrics['backend']} in {metrics['decode_s']:.1f}s → {args.output_path}", flush=True) + summary = { + "output_path": str(args.output_path.resolve()), + "prompt": args.prompt, + "prompt_used": prompt_for_encode, + "enhance_prompt": args.enhance_prompt, + "enhance_backend": enhance_backend, + "enhance_elapsed_s": round(enhance_elapsed_s, 3), + "height": args.height, + "width": args.width, + "fps": args.fps, + "target_frames": target_frames, + "generated_frames": args.num_frames, + "seed": args.seed, + "renoise_seed": args.renoise_seed, + "dmd_denoising_steps": steps, + "flow_shift": args.flow_shift, + "warp": not args.no_warp, + "spatial_mode": spatial_mode, + "fast": args.fast, + "fast_factor": args.fast_factor if args.fast else None, + "fast_spatial_scale": args.fast_spatial_scale if args.fast_spatial else None, + "refine_scale": args.refine_scale if args.refine else None, + "decode_backend": args.decode_backend, + "prompt_encode_s": round(prompt_encode_s, 3), + "dit_load_s": round(dit_load_s, 3), + "denoise_s": round(denoise_s, 3), + "decode_s": round(metrics["decode_s"], 3), + "rife_s": round(rife_s, 3), + "wall_total_s": round(time.perf_counter() - total_start, 3), + "peak_gib": round(peak, 3), + "latent_shape": [in_ch, lat_t, lat_h, lat_w], + "stage2_latent_shape": [in_ch, lat_t, active_plan.stage2_latent_height, active_plan.stage2_latent_width], + "mlx_checkpoint": str(args.mlx_checkpoint.resolve()) if args.mlx_checkpoint else None, + } + if args.metrics_json is not None: + args.metrics_json.parent.mkdir(parents=True, exist_ok=True) + args.metrics_json.write_text(json.dumps(summary, indent=2) + "\n") + print(f"[5B] wrote metrics {args.metrics_json}", flush=True) + print(json.dumps(summary, indent=2), flush=True) + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/mlx_wan_decode_benchmark.py b/examples/inference/basic/mlx_wan_decode_benchmark.py new file mode 100644 index 000000000..1621bd774 --- /dev/null +++ b/examples/inference/basic/mlx_wan_decode_benchmark.py @@ -0,0 +1,130 @@ +"""Compare Wan VAE and TAEHV decode on saved FastWan latents.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path + +import numpy as np + +from examples.inference.basic.mlx_wan_prompt_to_video import DEFAULT_MODEL_ROOT, decode_latents_to_video + + +def _torch_mps_memory() -> dict[str, int | None]: + """ + Report current and recommended memory usage for the MPS backend. + + Returns: + dict[str, int | None]: Memory metrics in bytes, or `None` values when + PyTorch or MPS is unavailable. + """ + try: + import torch + except ImportError: + return { + "current_allocated_bytes": None, + "driver_allocated_bytes": None, + "recommended_max_bytes": None, + } + if not torch.backends.mps.is_available(): + return { + "current_allocated_bytes": None, + "driver_allocated_bytes": None, + "recommended_max_bytes": None, + } + return { + "current_allocated_bytes": int(torch.mps.current_allocated_memory()), + "driver_allocated_bytes": int(torch.mps.driver_allocated_memory()), + "recommended_max_bytes": int(torch.mps.recommended_max_memory()), + } + + +def _parse_backends(raw: str) -> list[str]: + """ + Parse and validate a comma-separated list of decoding backends. + + Parameters: + raw (str): Comma-separated backend names. + + Returns: + list[str]: Trimmed, supported backend names in input order. + + Raises: + ValueError: If the input contains an unsupported backend. + """ + backends = [backend.strip() for backend in raw.split(",") if backend.strip()] + allowed = {"wan-vae", "taehv"} + unknown = sorted(set(backends) - allowed) + if unknown: + raise ValueError(f"Unsupported decode backends: {unknown}") + return backends + + +def main() -> None: + """ + Benchmark selected Wan latent decoding backends and record their performance metrics. + + Loads the specified latent array, decodes it with each selected backend, exports the + results as MP4 files, and writes per-backend timing and Torch MPS memory metrics to + `metrics.json`. + """ + parser = argparse.ArgumentParser(description="Benchmark decode backends on saved Wan/FastWan latents.") + parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT) + parser.add_argument("--latents-path", type=Path, required=True) + parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_decode_benchmark")) + parser.add_argument("--backends", default="wan-vae,taehv") + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--torch-device", default="auto") + parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16") + parser.add_argument("--taehv-source-path", type=Path, default=None) + parser.add_argument("--taehv-checkpoint-path", type=Path, default=None) + parser.add_argument("--taehv-parallel", action="store_true") + args = parser.parse_args() + + latents = np.load(args.latents_path) + args.output_dir.mkdir(parents=True, exist_ok=True) + rows = [] + + for backend in _parse_backends(args.backends): + print(f"=== Decode backend: {backend} ===") + before = _torch_mps_memory() + start = time.perf_counter() + output_path = args.output_dir / f"{args.latents_path.stem}_{backend}.mp4" + decode_latents_to_video( + model_root=args.model_root, + latents_np=latents, + output_path=output_path, + fps=args.fps, + device_arg=args.torch_device, + dtype_arg=args.torch_dtype, + backend=backend, + taehv_source_path=args.taehv_source_path, + taehv_checkpoint_path=args.taehv_checkpoint_path, + taehv_parallel=args.taehv_parallel, + ) + elapsed = time.perf_counter() - start + after = _torch_mps_memory() + metrics = { + "backend": backend, + "latents_path": str(args.latents_path), + "latents_shape": list(latents.shape), + "decode_export_s": elapsed, + "torch_mps_current_before_bytes": before["current_allocated_bytes"], + "torch_mps_current_after_bytes": after["current_allocated_bytes"], + "torch_mps_driver_before_bytes": before["driver_allocated_bytes"], + "torch_mps_driver_after_bytes": after["driver_allocated_bytes"], + "torch_mps_recommended_max_bytes": after["recommended_max_bytes"], + "output_path": str(output_path), + } + rows.append(metrics) + print(json.dumps(metrics, indent=2)) + + metrics_path = args.output_dir / "metrics.json" + metrics_path.write_text(json.dumps(rows, indent=2)) + print(f"Wrote decode metrics to: {metrics_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/mlx_wan_prompt_to_video.py b/examples/inference/basic/mlx_wan_prompt_to_video.py new file mode 100644 index 000000000..144e65861 --- /dev/null +++ b/examples/inference/basic/mlx_wan_prompt_to_video.py @@ -0,0 +1,1114 @@ +"""Generate a FastWan text-to-video clip with the Apple Silicon MLX runtime. + +This is the supported source-tree entrypoint for the FastWan-QAD-INT8-1.3B +Apple release: + +- Hugging Face/torch encodes the prompt with UMT5 (bf16 by default: fp32 + exponent range without fp16 overflow risk, at fp16 memory cost). +- MLX runs the FastWan DiT denoising loop (INT8 by default, compiled with + ``mx.compile`` unless ``--no-mlx-compile``). +- TAEHV (default, fast/low-memory) or the full Wan VAE (``--decode-backend + wan-vae``, higher fidelity, bf16) decodes the final latents. + +Defaults produce the validated release shape: 480x832, 81 frames, 3-step DMD. + +Optional quality / speed levers (compose freely): + +* ``--refine`` — H3 / LTX-2 two-pass: denoise at base res, upsample + + re-noise, re-denoise at target res with the same DiT. +* ``--fast`` — RIFE temporal fast mode (fewer frames → interpolate). +* ``--fast-spatial`` — spatial twin of RIFE: denoise *and decode* at half + res, then resample the decoded frames to the target size (no second + denoise). Orthogonal to ``--fast``. Attention is O(tokens²), so this is + the largest single denoise lever available: 86.1s → 10.3s at 1.3B. +* ``--enhance-prompt`` — local Context-IR-style prompt enrichment + (template or mlx-lm) before UMT5 encode. + +``--fast`` + ``--refine`` is the B composition: fewer frames at base res, +then a full-res refine pass. Works for Wan2.1-1.3B/14B and Wan2.2-5B. +""" + +from __future__ import annotations + +import argparse +import json +import subprocess +import sys +import tempfile +import time +from dataclasses import replace +from pathlib import Path +from typing import TYPE_CHECKING + +import numpy as np + +from fastvideo.mlx_runtime.fast_spatial import DEFAULT_FAST_SPATIAL_SHARPEN +from fastvideo.mlx_runtime.frame_upsample import DEFAULT_PIXEL_UPSAMPLE_MODE, PIXEL_UPSAMPLE_MODES +from fastvideo.mlx_runtime.memory import ( + add_memory_limit_args, + apply_memory_limits, + cleanup_mlx, + cleanup_torch_mps, +) +from fastvideo.mlx_runtime.prompt_cache import ( + fingerprint_digest, + load_prompt_cache, + save_prompt_cache, + text_encoder_fingerprint, +) +from fastvideo.mlx_runtime.rife_interp import aligned_keyframe_count + +if TYPE_CHECKING: # pragma: no cover - typing only + from fastvideo.mlx_runtime.fast_spatial import FastSpatialPlan + + +DEFAULT_MODEL_ID = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers" + +# Legacy pinned-snapshot location, kept for callers that import it (the MLX +# benchmark harness). New code should prefer resolve_model_root(None), which +# resolves whatever snapshot the local HF cache has (downloading if needed). +DEFAULT_MODEL_ROOT = ( + Path.home() + / ".cache/huggingface/hub/models--FastVideo--FastWan2.1-T2V-1.3B-Diffusers/" + "snapshots/25e7ed7f41fd8ce2fdd108688c65e8caf0ce3aef" +) + + +def resolve_model_root( + model_root: Path | None, + *, + model_id: str = DEFAULT_MODEL_ID, + include_transformer: bool = True, +) -> Path: + """Return a usable model directory, resolving via the HF cache if unset. + + A user-supplied ``model_root`` is returned as-is. Otherwise the model is + resolved through ``huggingface_hub.snapshot_download``, which reuses the + local cache when present and downloads the current snapshot when not — + no hardcoded snapshot hash. + """ + if model_root is not None: + return model_root + from huggingface_hub import snapshot_download + + # Always fetch the auxiliary assets. A pre-quantized MLX checkpoint does + # not need the raw transformer weights; the default Diffusers path does. + allow_patterns = [ + "model_index.json", + "scheduler/*", + "tokenizer/*", + "text_encoder/*", + "vae/*", + "transformer/*" if include_transformer else "transformer/config.json", + ] + return Path(snapshot_download( + model_id, + allow_patterns=allow_patterns, + )) + + +def _torch_device(device_arg: str): + import torch + + if device_arg == "auto": + return torch.device("mps" if torch.backends.mps.is_available() else "cpu") + return torch.device(device_arg) + + +def _torch_dtype(dtype_arg: str): + import torch + + return {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32}[dtype_arg] + + +def encode_prompt( + *, + model_root: Path, + prompt: str, + max_sequence_length: int, + device_arg: str, + dtype_arg: str, +): + import torch + from transformers import AutoTokenizer, UMT5EncoderModel + + device = _torch_device(device_arg) + dtype = _torch_dtype(dtype_arg) + tokenizer = AutoTokenizer.from_pretrained(model_root / "tokenizer", local_files_only=True) + text_encoder = UMT5EncoderModel.from_pretrained( + model_root / "text_encoder", + torch_dtype=dtype, + local_files_only=True, + ).to(device) + text_encoder.eval() + + text_inputs = tokenizer( + [prompt], + padding="max_length", + max_length=max_sequence_length, + truncation=True, + add_special_tokens=True, + return_attention_mask=True, + return_tensors="pt", + ) + text_input_ids = text_inputs.input_ids.to(device) + mask = text_inputs.attention_mask.to(device) + seq_lens = mask.gt(0).sum(dim=1).long() + + with torch.no_grad(): + prompt_embeds = text_encoder(text_input_ids, mask).last_hidden_state + prompt_embeds = prompt_embeds.to(dtype=dtype) + prompt_embeds = [u[:v] for u, v in zip(prompt_embeds, seq_lens, strict=False)] + prompt_embeds = torch.stack( + [ + torch.cat([u, u.new_zeros(max_sequence_length - u.size(0), u.size(1))]) + for u in prompt_embeds + ], + dim=0, + ) + if prompt_embeds.dtype == torch.bfloat16: + # NumPy (and the .npy cache/subprocess transport) has no bfloat16; + # fp32 is exact for every bf16 value. + prompt_embeds = prompt_embeds.float() + prompt_embeds = prompt_embeds.cpu().contiguous() + del text_encoder, tokenizer, text_inputs, text_input_ids, mask, seq_lens + cleanup_torch_mps() + return prompt_embeds + + +def encode_prompt_subprocess( + *, + model_root: Path, + prompt: str, + max_sequence_length: int, + device_arg: str, + dtype_arg: str, +): + import torch + + with tempfile.TemporaryDirectory(prefix="fastvideo_prompt_embeds_") as tmpdir: + output_path = Path(tmpdir) / "prompt_embeds.npy" + subprocess.run( + [ + sys.executable, + str(Path(__file__).resolve()), + "--model-root", + str(model_root), + "--prompt", + prompt, + "--max-sequence-length", + str(max_sequence_length), + "--torch-device", + device_arg, + "--text-encoder-dtype", + dtype_arg, + "--encode-prompt-only", + str(output_path), + ], + check=True, + ) + prompt_embeds = np.load(output_path) + return torch.from_numpy(prompt_embeds).contiguous() + + +def _prompt_cache_fingerprint( + *, + model_root: Path, + prompt: str, + max_sequence_length: int, + dtype_arg: str, +) -> dict[str, object]: + return { + "prompt": prompt, + "text_encoder": text_encoder_fingerprint(model_root), + "max_sequence_length": max_sequence_length, + "dtype": dtype_arg, + } + + +def _default_prompt_cache_path( + *, + model_root: Path, + prompt: str, + max_sequence_length: int, + dtype_arg: str, +) -> Path: + fingerprint = _prompt_cache_fingerprint( + model_root=model_root, + prompt=prompt, + max_sequence_length=max_sequence_length, + dtype_arg=dtype_arg, + ) + digest = fingerprint_digest(fingerprint)[:32] + return Path.home() / ".cache" / "fastvideo" / "prompt_embeds" / f"{digest}.npy" + + +def get_prompt_embeds( + *, + model_root: Path, + prompt: str, + max_sequence_length: int, + device_arg: str, + dtype_arg: str, + encode_mode: str, + cache_path: Path | None, +): + import torch + + fingerprint = None + if cache_path is not None: + fingerprint = _prompt_cache_fingerprint( + model_root=model_root, + prompt=prompt, + max_sequence_length=max_sequence_length, + dtype_arg=dtype_arg, + ) + cached = load_prompt_cache(cache_path, fingerprint) + if cached is not None: + return torch.from_numpy(cached).contiguous() + + if encode_mode == "subprocess": + prompt_embeds = encode_prompt_subprocess( + model_root=model_root, + prompt=prompt, + max_sequence_length=max_sequence_length, + device_arg=device_arg, + dtype_arg=dtype_arg, + ) + elif encode_mode == "inline": + prompt_embeds = encode_prompt( + model_root=model_root, + prompt=prompt, + max_sequence_length=max_sequence_length, + device_arg=device_arg, + dtype_arg=dtype_arg, + ) + else: + raise ValueError(f"Unsupported prompt encode mode: {encode_mode}") + + if cache_path is not None and fingerprint is not None: + save_prompt_cache(cache_path, prompt_embeds.cpu().numpy(), fingerprint) + return prompt_embeds + + +def make_rotary_embeddings(config: dict, *, latent_frames: int, latent_height: int, latent_width: int): + import mlx.core as mx + import torch + + from fastvideo.layers.rotary_embedding import get_rotary_pos_embed + + num_heads = int(config["num_attention_heads"]) + head_dim = int(config["attention_head_dim"]) + hidden_size = num_heads * head_dim + patch_size = tuple(config["patch_size"]) + post_patch = ( + latent_frames // patch_size[0], + latent_height // patch_size[1], + latent_width // patch_size[2], + ) + rope_dim_list = [head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + post_patch, + hidden_size, + num_heads, + rope_dim_list, + dtype=torch.float32, + rope_theta=10000, + ) + return ( + mx.array(freqs_cos.numpy()).astype(mx.float32), + mx.array(freqs_sin.numpy()).astype(mx.float32), + ) + + +def decode_latents_to_video( + *, + model_root: Path, + latents_np: np.ndarray, + output_path: Path, + fps: int, + device_arg: str, + dtype_arg: str, + backend: str, + taehv_source_path: Path | None, + taehv_checkpoint_path: Path | None, + taehv_parallel: bool, +) -> None: + if backend == "taehv": + if taehv_source_path is None: + from fastvideo.mlx_runtime.wan_vae import decode_latents_to_video as decode_latents_to_video_mlx + + decode_latents_to_video_mlx( + latents_np, + output_path, + fps=fps, + backend="taehv", + z_dim=latents_np.shape[1], + taehv_checkpoint=taehv_checkpoint_path, + torch_device=device_arg, + ) + return + + device = _torch_device(device_arg) + dtype = _torch_dtype(dtype_arg) + from fastvideo.mlx_runtime.taehv_decode import decode_latents_to_video_taehv + + decode_latents_to_video_taehv( + latents_np=latents_np, + output_path=output_path, + fps=fps, + device=device, + dtype=dtype, + parallel=taehv_parallel, + source_path=taehv_source_path, + checkpoint_path=taehv_checkpoint_path, + ) + cleanup_torch_mps() + return + + if backend != "wan-vae": + raise ValueError(f"Unsupported decode backend: {backend}") + + import torch + from diffusers import AutoencoderKLWan + from diffusers.video_processor import VideoProcessor + from diffusers.utils import export_to_video + + device = _torch_device(device_arg) + dtype = _torch_dtype(dtype_arg) + vae = AutoencoderKLWan.from_pretrained( + model_root / "vae", + torch_dtype=dtype, + local_files_only=True, + ).to(device) + vae.eval() + + latents = torch.from_numpy(latents_np).to(device=device, dtype=dtype) + latents_mean = torch.tensor(vae.config.latents_mean).view(1, vae.config.z_dim, 1, 1, 1).to(device, dtype) + latents_std = 1.0 / torch.tensor(vae.config.latents_std).view(1, vae.config.z_dim, 1, 1, 1).to(device, dtype) + latents = latents / latents_std + latents_mean + + with torch.no_grad(): + video = vae.decode(latents, return_dict=False)[0] + video = VideoProcessor(vae_scale_factor=vae.config.scale_factor_spatial).postprocess_video(video, output_type="np") + output_path.parent.mkdir(parents=True, exist_ok=True) + export_to_video(video[0], str(output_path), fps=fps) + del vae, latents, latents_mean, latents_std, video + cleanup_torch_mps() + + +def _unsharp(frame: np.ndarray, amount: float) -> np.ndarray: + """Light unsharp mask to counter RIFE's optical-flow softening.""" + from fastvideo.mlx_runtime.frame_upsample import unsharp + + return unsharp(frame, amount) + + +def _postprocess_video(*, video_path: Path, fps: int, rife: dict | None = None, + spatial: "FastSpatialPlan | None" = None) -> None: + """Apply the post-decode passes to the written mp4 in a single re-encode. + + ``--fast`` (RIFE frame interpolation) and ``--fast-spatial`` (pixel-space + upsample of a reduced-resolution decode) both operate on decoded frames. + Running them as separate in-place rewrites would put the video through two + lossy h264 round-trips, so they share one read/write here. + + RIFE runs first, at the smaller frame size: optical flow is estimated on + fewer pixels (cheaper) and the interpolated frames then ride through the + same upsample as the keyframes, which keeps the clip spatially uniform. + + Sharpening is applied once, at the end and at full resolution. Both passes + soften for the same reason (they synthesise pixels they do not have), so + stacking two unsharp masks over-crisps; the stronger of the two requested + amounts is used instead. + """ + import imageio.v3 as iio + + if rife is None and spatial is None: + return + + frames = [frame for frame in iio.imread(video_path)] + labels = [] + sharpen = 0.0 + + if rife is not None: + from fastvideo.mlx_runtime.rife_interp import interpolate as rife_interpolate, load_model + + factor, target_frames = rife["factor"], rife["target_frames"] + frames = rife_interpolate(frames, factor=factor, model=load_model()) + if len(frames) < target_frames: + raise RuntimeError(f"RIFE produced {len(frames)} frames, fewer than requested {target_frames}") + frames = frames[:target_frames] + sharpen = max(sharpen, float(rife["sharpen"] or 0.0)) + labels.append(f"RIFE {factor}x -> {len(frames)} frames") + + if spatial is not None and spatial.enabled: + from fastvideo.mlx_runtime.fast_spatial import apply_fast_spatial_upsample + + # The plan's own sharpen is folded into the single pass below. + frames = apply_fast_spatial_upsample(frames, replace(spatial, sharpen=0.0)) + sharpen = max(sharpen, float(spatial.sharpen)) + labels.append(f"upsample {spatial.scale}x -> {spatial.target_width}x{spatial.target_height} " + f"({spatial.upsample_mode})") + + if sharpen > 0.0: + frames = [_unsharp(frame, sharpen) for frame in frames] + labels.append(f"unsharp {sharpen:.2f}") + + iio.imwrite(video_path, np.stack(frames), fps=fps, codec="libx264") + print(f"[post] {', '.join(labels)} written to {video_path}") + + +def _rife_interpolate_video(*, video_path: Path, target_frames: int, factor: int, + sharpen: float, fps: int) -> None: + """Read the reduced-frame mp4, RIFE-interpolate up to ``target_frames`` on + Apple Silicon, optionally light-sharpen, and rewrite the file in place.""" + _postprocess_video( + video_path=video_path, + fps=fps, + rife={"factor": factor, "target_frames": target_frames, "sharpen": sharpen}, + ) + + +def main() -> None: + parser = argparse.ArgumentParser(description="Prompt-to-video FastWan generation using MLX for the DiT") + parser.add_argument("--model-root", type=Path, default=None, + help=f"Model directory. Defaults to the local HF cache for {DEFAULT_MODEL_ID} " + "(downloading it if missing).") + parser.add_argument("--prompt", default="A paper boat sails through a shallow stream in a mossy forest.") + parser.add_argument("--output-path", type=Path, default=Path("video_samples/mlx_fastwan_prompt_to_video.mp4")) + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--num-frames", type=int, default=81) + parser.add_argument("--num-inference-steps", type=int, default=3) + parser.add_argument("--dmd-denoising-steps", default="1000,757,522") + parser.add_argument("--denoising-mode", choices=("dmd", "scheduler"), default="dmd") + parser.add_argument("--flow-shift", type=float, default=8.0) + parser.add_argument( + "--refine", + action=argparse.BooleanOptionalAction, + default=False, + help="Two-pass H3/LTX-2 refine: denoise at height/width // refine-scale, then " + "upsample latents and re-denoise at the full target resolution with the same " + "DiT (no new model / no training). DMD mode only.", + ) + parser.add_argument( + "--refine-scale", + type=int, + default=2, + help="Spatial upsample factor between stage-1 and stage-2 (default: 2).", + ) + parser.add_argument( + "--refine-dmd-denoising-steps", + default=None, + help="Optional stage-2 DMD timesteps (comma-separated). Defaults to " + "--dmd-denoising-steps when unset.", + ) + parser.add_argument( + "--refine-sigma", + type=float, + default=None, + help="Override the stage-2 hand-off noise level (0-1). Lower keeps more " + "of the stage-1 draft. Normally derived from the first stage-2 timestep, " + "which keeps the noise level and the timestep the DiT is told consistent; " + "setting this decouples them, so the DiT sees a latent noised differently " + "than its timestep implies. Experimental — for exploring schedules whose " + "grid bottoms out too high (see the Wan2.2-5B note in the design doc).", + ) + parser.add_argument( + "--refine-add-noise", + action=argparse.BooleanOptionalAction, + default=True, + help="Re-noise upsampled stage-1 latents before the stage-2 denoise " + "(default: on; disable with --no-refine-add-noise).", + ) + parser.add_argument( + "--refine-upsample-mode", + choices=("bilinear", "nearest"), + default="bilinear", + help="Latent spatial upsample mode for the refine hand-off.", + ) + parser.add_argument( + "--save-stage1-latents", + action="store_true", + help="When --refine is set, also dump stage-1 clean latents next to the output.", + ) + parser.add_argument( + "--fast-spatial", + action=argparse.BooleanOptionalAction, + default=False, + help="Spatial fast mode: denoise and decode at height/width // " + "fast-spatial-scale, then resample the decoded frames up to the target " + "size (no second denoise — that is --refine). Composes with --fast " + "(RIFE). If both --fast-spatial and --refine are set, refine wins " + "(quality path).", + ) + parser.add_argument( + "--fast-spatial-scale", + type=int, + default=2, + help="Spatial downsample factor for --fast-spatial (default: 2).", + ) + parser.add_argument( + "--fast-spatial-upsample-mode", + choices=PIXEL_UPSAMPLE_MODES, + default=DEFAULT_PIXEL_UPSAMPLE_MODE, + help="Pixel-space interpolation kernel used to resample the decoded " + f"stage-1 frames (default: {DEFAULT_PIXEL_UPSAMPLE_MODE}).", + ) + parser.add_argument( + "--fast-spatial-sharpen", + type=float, + default=DEFAULT_FAST_SPATIAL_SHARPEN, + help="Light unsharp strength to counter resampling softness " + f"(default: {DEFAULT_FAST_SPATIAL_SHARPEN}; 0 disables).", + ) + parser.add_argument( + "--enhance-prompt", + action=argparse.BooleanOptionalAction, + default=False, + help="Local H3 Context-IR-style prompt enrichment before UMT5 encode. " + "Uses --enhance-prompt-backend (template always available; mlx-lm optional).", + ) + parser.add_argument( + "--enhance-prompt-backend", + choices=("auto", "template", "mlx-lm"), + default="auto", + help="auto: try mlx-lm then fall back to template. template: deterministic " + "cinematic expansion. mlx-lm: require a local instruct model.", + ) + parser.add_argument( + "--enhance-prompt-model", + default=None, + help="mlx-lm model id/path (default: mlx-community/Qwen2.5-0.5B-Instruct-4bit).", + ) + parser.add_argument( + "--enhance-prompt-cache", + action=argparse.BooleanOptionalAction, + default=True, + help="Cache enhanced prompts under ~/.cache/fastvideo/enhanced_prompts (default: on).", + ) + parser.add_argument("--max-sequence-length", type=int, default=512) + parser.add_argument("--seed", type=int, default=1024) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--fast", action=argparse.BooleanOptionalAction, default=False, + help="Fast mode: generate 1/factor of the frames, then RIFE-interpolate up " + "to --num-frames on Apple Silicon (~2.7x faster denoise, reconstruction " + "MS-SSIM ~0.97). Composes with --refine (B: fewer frames at base res, " + "full-res refine) and --fast-spatial. Uses the vendored MLX RIFE backend. " + "See docs/experiments/rife-speedup-summary.md.") + parser.add_argument("--fast-factor", type=int, default=2, + help="Fast-mode interpolation factor (2 = generate half the frames).") + parser.add_argument("--fast-sharpen", type=float, default=0.6, + help="Light unsharp strength to counter RIFE softness (0 disables).") + parser.add_argument("--torch-device", default="auto", help="'auto', 'mps', or 'cpu' for text/VAE components.") + parser.add_argument("--torch-dtype", choices=("fp16", "bf16", "fp32"), default="fp16", + help="Dtype for the TAEHV decode path (and legacy callers).") + parser.add_argument("--text-encoder-dtype", choices=("bf16", "fp16", "fp32"), default="bf16", + help="UMT5 prompt-encode dtype. bf16 keeps the fp32 exponent range " + "(no fp16 overflow risk in the T5 stack) at fp16 memory cost; the " + "reference CUDA pipeline encodes in fp32. Pass fp16 if your " + "macOS/torch build lacks bf16 on MPS.") + parser.add_argument("--vae-decode-dtype", choices=("bf16", "fp16", "fp32"), default="bf16", + help="Wan VAE decode dtype (wan-vae backend only). bf16 matches the " + "reference pipeline's effectively-lossless decode default.") + parser.add_argument("--mlx-dtype", choices=("fp16", "bf16", "fp32"), default="fp16") + parser.add_argument( + "--mlx-quantization", + choices=("none", "int8", "int4", "mxfp8", "mxfp4", "nvfp4"), + default="int8", + ) + parser.add_argument("--mlx-compile", action=argparse.BooleanOptionalAction, default=True, + help="Compile the DiT forward with mx.compile (bit-identical to eager, " + "~1.4x faster denoise; falls back to eager if tracing fails). " + "Disable with --no-mlx-compile.") + parser.add_argument("--metrics-json", type=Path, default=None) + parser.add_argument("--save-latents", action="store_true") + parser.add_argument("--decode-backend", choices=("wan-vae", "taehv"), default="taehv", + help="taehv (default): fast, low-memory tiny decoder. " + "wan-vae: full Wan VAE in bf16 — slower and heavier but higher fidelity.") + parser.add_argument("--taehv-source-path", type=Path, default=None) + parser.add_argument("--taehv-checkpoint-path", type=Path, default=None) + parser.add_argument("--taehv-parallel", action="store_true", help="Decode all TAEHV frames at once; faster but higher memory.") + parser.add_argument("--prompt-encode-mode", choices=("inline", "subprocess"), default="inline") + parser.add_argument("--prompt-embeds-cache", type=Path, default=None, + help="Explicit prompt-embedding cache file. Overrides the " + "automatic content-addressed cache (--prompt-cache).") + parser.add_argument("--prompt-cache", action=argparse.BooleanOptionalAction, default=True, + help="Cache prompt embeddings under ~/.cache/fastvideo/prompt_embeds " + "keyed by (model, prompt, length, dtype), so repeat runs skip " + "the text encoder entirely. Default: on.") + parser.add_argument("--mlx-checkpoint", type=Path, default=None, + help="Load the DiT from a pre-quantized MLX checkpoint directory " + "(created with --save-mlx-checkpoint) instead of casting/quantizing " + "the Diffusers weights on every run.") + parser.add_argument("--save-mlx-checkpoint", type=Path, default=None, + help="After loading the DiT, save it (cast + quantized) as an MLX " + "checkpoint directory for fast reloads via --mlx-checkpoint.") + add_memory_limit_args(parser) + parser.add_argument("--encode-prompt-only", type=Path, default=None, help=argparse.SUPPRESS) + args = parser.parse_args() + + # Fast mode: generate fewer frames now, RIFE-interpolate back up after decode. + # Composes with --refine (B): stage-1/2 both see the reduced frame count; + # RIFE restores the target length after the final decode. + fast_target_frames = None + if args.fast: + if args.fast_factor < 2: + parser.error("--fast-factor must be >= 2") + fast_target_frames = args.num_frames + + if args.fast_spatial and args.fast_spatial_scale < 2: + parser.error("--fast-spatial-scale must be >= 2 when --fast-spatial is set") + if args.refine and args.fast_spatial: + print("[fast-spatial] note: --refine is set; refine wins (quality path). " + "--fast-spatial upsample-only path is skipped.") + + runtime_limits = apply_memory_limits( + mlx_memory_limit_gib=args.mlx_memory_limit_gib, + mlx_cache_limit_gib=args.mlx_cache_limit_gib, + mlx_disable_cache=args.mlx_disable_cache, + mlx_wired_limit_gib=args.mlx_wired_limit_gib, + torch_mps_high_watermark_ratio=args.torch_mps_high_watermark_ratio, + torch_mps_low_watermark_ratio=args.torch_mps_low_watermark_ratio, + ).as_metrics() + + model_root = resolve_model_root( + args.model_root, + include_transformer=args.mlx_checkpoint is None and args.encode_prompt_only is None, + ) + + if args.encode_prompt_only is not None: + prompt_embeds = encode_prompt( + model_root=model_root, + prompt=args.prompt, + max_sequence_length=args.max_sequence_length, + device_arg=args.torch_device, + dtype_arg=args.text_encoder_dtype, + ) + args.encode_prompt_only.parent.mkdir(parents=True, exist_ok=True) + np.save(args.encode_prompt_only, prompt_embeds.cpu().numpy()) + return + + import mlx.core as mx + import torch + from diffusers import UniPCMultistepScheduler + + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + from fastvideo.mlx_runtime.fast_spatial import ( + plan_fast_spatial, + resolve_spatial_mode, + ) + from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors + from fastvideo.mlx_runtime.prompt_enhance import ( + enhance_result_as_metrics, + load_or_enhance_prompt, + ) + from fastvideo.mlx_runtime.refine import plan_refine_resolutions, run_two_pass_dmd + from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step + + mx.random.seed(args.seed) + torch.manual_seed(args.seed) + + if args.refine and args.denoising_mode != "dmd": + raise SystemExit("--refine currently requires --denoising-mode dmd") + if args.refine and args.refine_scale < 2: + raise SystemExit("--refine-scale must be >= 2 when --refine is set") + + spatial_mode = resolve_spatial_mode(refine=args.refine, fast_spatial=args.fast_spatial) + + config_path = model_root / "transformer/config.json" + checkpoint_path = model_root / "transformer/diffusion_pytorch_model.safetensors" + config = json.loads(config_path.read_text()) + # A pre-quantized MLX DiT can be paired with a lightweight asset root for + # UMT5/TAEHV. In that case the model root is *not* the architecture + # authority: use the checkpoint's embedded transformer config for the + # sampler guard and latent geometry. + dit_config = config + if args.mlx_checkpoint is not None: + mlx_config_path = Path(args.mlx_checkpoint) / "mlx_dit.json" + if mlx_config_path.is_file(): + mlx_checkpoint_config = json.loads(mlx_config_path.read_text()) + dit_config = mlx_checkpoint_config.get("config", mlx_checkpoint_config) + if int(dit_config.get("in_channels", 0)) == 48 and int(dit_config.get("out_channels", 0)) == 48: + raise SystemExit( + "Wan2.2-TI2V-5B uses 48-channel, per-token timestep conditioning. " + "Use examples/inference/basic/mlx_wan22_generate.py instead; this " + "generic Wan2.1 sampler produces invalid outputs for that checkpoint." + ) + # Latent geometry follows the model's own VAE: Wan2.1 compresses 4x8x8, + # the Wan2.2 TI2V VAE (48 channels) compresses 4x16x16. Read the factors + # from the checkpoint's vae config rather than assuming Wan2.1. + vae_config_path = model_root / "vae/config.json" + vae_config = json.loads(vae_config_path.read_text()) if vae_config_path.is_file() else {} + # Wan2.1 uses 16 latent channels with 4x8x8 compression. Do not inherit + # the 4x16x16 VAE geometry merely because the asset root belongs to a + # Wan2.2 checkpoint. + is_wan21 = int(dit_config.get("in_channels", 0)) == 16 + vae_temporal_factor = 4 if is_wan21 else int(vae_config.get("scale_factor_temporal", 4)) + vae_spatial_factor = 8 if is_wan21 else int(vae_config.get("scale_factor_spatial", 8)) + patch_size = tuple(dit_config.get("patch_size", (1, 2, 2))) + if fast_target_frames is not None: + args.num_frames = aligned_keyframe_count( + fast_target_frames, + args.fast_factor, + temporal_compression=vae_temporal_factor, + ) + print(f"[fast] generating {args.num_frames} frames, RIFE {args.fast_factor}x -> {fast_target_frames}") + + # Spatial plan: refine (quality two-pass) or fast-spatial (upsample-only). + # Both reuse the same resolution splitter; only the post-denoise path differs. + refine_plan = plan_refine_resolutions( + height=args.height, + width=args.width, + num_frames=args.num_frames, + spatial_scale=args.refine_scale if spatial_mode == "refine" else 1, + vae_spatial_compression=vae_spatial_factor, + vae_temporal_compression=vae_temporal_factor, + patch_size=patch_size, + enabled=(spatial_mode == "refine"), + ) + fast_spatial_plan = plan_fast_spatial( + height=args.height, + width=args.width, + num_frames=args.num_frames, + spatial_scale=args.fast_spatial_scale, + vae_spatial_compression=vae_spatial_factor, + vae_temporal_compression=vae_temporal_factor, + patch_size=patch_size, + upsample_mode=args.fast_spatial_upsample_mode, + sharpen=args.fast_spatial_sharpen, + enabled=(spatial_mode == "fast_spatial"), + ) + active_plan = refine_plan if spatial_mode == "refine" else fast_spatial_plan.plan + # Stage-1 geometry drives the first denoise (and the only denoise when + # refine is off). Stage-2 / target geometry is used after the hand-off. + latent_frames = active_plan.latent_frames + latent_height = active_plan.stage1_latent_height + latent_width = active_plan.stage1_latent_width + mx_dtype = {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[args.mlx_dtype] + quantization = None if args.mlx_quantization == "none" else args.mlx_quantization + + total_start = time.perf_counter() + + # C: optional local prompt enrichment (template or mlx-lm) before UMT5. + enhance_result = None + enhance_time = 0.0 + prompt_for_encode = args.prompt + if args.enhance_prompt: + enhance_start = time.perf_counter() + enhance_result = load_or_enhance_prompt( + args.prompt, + backend=args.enhance_prompt_backend, + model=args.enhance_prompt_model, + cache=args.enhance_prompt_cache, + ) + enhance_time = time.perf_counter() - enhance_start + prompt_for_encode = enhance_result.enhanced + print( + f"[enhance] backend={enhance_result.backend} " + f"({enhance_result.elapsed_s:.2f}s cached={enhance_result.backend == 'cache'})" + ) + print(f"[enhance] original: {enhance_result.original}") + print(f"[enhance] enhanced: {enhance_result.enhanced}") + if args.enhance_prompt_backend != "template": + cleanup_mlx() + + prompt_start = time.perf_counter() + prompt_embeds = get_prompt_embeds( + model_root=model_root, + prompt=prompt_for_encode, + max_sequence_length=args.max_sequence_length, + device_arg=args.torch_device, + dtype_arg=args.text_encoder_dtype, + encode_mode=args.prompt_encode_mode, + cache_path=args.prompt_embeds_cache + or (_default_prompt_cache_path( + model_root=model_root, + prompt=prompt_for_encode, + max_sequence_length=args.max_sequence_length, + dtype_arg=args.text_encoder_dtype, + ) if args.prompt_cache else None), + ) + prompt_time = time.perf_counter() - prompt_start + + load_start = time.perf_counter() + mx.clear_cache() + mx.reset_peak_memory() + if args.mlx_checkpoint is not None: + from fastvideo.mlx_runtime.checkpoint import load_mlx_dit_checkpoint + + dit = load_mlx_dit_checkpoint(args.mlx_checkpoint, compile=args.mlx_compile) + config = dit.config + else: + dit = mlx_dit_from_diffusers_safetensors( + checkpoint_path, + config_path, + dtype=args.mlx_dtype, + quantization=quantization, + compile=args.mlx_compile, + ) + load_time = time.perf_counter() - load_start + load_peak_memory = mx.get_peak_memory() + + if args.save_mlx_checkpoint is not None: + from fastvideo.mlx_runtime.checkpoint import save_mlx_dit_checkpoint + + save_mlx_dit_checkpoint(dit, args.save_mlx_checkpoint) + + if args.denoising_mode == "dmd": + scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift) + denoising_steps = [int(step.strip()) for step in args.dmd_denoising_steps.split(",") if step.strip()] + timesteps = torch.tensor(denoising_steps, dtype=torch.long) + else: + scheduler = UniPCMultistepScheduler.from_pretrained(model_root / "scheduler", local_files_only=True) + scheduler.set_timesteps(args.num_inference_steps, device="cpu") + scheduler.set_begin_index(0) + timesteps = scheduler.timesteps + + generator = torch.Generator(device="cpu").manual_seed(args.seed) + latents_torch = torch.randn( + (1, int(config["in_channels"]), latent_frames, latent_height, latent_width), + generator=generator, + dtype=torch.float32, + ) + latents = mx.array(latents_torch.numpy()).astype(mx_dtype) + encoder_hidden_states = mx.array(prompt_embeds.numpy()).astype(mx_dtype) + freqs_cis = make_rotary_embeddings( + config, + latent_frames=latent_frames, + latent_height=latent_height, + latent_width=latent_width, + ) + + # DMD keeps the whole update on the MLX device via the native sampler. Only + # the (non-distilled) diffusers scheduler path still round-trips to torch. + dmd_schedule = MLXDMDSchedule.from_torch_scheduler(scheduler) if args.denoising_mode == "dmd" else None + + denoise_start = time.perf_counter() + mx.reset_peak_memory() + stage1_latents_np = None + refine_sigma = None + if args.refine: + assert dmd_schedule is not None # guarded above + # Left unset, the stage-2 grid is derived from the stage-1 one with the + # leading full-noise step dropped. Reusing --dmd-denoising-steps + # verbatim would start stage 2 at sigma=1 and discard stage 1. + refine_steps = None + if args.refine_dmd_denoising_steps: + refine_steps = [ + float(step.strip()) for step in args.refine_dmd_denoising_steps.split(",") if step.strip() + ] + freqs_cis_stage2 = make_rotary_embeddings( + config, + latent_frames=latent_frames, + latent_height=refine_plan.stage2_latent_height, + latent_width=refine_plan.stage2_latent_width, + ) + print( + f"[refine] stage1={refine_plan.stage1_width}x{refine_plan.stage1_height} " + f"-> stage2={refine_plan.target_width}x{refine_plan.target_height} " + f"(scale={refine_plan.spatial_scale}x, mode={args.refine_upsample_mode})" + ) + two_pass = run_two_pass_dmd( + dit=dit, + encoder_hidden_states=encoder_hidden_states, + noise_latents_stage1=latents, + freqs_cis_stage1=freqs_cis, + freqs_cis_stage2=freqs_cis_stage2, + plan=refine_plan, + schedule=dmd_schedule, + timesteps=[float(t.item()) for t in timesteps], + refine_timesteps=refine_steps, + mx_dtype=mx_dtype, + seed=args.seed, + add_noise_flag=args.refine_add_noise, + upsample_mode=args.refine_upsample_mode, + refine_sigma=args.refine_sigma, + ) + latents = two_pass.latents + stage1_latents_np = np.array(two_pass.stage1_latents.astype(mx.float32)) + refine_sigma = two_pass.refine_sigma + print(f"[refine] stage-2 hand-off sigma={refine_sigma:.4f} " + f"(stage-1 weight {1.0 - refine_sigma:.4f})") + del two_pass, freqs_cis_stage2 + else: + for step_index, timestep in enumerate(timesteps): + noise_input_latent = latents + timestep_mx = mx.array([float(timestep.item())]).astype(mx.float32) + noise_pred = dit(latents.astype(mx_dtype), encoder_hidden_states, timestep_mx, freqs_cis) + + if args.denoising_mode == "dmd": + # On-device DMD update: no per-step MLX->torch->MLX round-trip. The + # affine math runs in fp32 to match the torch reference precision, + # then casts back to the runtime dtype. Re-noise is drawn with MLX's + # RNG (seeded above) instead of the torch CPU generator. + ts_val = float(timestep.item()) + noise_input_f32 = noise_input_latent.astype(mx.float32) + pred_noise_f32 = noise_pred.astype(mx.float32) + if step_index < len(timesteps) - 1: + next_ts: float | None = float(timesteps[step_index + 1].item()) + renoise = mx.random.normal(noise_input_f32.shape).astype(mx.float32) + else: + next_ts, renoise = None, None + latents = dmd_step( + latents=noise_input_f32, + noise_input_latent=noise_input_f32, + pred_noise=pred_noise_f32, + schedule=dmd_schedule, + timestep=ts_val, + next_timestep=next_ts, + noise=renoise, + ).astype(mx_dtype) + else: + mx.eval(noise_pred) + noise_pred_torch = torch.from_numpy(np.array(noise_pred.astype(mx.float32))) + latents_torch = torch.from_numpy(np.array(latents.astype(mx.float32))) + latents_torch = scheduler.step(noise_pred_torch, timestep, latents_torch, return_dict=False)[0] + latents = mx.array(latents_torch.numpy()).astype(mx_dtype) + + mx.eval(latents) + print(f"denoise step {step_index + 1}/{len(timesteps)} complete") + if args.denoising_mode == "dmd": + del noise_input_f32, pred_noise_f32, renoise + else: + del noise_pred_torch, latents_torch + del noise_input_latent, noise_pred, timestep_mx + + # A: spatial fast mode leaves the latents alone. They stay on the stage-1 + # grid through decode, and the resample to the target size happens on the + # decoded frames in _postprocess_video — interpolating a Wan latent puts it + # off the decoder's manifold and returns a blurred veil. + denoise_time = time.perf_counter() - denoise_start + denoise_peak_memory = mx.get_peak_memory() + active_memory = mx.get_active_memory() + + latents_np = np.array(latents.astype(mx.float32)) + if args.save_latents: + latent_path = args.output_path.with_suffix(".latents.npy") + latent_path.parent.mkdir(parents=True, exist_ok=True) + np.save(latent_path, latents_np) + print(f"Saved latents to: {latent_path}") + if args.save_stage1_latents and stage1_latents_np is not None: + stage1_path = args.output_path.with_name(args.output_path.stem + ".stage1.latents.npy") + stage1_path.parent.mkdir(parents=True, exist_ok=True) + np.save(stage1_path, stage1_latents_np) + print(f"Saved stage-1 latents to: {stage1_path}") + + del dit, latents, encoder_hidden_states, freqs_cis + cleanup_mlx() + + decode_start = time.perf_counter() + decode_latents_to_video( + model_root=model_root, + latents_np=latents_np, + output_path=args.output_path, + fps=args.fps, + device_arg=args.torch_device, + dtype_arg=(args.vae_decode_dtype if args.decode_backend == "wan-vae" else args.torch_dtype), + backend=args.decode_backend, + taehv_source_path=args.taehv_source_path, + taehv_checkpoint_path=args.taehv_checkpoint_path, + taehv_parallel=args.taehv_parallel, + ) + decode_time = time.perf_counter() - decode_start + + # Post-decode passes share one read/write so the clip takes a single h264 + # round-trip even when --fast and --fast-spatial are combined. + postprocess_time = 0.0 + rife_request = None + if fast_target_frames is not None: + rife_request = { + "factor": args.fast_factor, + "target_frames": fast_target_frames, + "sharpen": args.fast_sharpen, + } + spatial_request = fast_spatial_plan if spatial_mode == "fast_spatial" else None + if rife_request is not None or spatial_request is not None: + postprocess_start = time.perf_counter() + _postprocess_video( + video_path=args.output_path, + fps=args.fps, + rife=rife_request, + spatial=spatial_request, + ) + postprocess_time = time.perf_counter() - postprocess_start + print(f"Post-decode (RIFE/upsample) time: {postprocess_time:.2f}s") + rife_time = postprocess_time + + total_time = time.perf_counter() - total_start + + if enhance_result is not None: + print(f"Prompt enhance time: {enhance_time:.2f}s ({enhance_result.backend})") + print(f"Prompt encode time: {prompt_time:.2f}s") + print(f"MLX DiT load time: {load_time:.2f}s") + print(f"MLX denoise time: {denoise_time:.2f}s") + print(f"Decode/export time: {decode_time:.2f}s") + if postprocess_time: + print(f"Post-decode time: {postprocess_time:.2f}s") + print(f"Total prompt-to-video time: {total_time:.2f}s") + print(f"MLX load peak memory: {load_peak_memory / (1024 ** 3):.2f} GiB") + print(f"MLX denoise peak memory: {denoise_peak_memory / (1024 ** 3):.2f} GiB") + print(f"MLX active memory after denoise: {active_memory / (1024 ** 3):.2f} GiB") + print(f"Spatial mode: {spatial_mode}") + print(f"Output written to: {args.output_path}") + + if args.metrics_json is not None: + metrics = { + "prompt": prompt_for_encode, + "prompt_user": args.prompt, + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "num_frames_target": fast_target_frames if fast_target_frames is not None else args.num_frames, + "denoising_mode": args.denoising_mode, + "dmd_denoising_steps": [int(step.strip()) for step in args.dmd_denoising_steps.split(",") if step.strip()], + "spatial_mode": spatial_mode, + "fast": bool(args.fast), + "fast_factor": args.fast_factor if args.fast else None, + "fast_spatial": spatial_mode == "fast_spatial", + "fast_spatial_scale": fast_spatial_plan.scale if spatial_mode == "fast_spatial" else 1, + "fast_spatial_stage1_height": fast_spatial_plan.stage1_height if spatial_mode == "fast_spatial" else None, + "fast_spatial_stage1_width": fast_spatial_plan.stage1_width if spatial_mode == "fast_spatial" else None, + "fast_spatial_upsample_mode": fast_spatial_plan.upsample_mode if spatial_mode == "fast_spatial" else None, + "fast_spatial_sharpen": fast_spatial_plan.sharpen if spatial_mode == "fast_spatial" else None, + "refine": spatial_mode == "refine", + "refine_scale": refine_plan.spatial_scale if spatial_mode == "refine" else 1, + "refine_stage1_height": refine_plan.stage1_height if spatial_mode == "refine" else None, + "refine_stage1_width": refine_plan.stage1_width if spatial_mode == "refine" else None, + "refine_sigma": refine_sigma, + "refine_upsample_mode": args.refine_upsample_mode if spatial_mode == "refine" else None, + "refine_add_noise": args.refine_add_noise if spatial_mode == "refine" else None, + "mlx_dtype": args.mlx_dtype, + "mlx_quantization": args.mlx_quantization, + "mlx_compile": args.mlx_compile, + "text_encoder_dtype": args.text_encoder_dtype, + "vae_decode_dtype": args.vae_decode_dtype if args.decode_backend == "wan-vae" else None, + "model_root": str(model_root), + "decode_backend": args.decode_backend, + "taehv_parallel": args.taehv_parallel if args.decode_backend == "taehv" else None, + "prompt_encode_mode": args.prompt_encode_mode, + "prompt_embeds_cache": str(args.prompt_embeds_cache) if args.prompt_embeds_cache else None, + "prompt_enhance_s": enhance_time, + "prompt_encode_s": prompt_time, + "mlx_dit_load_s": load_time, + "mlx_denoise_s": denoise_time, + "vae_decode_export_s": decode_time, + "decode_export_s": decode_time, + "rife_interpolate_s": rife_time, + "postprocess_s": postprocess_time, + "total_s": total_time, + **enhance_result_as_metrics(enhance_result), + "mlx_load_peak_bytes": int(load_peak_memory), + "mlx_denoise_peak_bytes": int(denoise_peak_memory), + "mlx_active_after_denoise_bytes": int(active_memory), + "output_path": str(args.output_path), + **runtime_limits, + } + args.metrics_json.parent.mkdir(parents=True, exist_ok=True) + args.metrics_json.write_text(json.dumps(metrics, indent=2)) + print(f"Metrics written to: {args.metrics_json}") + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/mlx_wan_quant_benchmark.py b/examples/inference/basic/mlx_wan_quant_benchmark.py new file mode 100644 index 000000000..803045319 --- /dev/null +++ b/examples/inference/basic/mlx_wan_quant_benchmark.py @@ -0,0 +1,365 @@ +"""Benchmark MLX FastWan quantization modes with one shared prompt encode.""" + +from __future__ import annotations + +import argparse +import json +import time +from pathlib import Path +from typing import cast + +import numpy as np + +from examples.inference.basic.mlx_wan_prompt_to_video import ( + DEFAULT_MODEL_ROOT, + decode_latents_to_video, + encode_prompt, + make_rotary_embeddings, +) +from fastvideo.mlx_runtime.memory import cleanup_mlx + + +def _parse_modes(raw: str) -> list[str]: + """ + Parse and validate a comma-separated list of quantization modes. + + Parameters: + raw (str): Comma-separated mode names. + + Returns: + list[str]: Normalized, whitespace-trimmed mode names. + + Raises: + ValueError: If any mode is unsupported. + """ + modes = [mode.strip() for mode in raw.split(",") if mode.strip()] + allowed = {"none", "int8", "int4", "mxfp8", "mxfp4", "nvfp4"} + unknown = sorted(set(modes) - allowed) + if unknown: + raise ValueError(f"Unsupported modes: {unknown}") + return modes + + +def _latent_delta_metrics(candidate: np.ndarray, baseline: np.ndarray) -> dict[str, float]: + """ + Compare candidate and baseline latent arrays using error and signal-quality metrics. + + Parameters: + candidate (np.ndarray): Latent array to evaluate. + baseline (np.ndarray): Reference latent array for comparison. + + Returns: + dict[str, float]: Mean squared error, mean absolute error, maximum absolute + error, and signal-to-noise ratio in decibels between the arrays. + """ + diff = candidate.astype(np.float32) - baseline.astype(np.float32) + mse = float(np.mean(np.square(diff))) + mae = float(np.mean(np.abs(diff))) + max_abs = float(np.max(np.abs(diff))) + signal = float(np.mean(np.square(baseline.astype(np.float32)))) + return { + "latent_mse_vs_fp16": mse, + "latent_mae_vs_fp16": mae, + "latent_max_abs_vs_fp16": max_abs, + "latent_snr_db_vs_fp16": float(10.0 * np.log10(signal / mse)) if mse > 0 else float("inf"), + } + + +def _torch_mps_memory() -> dict[str, int | None]: + """ + Report PyTorch MPS memory statistics when PyTorch MPS is available. + + Returns: + dict[str, int | None]: A mapping of MPS memory metric names to byte counts, or `None` values when PyTorch or MPS is unavailable. + """ + try: + import torch + except ImportError: + return { + "torch_mps_current_allocated_bytes": None, + "torch_mps_driver_allocated_bytes": None, + "torch_mps_recommended_max_bytes": None, + } + if not torch.backends.mps.is_available(): + return { + "torch_mps_current_allocated_bytes": None, + "torch_mps_driver_allocated_bytes": None, + "torch_mps_recommended_max_bytes": None, + } + return { + "torch_mps_current_allocated_bytes": int(torch.mps.current_allocated_memory()), + "torch_mps_driver_allocated_bytes": int(torch.mps.driver_allocated_memory()), + "torch_mps_recommended_max_bytes": int(torch.mps.recommended_max_memory()), + } + + +def _decode_with_metrics(*, args, latents: np.ndarray, output_path: Path) -> dict[str, float | int | None | str]: + """ + Decode latents to a video and collect export timing and PyTorch MPS memory metrics. + + Parameters: + args: Configuration values for decoding and video export. + latents (np.ndarray): Latent representation to decode. + output_path (Path): Destination path for the exported video. + + Returns: + dict[str, float | int | None | str]: Video export duration and PyTorch MPS memory measurements. + """ + before = _torch_mps_memory() + decode_start = time.perf_counter() + decode_latents_to_video( + model_root=args.model_root, + latents_np=latents, + output_path=output_path, + fps=args.fps, + device_arg=args.torch_device, + dtype_arg=args.torch_dtype, + backend=args.decode_backend, + taehv_source_path=args.taehv_source_path, + taehv_checkpoint_path=args.taehv_checkpoint_path, + taehv_parallel=args.taehv_parallel, + ) + decode_time = time.perf_counter() - decode_start + after = _torch_mps_memory() + return { + "decode_export_s": decode_time, + "decode_torch_mps_current_before_bytes": before["torch_mps_current_allocated_bytes"], + "decode_torch_mps_current_after_bytes": after["torch_mps_current_allocated_bytes"], + "decode_torch_mps_driver_before_bytes": before["torch_mps_driver_allocated_bytes"], + "decode_torch_mps_driver_after_bytes": after["torch_mps_driver_allocated_bytes"], + "decode_torch_mps_recommended_max_bytes": after["torch_mps_recommended_max_bytes"], + } + + +def _run_one_mode( + *, + mode: str, + args, + config: dict, + checkpoint_path: Path, + config_path: Path, + prompt_embeds, + freqs_cis, +): + """ + Run denoising for one quantization mode and collect performance and memory metrics. + + Parameters: + mode (str): Quantization mode to benchmark. + args: Benchmark configuration, including dtype, dimensions, seed, scheduler, and denoising settings. + config (dict): Model configuration containing the input channel count. + checkpoint_path (Path): Path to the transformer checkpoint. + config_path (Path): Path to the transformer configuration. + prompt_embeds: Encoded prompt embeddings shared across benchmark modes. + freqs_cis: Rotary positional embeddings used during denoising. + + Returns: + dict: The mode name, generated latent array, and metrics for model loading, + denoising, step timing, and MLX memory usage. + """ + import mlx.core as mx + import torch + + from fastvideo.benchmarks.mlx_fastwan_bench import denoise_dmd_on_device + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors + from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step + + mx_dtype = mx.float16 if args.mlx_dtype == "fp16" else mx.float32 + quantization = None if mode == "none" else mode + latent_frames = (args.num_frames - 1) // 4 + 1 + latent_height = args.height // 8 + latent_width = args.width // 8 + + load_start = time.perf_counter() + mx.clear_cache() + mx.reset_peak_memory() + dit = mlx_dit_from_diffusers_safetensors( + checkpoint_path, + config_path, + dtype=args.mlx_dtype, + quantization=quantization, + ) + load_time = time.perf_counter() - load_start + load_peak_memory = mx.get_peak_memory() + + scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift) + schedule = MLXDMDSchedule.from_torch_scheduler(scheduler) + timesteps = [int(step.strip()) for step in args.dmd_denoising_steps.split(",") if step.strip()] + # Same torch generator sequence as the original host-round-trip loop + # (initial latents first, then one re-noise draw per intermediate step), + # so every mode still shares identical stochasticity. + generator = torch.Generator(device="cpu").manual_seed(args.seed) + latents_seed = torch.randn( + (1, int(config["in_channels"]), latent_frames, latent_height, latent_width), + generator=generator, + dtype=torch.float32, + ).numpy() + renoise_by_step = [ + torch.randn(latents_seed.shape, generator=generator, dtype=torch.float32).numpy() + for _ in range(max(0, len(timesteps) - 1)) + ] + latents = mx.array(latents_seed).astype(mx_dtype) + encoder_hidden_states = mx.array(prompt_embeds.numpy()).astype(mx_dtype) + + denoise_start = time.perf_counter() + mx.reset_peak_memory() + latents_np, step_times = denoise_dmd_on_device( + mx=mx, + dit=dit, + latents=latents, + encoder_hidden_states=encoder_hidden_states, + freqs_cis=freqs_cis, + timesteps=timesteps, + renoise_by_step=renoise_by_step, + schedule=schedule, + dmd_step=dmd_step, + mx_dtype=mx_dtype, + ) + denoise_time = time.perf_counter() - denoise_start + denoise_peak_memory = mx.get_peak_memory() + active_memory = mx.get_active_memory() + return { + "mode": mode, + "latents": latents_np, + "metrics": { + "mlx_dit_load_s": load_time, + "mlx_denoise_s": denoise_time, + "mlx_denoise_first_step_s": step_times[0] if step_times else None, + "mlx_load_peak_bytes": int(load_peak_memory), + "mlx_denoise_peak_bytes": int(denoise_peak_memory), + "mlx_active_after_denoise_bytes": int(active_memory), + }, + } + + +def main() -> None: + """ + Run the MLX FastWan quantization benchmark for the selected modes and write latency, memory, output, and latent-difference metrics to the output directory. + """ + parser = argparse.ArgumentParser(description="Benchmark MLX FastWan quantization modes.") + parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT) + parser.add_argument("--prompt", default="A snow leopard walks across a windy mountain ridge.") + parser.add_argument("--height", type=int, default=192) + parser.add_argument("--width", type=int, default=320) + parser.add_argument("--num-frames", type=int, default=17) + parser.add_argument("--dmd-denoising-steps", default="1000,757,522") + parser.add_argument("--flow-shift", type=float, default=8.0) + parser.add_argument("--max-sequence-length", type=int, default=256) + parser.add_argument("--seed", type=int, default=1024) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--torch-device", default="auto") + parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16") + parser.add_argument("--mlx-dtype", choices=("fp16", "fp32"), default="fp16") + parser.add_argument("--modes", default="none,int8,int4,mxfp8,mxfp4,nvfp4") + parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_quant_benchmark")) + parser.add_argument("--decode-backend", choices=("none", "wan-vae", "taehv"), default="taehv") + parser.add_argument("--taehv-source-path", type=Path, default=None) + parser.add_argument("--taehv-checkpoint-path", type=Path, default=None) + parser.add_argument("--taehv-parallel", action="store_true") + args = parser.parse_args() + + import mlx.core as mx + import torch + + mx.random.seed(args.seed) + torch.manual_seed(args.seed) + args.output_dir.mkdir(parents=True, exist_ok=True) + + config_path = args.model_root / "transformer/config.json" + checkpoint_path = args.model_root / "transformer/diffusion_pytorch_model.safetensors" + config = json.loads(config_path.read_text()) + latent_frames = (args.num_frames - 1) // 4 + 1 + latent_height = args.height // 8 + latent_width = args.width // 8 + + prompt_start = time.perf_counter() + prompt_embeds = encode_prompt( + model_root=args.model_root, + prompt=args.prompt, + max_sequence_length=args.max_sequence_length, + device_arg=args.torch_device, + dtype_arg=args.torch_dtype, + ) + prompt_time = time.perf_counter() - prompt_start + freqs_cis = make_rotary_embeddings( + config, + latent_frames=latent_frames, + latent_height=latent_height, + latent_width=latent_width, + ) + + from fastvideo.mlx_runtime.fastwan import UnsupportedMLXQuantizationError + + baseline_latents = None + rows = [] + for mode in _parse_modes(args.modes): + print(f"=== MLX quant mode: {mode} ===") + mode_start = time.perf_counter() + try: + result = _run_one_mode( + mode=mode, + args=args, + config=config, + checkpoint_path=checkpoint_path, + config_path=config_path, + prompt_embeds=prompt_embeds, + freqs_cis=freqs_cis, + ) + except UnsupportedMLXQuantizationError as exc: + print(f"skipping mode (unsupported by this MLX build): {exc}") + rows.append({"mode": mode, "status": "unsupported_by_mlx", "error": str(exc)}) + continue + cleanup_mlx(mx) + latents = result["latents"] + if baseline_latents is None: + baseline_latents = latents + latent_path = args.output_dir / f"latents_{mode}.npy" + np.save(latent_path, latents) + + decode_time = 0.0 + decode_metrics = {} + output_path = None + if args.decode_backend != "none": + output_path = args.output_dir / f"video_{mode}_{args.decode_backend}_{args.height}x{args.width}x{args.num_frames}.mp4" + decode_metrics = _decode_with_metrics(args=args, latents=latents, output_path=output_path) + decode_time = cast(float, decode_metrics["decode_export_s"]) + + mode_total = time.perf_counter() - mode_start + mlx_denoise_peak_bytes = int(result["metrics"]["mlx_denoise_peak_bytes"]) + mlx_active_bytes = int(result["metrics"]["mlx_active_after_denoise_bytes"]) + metrics = { + "mode": mode, + "status": "ok", + "prompt_encode_shared_s": prompt_time, + "height": args.height, + "width": args.width, + "num_frames": args.num_frames, + "decode_backend": args.decode_backend, + "decode_export_s": decode_time, + "mode_total_excluding_shared_prompt_s": mode_total, + "mode_total_including_shared_prompt_s": mode_total + prompt_time, + "latents_path": str(latent_path), + "output_path": str(output_path) if output_path else None, + "mlx_denoise_peak_gib": mlx_denoise_peak_bytes / (1024**3), + "mlx_active_after_denoise_gib": mlx_active_bytes / (1024**3), + "mlx_dit_peak_under_16gb": mlx_denoise_peak_bytes < 16 * 1024**3, + "mlx_dit_active_under_16gb": mlx_active_bytes < 16 * 1024**3, + "mac_16gb_status": ( + "dit_memory_fits_16gb_measured_decode_separately" + if mlx_denoise_peak_bytes < 16 * 1024**3 else "dit_memory_exceeds_16gb" + ), + **result["metrics"], + **decode_metrics, + **_latent_delta_metrics(latents, baseline_latents), + } + rows.append(metrics) + print(json.dumps(metrics, indent=2)) + + metrics_path = args.output_dir / "metrics.json" + metrics_path.write_text(json.dumps(rows, indent=2)) + print(f"Wrote benchmark metrics to: {metrics_path}") + + +if __name__ == "__main__": + main() diff --git a/examples/inference/basic/mlx_wan_video_quality.py b/examples/inference/basic/mlx_wan_video_quality.py new file mode 100644 index 000000000..8de563a21 --- /dev/null +++ b/examples/inference/basic/mlx_wan_video_quality.py @@ -0,0 +1,104 @@ +"""Compare generated MP4s against a reference MP4 with simple pixel metrics.""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + + +def _read_video(path: Path) -> np.ndarray: + """ + Read all frames from a video file as an RGB NumPy array. + + Parameters: + path (Path): Path to the video file. + + Returns: + np.ndarray: Video frames stacked along the first axis. + + Raises: + ValueError: If the video contains no readable frames. + """ + import cv2 + + cap = cv2.VideoCapture(str(path)) + frames = [] + try: + while True: + ok, frame_bgr = cap.read() + if not ok: + break + frame_rgb = cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB) + frames.append(frame_rgb) + finally: + cap.release() + if not frames: + raise ValueError(f"No frames read from {path}") + return np.stack(frames, axis=0) + + +def _metrics(candidate: np.ndarray, reference: np.ndarray) -> dict[str, float | int | list[int]]: + """ + Compute pixel-level comparison metrics between candidate and reference video frames. + + Parameters: + candidate (np.ndarray): Candidate video frames in frame, height, width, and channel order. + reference (np.ndarray): Reference video frames with the same shape as the candidate. + + Returns: + dict[str, float | int | list[int]]: Frame dimensions and pixel comparison metrics, including MSE, MAE, maximum absolute difference, and PSNR in decibels. + + Raises: + ValueError: If the candidate and reference arrays have different shapes. + """ + if candidate.shape != reference.shape: + raise ValueError(f"Shape mismatch: candidate={candidate.shape}, reference={reference.shape}") + candidate_f = candidate.astype(np.float32) + reference_f = reference.astype(np.float32) + diff = candidate_f - reference_f + mse = float(np.mean(np.square(diff))) + mae = float(np.mean(np.abs(diff))) + max_abs = float(np.max(np.abs(diff))) + psnr = float(20.0 * np.log10(255.0 / np.sqrt(mse))) if mse > 0 else float("inf") + return { + "frames": int(candidate.shape[0]), + "height": int(candidate.shape[1]), + "width": int(candidate.shape[2]), + "channels": int(candidate.shape[3]), + "mse_vs_reference": mse, + "mae_vs_reference": mae, + "max_abs_vs_reference": max_abs, + "psnr_db_vs_reference": psnr, + } + + +def main() -> None: + """Compare candidate MP4 videos with a reference and write pixel-level metrics to a JSON file.""" + parser = argparse.ArgumentParser(description="Compare MP4s against a reference MP4.") + parser.add_argument("--reference", type=Path, required=True) + parser.add_argument("--candidates", type=Path, nargs="+", required=True) + parser.add_argument("--metrics-json", type=Path, required=True) + args = parser.parse_args() + + reference = _read_video(args.reference) + rows = [] + for candidate_path in args.candidates: + candidate = _read_video(candidate_path) + row = { + "reference_path": str(args.reference), + "candidate_path": str(candidate_path), + **_metrics(candidate, reference), + } + rows.append(row) + print(json.dumps(row, indent=2)) + + args.metrics_json.parent.mkdir(parents=True, exist_ok=True) + args.metrics_json.write_text(json.dumps(rows, indent=2)) + print(f"Wrote video quality metrics to: {args.metrics_json}") + + +if __name__ == "__main__": + main() diff --git a/fastvideo/benchmarks/__init__.py b/fastvideo/benchmarks/__init__.py new file mode 100644 index 000000000..988131360 --- /dev/null +++ b/fastvideo/benchmarks/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/benchmarks/eval_metalfx_rife.py b/fastvideo/benchmarks/eval_metalfx_rife.py new file mode 100644 index 000000000..833f040e7 --- /dev/null +++ b/fastvideo/benchmarks/eval_metalfx_rife.py @@ -0,0 +1,375 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Benchmark generate-fewer-frames plus MLX RIFE interpolation. + +This script keeps the video diffusion path delegated to +``fastvideo.benchmarks.mlx_fastwan_bench._generate_cell``. It only orchestrates +two frame-count cells and the postprocess interpolation step. +""" + +from __future__ import annotations + +import argparse +import json +import shutil +import time +from pathlib import Path +from types import SimpleNamespace + +import imageio.v2 as imageio +import numpy as np + +from examples.inference.basic.mlx_wan_prompt_to_video import encode_prompt, make_rotary_embeddings +from fastvideo.benchmarks.mlx_fastwan_bench import _generate_cell, _ms_ssim +from fastvideo.mlx_runtime.rife_interp import RIFEBackendError, interpolate, load_model +from fastvideo.mlx_runtime.memory import add_memory_limit_args, apply_memory_limits + +FOX_PROMPT = "A fox runs through a misty pine forest, leaves kicking up behind it." +DEFAULT_MODEL_ROOT = Path("/Users/aryank/models/qad_int8_v2") + + +def _parse_timesteps(raw: str) -> list[int]: + timesteps = [int(part.strip()) for part in raw.split(",") if part.strip()] + if not timesteps: + raise SystemExit("No DMD timesteps parsed from --dmd-denoising-steps") + return timesteps + + +def _read_video(path: Path) -> list[np.ndarray]: + if not path.is_file(): + raise FileNotFoundError(f"Video does not exist: {path}") + frames = [np.asarray(frame[:, :, :3], dtype=np.uint8) for frame in imageio.mimread(path)] + if not frames: + raise RuntimeError(f"No frames decoded from {path}") + return frames + + +def _write_video(path: Path, frames: list[np.ndarray], fps: int) -> None: + if not frames: + raise ValueError("Cannot write an empty frame list") + path.parent.mkdir(parents=True, exist_ok=True) + with imageio.get_writer(str(path), fps=fps, macro_block_size=1, codec="libx264", quality=8) as writer: + for frame in frames: + writer.append_data(np.asarray(frame, dtype=np.uint8)) + + +def _copy_video(src: Path, dst: Path) -> Path: + dst.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(src, dst) + return dst + + +def _make_generation_inputs(args, num_frames: int): + import mlx.core as mx + import torch + + config_path = args.model_root / "transformer" / "config.json" + checkpoint_path = args.model_root / "transformer" / "diffusion_pytorch_model.safetensors" + if not config_path.is_file(): + raise SystemExit(f"Missing DiT config: {config_path}") + if not checkpoint_path.is_file(): + raise SystemExit(f"Missing DiT checkpoint: {checkpoint_path}") + + config = json.loads(config_path.read_text()) + latent_frames = (num_frames - 1) // 4 + 1 + latent_height = args.height // 8 + latent_width = args.width // 8 + freqs_cis = make_rotary_embeddings( + config, + latent_frames=latent_frames, + latent_height=latent_height, + latent_width=latent_width, + ) + + generator = torch.Generator(device="cpu").manual_seed(args.seed) + latents_seed = torch.randn( + (1, int(config["in_channels"]), latent_frames, latent_height, latent_width), + generator=generator, + dtype=torch.float32, + ).numpy() + timesteps = _parse_timesteps(args.dmd_denoising_steps) + renoise_by_step = [ + torch.randn(latents_seed.shape, generator=generator, dtype=torch.float32).numpy() + for _ in range(max(0, + len(timesteps) - 1)) + ] + prompt_embeds = encode_prompt( + model_root=args.model_root, + prompt=args.prompt, + max_sequence_length=args.max_sequence_length, + device_arg=args.torch_device, + dtype_arg=args.torch_dtype, + ) + return { + "checkpoint_path": checkpoint_path, + "config_path": config_path, + "encoder_hidden_states": mx.array(prompt_embeds.numpy()), + "freqs_cis": freqs_cis, + "timesteps": timesteps, + "latents_seed": latents_seed, + "renoise_by_step": renoise_by_step, + "latent_frames": latent_frames, + } + + +def _generate(args, num_frames: int): + cell_args = SimpleNamespace( + model_root=args.model_root, + height=args.height, + width=args.width, + num_frames=num_frames, + fps=args.fps, + flow_shift=args.flow_shift, + torch_device=args.torch_device, + torch_dtype=args.torch_dtype, + taehv_source_path=args.taehv_source_path, + taehv_checkpoint_path=args.taehv_checkpoint_path, + taehv_parallel=args.taehv_parallel, + mlx_checkpoint_cache=args.mlx_checkpoint_cache, + mlx_memory_limit_gib=args.mlx_memory_limit_gib, + mlx_cache_limit_gib=args.mlx_cache_limit_gib, + mlx_disable_cache=args.mlx_disable_cache, + mlx_wired_limit_gib=args.mlx_wired_limit_gib, + torch_mps_high_watermark_ratio=args.torch_mps_high_watermark_ratio, + torch_mps_low_watermark_ratio=args.torch_mps_low_watermark_ratio, + benchmark_preset="metalfx-rife", + current_prompt_id=f"fox-{num_frames}f", + current_prompt=args.prompt, + output_dir=args.output_dir, + ) + inputs = _make_generation_inputs(args, num_frames) + print(f"=== generate: {num_frames} frames ({inputs['latent_frames']} latent frames) ===", flush=True) + return _generate_cell( + args=cell_args, + mode=args.mode, + decoder=args.decoder, + checkpoint_path=inputs["checkpoint_path"], + config_path=inputs["config_path"], + encoder_hidden_states=inputs["encoder_hidden_states"], + freqs_cis=inputs["freqs_cis"], + timesteps=inputs["timesteps"], + latents_seed=inputs["latents_seed"], + renoise_by_step=inputs["renoise_by_step"], + ) + + +def _relative_to_output(path: Path, output_dir: Path) -> str: + try: + return str(path.relative_to(output_dir)) + except ValueError: + return str(path) + + +def _write_report(args, result: dict) -> Path: + report_path = args.output_dir / "metalfx_rife_report.md" + rows = result["speed_rows"] + table_lines = [ + "| path | denoise_s | decode_s | gen_total_s | rife_s | net_s | speedup_vs_81 |", + "| --- | ---: | ---: | ---: | ---: | ---: | ---: |", + ] + for row in rows: + table_lines.append( + "| {path} | {denoise_s:.3f} | {decode_s:.3f} | {gen_total_s:.3f} | {rife_s:.3f} | {net_s:.3f} | {speedup_vs_81:.3f}x |" + .format(**row)) + text = f"""# MetalFX/RIFE Generate-Fewer-Frames Benchmark Run + +Prompt: {args.prompt} + +Resolution: {args.height}x{args.width}, fps={args.fps}, mode={args.mode}, decoder={args.decoder} + +| metric | value | +| --- | ---: | +| reconstruction_ms_ssim | {result['reconstruction_ms_ssim']:.6f} | +| reference_frames | {result['reference_frames']} | +| reduced_frames | {result['reduced_frames']} | +| interpolated_frames | {result['interpolated_frames']} | + +{chr(10).join(table_lines)} + +Videos: + +- reference: `{result['videos']['reference']}` +- reference drop-41 RIFE reconstruction: `{result['videos']['drop41_rife81']}` +- generated 41 RIFE to 81: `{result['videos']['generated41_rife81']}` +- generated 41 direct: `{result['videos']['generated41']}` + +Raw metrics are in `{_relative_to_output(args.output_dir / 'metrics.json', args.output_dir)}`. +""" + report_path.write_text(text) + return report_path + + +def main() -> None: + parser = argparse.ArgumentParser( + description="Evaluate 41-frame generation plus MLX RIFE interpolation vs 81-frame generation.") + parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT) + parser.add_argument("--prompt", default=FOX_PROMPT) + parser.add_argument("--height", type=int, default=480) + parser.add_argument("--width", type=int, default=832) + parser.add_argument("--reference-frames", type=int, default=81) + parser.add_argument("--reduced-frames", type=int, default=41) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--seed", type=int, default=1024) + parser.add_argument("--mode", default="int8", choices=("fp16", "bf16", "int8", "int4", "mxfp8", "mxfp4", "nvfp4")) + parser.add_argument("--decoder", default="taehv", choices=("taehv", "wan-vae")) + parser.add_argument("--dmd-denoising-steps", default="1000,757,522") + parser.add_argument("--flow-shift", type=float, default=8.0) + parser.add_argument("--max-sequence-length", type=int, default=512) + parser.add_argument("--torch-device", default="auto") + parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16") + parser.add_argument("--output-dir", type=Path, default=Path("bench/metalfx_rife")) + parser.add_argument("--mlx-checkpoint-cache", type=Path, default=None) + parser.add_argument("--compile", action="store_true", help="Enable FASTVIDEO_MLX_COMPILE=1 for DiT denoise.") + parser.add_argument("--rife-scale", type=float, default=1.0) + parser.add_argument("--taehv-source-path", type=Path, default=None) + parser.add_argument("--taehv-checkpoint-path", type=Path, default=None) + parser.add_argument("--taehv-parallel", action="store_true") + add_memory_limit_args(parser) + args = parser.parse_args() + + if args.reference_frames != (args.reduced_frames - 1) * 2 + 1: + raise SystemExit( + "--reference-frames must equal (--reduced-frames - 1) * 2 + 1 for the default every-other-frame test") + if args.compile: + import os + + os.environ["FASTVIDEO_MLX_COMPILE"] = "1" + + args.model_root = args.model_root.expanduser().resolve() + args.output_dir = args.output_dir.expanduser().resolve() + args.output_dir.mkdir(parents=True, exist_ok=True) + if args.mlx_checkpoint_cache is None: + args.mlx_checkpoint_cache = args.output_dir / "mlx_checkpoint_cache" + else: + args.mlx_checkpoint_cache = args.mlx_checkpoint_cache.expanduser().resolve() + + import mlx.core as mx + import torch + + runtime_limits = apply_memory_limits( + mlx_memory_limit_gib=args.mlx_memory_limit_gib, + mlx_cache_limit_gib=args.mlx_cache_limit_gib, + mlx_disable_cache=args.mlx_disable_cache, + mlx_wired_limit_gib=args.mlx_wired_limit_gib, + torch_mps_high_watermark_ratio=args.torch_mps_high_watermark_ratio, + torch_mps_low_watermark_ratio=args.torch_mps_low_watermark_ratio, + mx_module=mx, + ).as_metrics() + mx.random.seed(args.seed) + torch.manual_seed(args.seed) + + reference_cell = _generate(args, args.reference_frames) + reduced_cell = _generate(args, args.reduced_frames) + + reference_video = _copy_video(reference_cell.video_path, args.output_dir / "fox_reference_81.mp4") + generated41_video = _copy_video(reduced_cell.video_path, args.output_dir / "fox_generated_41.mp4") + + reference_frames = _read_video(reference_video) + dropped_reference_frames = reference_frames[::2] + if len(dropped_reference_frames) != args.reduced_frames: + raise RuntimeError(f"Expected {args.reduced_frames} dropped frames, got {len(dropped_reference_frames)}") + + try: + rife_model = load_model("4.25") + except RIFEBackendError: + raise + except Exception as exc: # noqa: BLE001 - preserve exact backend failure. + raise RIFEBackendError(f"Unexpected RIFE load failure: {exc}") from exc + + print("=== RIFE reconstruction: drop every other reference frame, 41 -> 81 ===", flush=True) + start = time.perf_counter() + reconstructed_frames = interpolate(dropped_reference_frames, factor=2, model=rife_model, scale=args.rife_scale) + recon_rife_s = time.perf_counter() - start + reconstructed_video = args.output_dir / "fox_reference_drop41_rife81.mp4" + _write_video(reconstructed_video, reconstructed_frames, args.fps) + reconstruction_ms_ssim = _ms_ssim(reference_video, reconstructed_video, required=True) + if reconstruction_ms_ssim is None: + raise RuntimeError("MS-SSIM returned None for reconstruction comparison") + + print("=== RIFE actual reduced path: generated 41 -> 81 ===", flush=True) + generated41_frames = _read_video(generated41_video) + start = time.perf_counter() + generated41_rife_frames = interpolate(generated41_frames, factor=2, model=rife_model, scale=args.rife_scale) + generated41_rife_s = time.perf_counter() - start + generated41_rife_video = args.output_dir / "fox_generated41_rife81.mp4" + _write_video(generated41_rife_video, generated41_rife_frames, args.fps) + + direct_denoise_s = float(reference_cell.metrics["denoise_s"]) + reduced_denoise_s = float(reduced_cell.metrics["denoise_s"]) + direct_gen_total_s = float(reference_cell.metrics["denoise_s"]) + float(reference_cell.metrics["decode_s"]) + reduced_gen_total_s = float(reduced_cell.metrics["denoise_s"]) + float(reduced_cell.metrics["decode_s"]) + net_denoise_rife_s = reduced_denoise_s + generated41_rife_s + net_gen_rife_s = reduced_gen_total_s + generated41_rife_s + + result = { + "prompt": + args.prompt, + "model_root": + str(args.model_root), + "rife_impl": + "rife-mlx vendored at fastvideo/third_party/rife_mlx, weights mlx-community/RIFE-4.25", + "runtime_limits": + runtime_limits, + "reference_frames": + len(reference_frames), + "reduced_frames": + len(generated41_frames), + "interpolated_frames": + len(generated41_rife_frames), + "reconstruction_ms_ssim": + reconstruction_ms_ssim, + "reconstruction_rife_s": + recon_rife_s, + "generated41_rife_s": + generated41_rife_s, + "reference_metrics": + reference_cell.metrics, + "reduced_metrics": + reduced_cell.metrics, + "speed_rows": [ + { + "path": "generate_81", + "denoise_s": direct_denoise_s, + "decode_s": float(reference_cell.metrics["decode_s"]), + "gen_total_s": direct_gen_total_s, + "rife_s": 0.0, + "net_s": direct_gen_total_s, + "speedup_vs_81": 1.0, + }, + { + "path": "generate_41_plus_rife81_denoise_only", + "denoise_s": reduced_denoise_s, + "decode_s": 0.0, + "gen_total_s": reduced_denoise_s, + "rife_s": generated41_rife_s, + "net_s": net_denoise_rife_s, + "speedup_vs_81": direct_denoise_s / net_denoise_rife_s, + }, + { + "path": "generate_41_plus_rife81_decode_included", + "denoise_s": reduced_denoise_s, + "decode_s": float(reduced_cell.metrics["decode_s"]), + "gen_total_s": reduced_gen_total_s, + "rife_s": generated41_rife_s, + "net_s": net_gen_rife_s, + "speedup_vs_81": direct_gen_total_s / net_gen_rife_s, + }, + ], + "videos": { + "reference": _relative_to_output(reference_video, args.output_dir), + "drop41_rife81": _relative_to_output(reconstructed_video, args.output_dir), + "generated41": _relative_to_output(generated41_video, args.output_dir), + "generated41_rife81": _relative_to_output(generated41_rife_video, args.output_dir), + }, + } + metrics_path = args.output_dir / "metrics.json" + metrics_path.write_text(json.dumps(result, indent=2)) + report_path = _write_report(args, result) + + print(json.dumps(result["speed_rows"], indent=2)) + print(f"reconstruction_ms_ssim={reconstruction_ms_ssim:.6f}") + print(f"wrote {metrics_path}") + print(f"wrote {report_path}") + + +if __name__ == "__main__": + main() diff --git a/fastvideo/benchmarks/mlx_fastwan_bench.py b/fastvideo/benchmarks/mlx_fastwan_bench.py new file mode 100644 index 000000000..15a096d3f --- /dev/null +++ b/fastvideo/benchmarks/mlx_fastwan_bench.py @@ -0,0 +1,826 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Prove-out benchmark for the MLX FastWan runtime (Apple Silicon). + +Sweeps ``{dtype/quant} x {decoder}``, generates a clip per cell, and records the +latency breakdown, peak unified memory, and MS-SSIM (optionally LPIPS) against a +reference video. It emits a JSON blob and a markdown table -- the artifact that +turns "int8 + TAEHV looks good" into defensible numbers, and (via +``--assert-min-ssim``) a regression gate for the ``mx.compile`` work. + +Design notes: +- Generation reuses the hybrid POC helpers in + ``examples/inference/basic/mlx_wan_prompt_to_video.py`` (torch-MPS UMT5 encode + and Wan-VAE/TAEHV decode) plus the on-device MLX DMD sampler + (``fastvideo/mlx_runtime/sampling.py``); the denoise loop never leaves the + device. +- Quality reuses the tested MS-SSIM primitive + ``fastvideo/tests/utils.py::compute_video_ssim_torchvision``. +- Reference: by default each cell is scored against the highest-fidelity cell + in the sweep (``fp16`` + ``wan-vae``), which needs no CUDA box and answers + "how much does int8/int4/TAEHV degrade vs the best local config". Pass + ``--reference PATH`` to score against an external clip instead (e.g. the + torch-MPS or CUDA FastVideo output of the same model) for a "vs. the original + model" column. + +Run on an Apple Silicon Mac (needs ``mlx`` + a torch build with MPS): + + python fastvideo/benchmarks/mlx_fastwan_bench.py \ + --modes fp16,bf16,int8,int4 --decoders taehv,wan-vae +""" + +from __future__ import annotations + +import argparse +import html +import json +import os +import time +from dataclasses import dataclass, field +from pathlib import Path + +import numpy as np + +from examples.inference.basic.mlx_wan_prompt_to_video import ( + DEFAULT_MODEL_ROOT, + decode_latents_to_video, + encode_prompt, + make_rotary_embeddings, +) +from fastvideo.mlx_runtime.memory import add_memory_limit_args, apply_memory_limits, cleanup_mlx + +# The highest-fidelity cell; used as the default SSIM reference when no external +# reference video is supplied. +REFERENCE_MODE = "fp16" +REFERENCE_DECODER = "wan-vae" + +ALLOWED_MODES = ("fp16", "bf16", "int8", "int4", "mxfp8", "mxfp4", "nvfp4") +ALLOWED_DECODERS = ("taehv", "wan-vae") + + +@dataclass(frozen=True) +class PromptCase: + id: str + prompt: str + + +@dataclass(frozen=True) +class BenchmarkPreset: + height: int + width: int + num_frames: int + modes: str + decoders: str + mlx_memory_limit_gib: float | None = None + mlx_disable_cache: bool = False + torch_mps_high_watermark_ratio: float | None = None + torch_mps_low_watermark_ratio: float | None = None + + +PROMPT_SETS = { + "motion7": ( + PromptCase("beach-sunset", "A slow cinematic sunset over ocean waves at a quiet beach."), + PromptCase("fox-forest", "A fox runs through a misty pine forest, leaves kicking up behind it."), + PromptCase("raccoon-sunflowers", "A raccoon walks through a sunflower field as petals move in the wind."), + PromptCase("surfing-cat", "A cat wearing sunglasses surfs across a bright blue ocean wave."), + PromptCase("burning-clock", "A vintage table clock burns on a wooden desk, flames flickering realistically."), + PromptCase("forest-walk", "Video game style footage of a man walking through a dense forest path."), + PromptCase("sea-dock-yachts", "Several yachts are parked at a sea dock while water ripples around them."), + ), +} + +BENCHMARK_PRESETS = { + "default": + BenchmarkPreset( + height=480, + width=832, + num_frames=81, + modes="fp16,bf16,int8,int4", + decoders="taehv,wan-vae", + ), + "mac-16gb": + BenchmarkPreset( + height=448, + width=832, + num_frames=61, + modes="int8", + decoders="taehv", + mlx_memory_limit_gib=16.0, + mlx_disable_cache=True, + torch_mps_high_watermark_ratio=0.57, + torch_mps_low_watermark_ratio=0.0, + ), + "mac-32gb": + BenchmarkPreset( + height=480, + width=832, + num_frames=81, + modes="int8,fp16", + decoders="taehv", + ), + "mac-64gb": + BenchmarkPreset( + height=480, + width=832, + num_frames=81, + modes="int8,fp16", + decoders="taehv,wan-vae", + ), +} + + +@dataclass +class Cell: + prompt_id: str + prompt: str + mode: str + decoder: str + video_path: Path + latents: np.ndarray + metrics: dict[str, float | int | str | bool | None] = field(default_factory=dict) + + +def _mode_to_dtype_quant(mode: str) -> tuple[str, str | None]: + """Map a sweep mode to (MLX compute dtype, quantization spec). + + Quantized modes keep fp16 activations and quantize only the DiT linear + weights (matching ``mlx_dit_from_diffusers_safetensors``). + """ + if mode == "bf16": + return "bf16", None + if mode == "fp16": + return "fp16", None + # int8/int4/mxfp*/nvfp4 -> fp16 activations + quantized weights. + return "fp16", mode + + +def _mx_dtype(mx, base: str): + return {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[base] + + +def _parse_list(raw: str, allowed: tuple[str, ...], label: str) -> list[str]: + items = [x.strip() for x in raw.split(",") if x.strip()] + unknown = sorted(set(items) - set(allowed)) + if unknown: + raise ValueError(f"Unsupported {label}: {unknown} (allowed: {list(allowed)})") + return items + + +def _safe_slug(value: str, *, fallback: str) -> str: + slug = "".join(ch.lower() if ch.isalnum() else "-" for ch in value.strip()) + slug = "-".join(part for part in slug.split("-") if part) + return slug[:64] or fallback + + +def _load_prompt_cases(prompt: str, prompt_file: Path | None, prompt_set: str = "single") -> list[PromptCase]: + """Load one prompt, a built-in prompt set, or a text/jsonl prompt file.""" + if prompt_file is None: + if prompt_set == "single": + return [PromptCase(id="prompt-001", prompt=prompt)] + if prompt_set not in PROMPT_SETS: + raise ValueError(f"Unsupported prompt set: {prompt_set} (allowed: {sorted(PROMPT_SETS) + ['single']})") + return list(PROMPT_SETS[prompt_set]) + cases: list[PromptCase] = [] + for line_index, raw_line in enumerate(prompt_file.read_text().splitlines(), start=1): + line = raw_line.strip() + if not line or line.startswith("#"): + continue + prompt_id = f"prompt-{len(cases) + 1:03d}" + prompt_text = line + if prompt_file.suffix.lower() == ".jsonl": + item = json.loads(line) + prompt_text = str(item.get("prompt") or item.get("text") or item.get("caption") or "").strip() + if not prompt_text: + raise ValueError(f"{prompt_file}:{line_index} has no prompt/text/caption field") + prompt_id = str(item.get("id") or item.get("name") or prompt_id) + cases.append(PromptCase(id=_safe_slug(prompt_id, fallback=f"prompt-{len(cases) + 1:03d}"), prompt=prompt_text)) + if not cases: + raise ValueError(f"No prompts found in {prompt_file}") + return cases + + +def denoise_dmd_on_device( + *, + mx, + dit, + latents, + encoder_hidden_states, + freqs_cis, + timesteps: list[int], + renoise_by_step: list[np.ndarray], + schedule, + dmd_step, + mx_dtype, +) -> tuple[np.ndarray, list[float]]: + """Run the FastWan DMD loop entirely on the MLX device. + + Mirrors the loop in ``mlx_wan_prompt_to_video.py`` (fp32 affine math, MLX RNG + re-noise) so the benchmark measures exactly the shipped path. + + Returns the final latents plus per-step wall times. The first step carries + one-time costs (mx.compile tracing, kernel warm-up), so first-vs-steady + step timing is how the benchmark separates cold-start from steady-state + denoise throughput. + + All host-side tensors (timesteps, re-noise draws) are uploaded before the + loop starts, so the per-step body performs no bulk host->device transfers + and step timings measure device work rather than staging copies. + """ + timesteps_mx = [mx.array([float(timestep)]).astype(mx.float32) for timestep in timesteps] + renoise_mx = [mx.array(renoise).astype(mx.float32) for renoise in renoise_by_step] + if timesteps_mx or renoise_mx: + mx.eval(*timesteps_mx, *renoise_mx) + + step_times: list[float] = [] + for step_index, timestep in enumerate(timesteps): + step_start = time.perf_counter() + noise_input_latent = latents + noise_pred = dit(latents.astype(mx_dtype), encoder_hidden_states, timesteps_mx[step_index], freqs_cis) + + noise_input_f32 = noise_input_latent.astype(mx.float32) + pred_noise_f32 = noise_pred.astype(mx.float32) + if step_index < len(timesteps) - 1: + next_ts: float | None = float(timesteps[step_index + 1]) + renoise = renoise_mx[step_index] + else: + next_ts, renoise = None, None + latents = dmd_step( + latents=noise_input_f32, + noise_input_latent=noise_input_f32, + pred_noise=pred_noise_f32, + schedule=schedule, + timestep=float(timestep), + next_timestep=next_ts, + noise=renoise, + ).astype(mx_dtype) + mx.eval(latents) + step_times.append(time.perf_counter() - step_start) + return np.array(latents.astype(mx.float32)), step_times + + +def _peak_memory_bytes(mx) -> int: + try: + return int(mx.get_peak_memory()) + except Exception: # noqa: BLE001 - best-effort telemetry only. + return 0 + + +def _latent_delta_metrics(candidate: np.ndarray, baseline: np.ndarray) -> dict[str, float]: + diff = candidate.astype(np.float32) - baseline.astype(np.float32) + mse = float(np.mean(np.square(diff))) + signal = float(np.mean(np.square(baseline.astype(np.float32)))) + return { + "latent_mse_vs_ref_mode": mse, + "latent_snr_db_vs_ref_mode": float(10.0 * np.log10(signal / mse)) if mse > 0 else float("inf"), + } + + +def _ms_ssim(reference_video: Path, candidate_video: Path, *, required: bool = False) -> float | None: + """Mean MS-SSIM between two mp4s, via the repo's tested helper.""" + if not reference_video.exists() or not candidate_video.exists(): + return None + try: + from fastvideo.tests.utils import compute_video_ssim_torchvision + except ImportError as exc: + message = ("MS-SSIM is unavailable because `pytorch-msssim` is not installed. " + "Install FastVideo with the test extra, e.g. `uv pip install -e '.[mlx,test]'`, " + "or run without an SSIM assertion.") + if required: + raise RuntimeError(message) from exc + print(f"{message} Skipping MS-SSIM.") + return None + + try: + ssim_values = compute_video_ssim_torchvision(str(reference_video), str(candidate_video), use_ms_ssim=True) + except ImportError as exc: + message = ("MS-SSIM is unavailable because `pytorch-msssim` is not installed. " + "Install FastVideo with the test extra, e.g. `uv pip install -e '.[mlx,test]'`, " + "or run without an SSIM assertion.") + if required: + raise RuntimeError(message) from exc + print(f"{message} Skipping MS-SSIM.") + return None + return float(ssim_values[0]) + + +def _markdown_table(rows: list[dict]) -> str: + columns = [ + ("prompt_id", "prompt"), + ("mode", "mode"), + ("decoder", "decoder"), + ("status", "status"), + ("denoise_s", "denoise s"), + ("decode_s", "decode s"), + ("total_s", "total s"), + ("peak_gib", "peak GiB"), + ("ms_ssim_vs_ref", "MS-SSIM"), + ("lpips_vs_ref", "LPIPS"), + ] + header = "| " + " | ".join(label for _, label in columns) + " |" + sep = "| " + " | ".join("---" for _ in columns) + " |" + lines = [header, sep] + for row in rows: + cells = [] + for key, _ in columns: + value = row.get(key) + if isinstance(value, float): + cells.append(f"{value:.3f}") + elif value is None: + cells.append("-") + else: + cells.append(str(value)) + lines.append("| " + " | ".join(cells) + " |") + return "\n".join(lines) + + +def _format_metric(value) -> str: + if isinstance(value, float): + return f"{value:.3f}" + if value is None: + return "-" + return str(value) + + +def _html_grid(rows: list[dict]) -> str: + groups: dict[str, list[dict]] = {} + for row in rows: + groups.setdefault(str(row.get("prompt_id", "prompt")), []).append(row) + + sections = [] + for prompt_id, group_rows in groups.items(): + prompt = next((str(row.get("prompt", "")) for row in group_rows if row.get("prompt")), "") + cards = [] + for row in group_rows: + title = f"{row.get('mode', '-')} / {row.get('decoder', '-')}" + status = row.get("status", "-") + video_path = row.get("video_path") + if video_path: + media = f'' + else: + media = f'
No video
{html.escape(str(row.get("error", "")))}
' + metrics = (f"status={status} · total={_format_metric(row.get('total_s'))}s · " + f"denoise={_format_metric(row.get('denoise_s'))}s · " + f"decode={_format_metric(row.get('decode_s'))}s · " + f"peak={_format_metric(row.get('peak_gib'))}GiB") + cards.append("
" + f"

{html.escape(title)}

" + f"{media}" + f"

{html.escape(metrics)}

" + "
") + sections.append("
" + f"

{html.escape(prompt_id)}

" + f"

{html.escape(prompt)}

" + f"
{''.join(cards)}
" + "
") + + return """ + + + + + FastVideo MLX benchmark grid + + + +

FastVideo MLX benchmark grid

+

Use the controls below to start/stop every clip together for side-by-side inspection.

+ + +""" + "\n".join(sections) + """ + + +""" + + +def _write_html_grid(rows: list[dict], output_dir: Path) -> Path: + html_path = output_dir / "index.html" + html_path.write_text(_html_grid(rows)) + return html_path + + +def _generate_cell( + *, + args, + mode: str, + decoder: str, + checkpoint_path: Path, + config_path: Path, + encoder_hidden_states, + freqs_cis, + timesteps: list[int], + latents_seed: np.ndarray, + renoise_by_step: list[np.ndarray], +) -> Cell: + import mlx.core as mx + + from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors + from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + + base_dtype, quantization = _mode_to_dtype_quant(mode) + mx_dtype = _mx_dtype(mx, base_dtype) + + mx.clear_cache() + mx.reset_peak_memory() + load_start = time.perf_counter() + load_source = "diffusers" + if args.mlx_checkpoint_cache is not None: + from fastvideo.mlx_runtime.checkpoint import ( + load_mlx_dit_checkpoint, + save_mlx_dit_checkpoint, + ) + + mode_ckpt_dir = args.mlx_checkpoint_cache / mode + if (mode_ckpt_dir / "mlx_dit.json").exists(): + dit = load_mlx_dit_checkpoint(mode_ckpt_dir) + load_source = "mlx_checkpoint" + else: + dit = mlx_dit_from_diffusers_safetensors( + checkpoint_path, + config_path, + dtype=base_dtype, + quantization=quantization, + ) + save_mlx_dit_checkpoint(dit, mode_ckpt_dir) + load_source = "diffusers_then_saved" + else: + dit = mlx_dit_from_diffusers_safetensors( + checkpoint_path, + config_path, + dtype=base_dtype, + quantization=quantization, + ) + load_s = time.perf_counter() - load_start + load_peak = _peak_memory_bytes(mx) + + scheduler = FlowMatchEulerDiscreteScheduler(shift=args.flow_shift) + schedule = MLXDMDSchedule.from_torch_scheduler(scheduler) + + latents = mx.array(latents_seed).astype(mx_dtype) + mx.reset_peak_memory() + denoise_start = time.perf_counter() + latents_np, step_times = denoise_dmd_on_device( + mx=mx, + dit=dit, + latents=latents, + encoder_hidden_states=encoder_hidden_states.astype(mx_dtype), + freqs_cis=freqs_cis, + timesteps=timesteps, + renoise_by_step=renoise_by_step, + schedule=schedule, + dmd_step=dmd_step, + mx_dtype=mx_dtype, + ) + denoise_s = time.perf_counter() - denoise_start + denoise_peak = _peak_memory_bytes(mx) + del dit, latents + cleanup_mlx(mx) + + video_path = (args.output_dir / f"{args.current_prompt_id}" / + f"video_{mode}_{decoder}_{args.height}x{args.width}x{args.num_frames}.mp4") + decode_start = time.perf_counter() + decode_latents_to_video( + model_root=args.model_root, + latents_np=latents_np, + output_path=video_path, + fps=args.fps, + device_arg=args.torch_device, + dtype_arg=args.torch_dtype, + backend=decoder, + taehv_source_path=args.taehv_source_path, + taehv_checkpoint_path=args.taehv_checkpoint_path, + taehv_parallel=args.taehv_parallel, + ) + decode_s = time.perf_counter() - decode_start + + metrics: dict[str, float | int | str | bool | None] = { + "prompt_id": args.current_prompt_id, + "prompt": args.current_prompt, + "benchmark_preset": args.benchmark_preset, + "mode": mode, + "decoder": decoder, + "status": "ok", + "video_path": str(video_path.relative_to(args.output_dir)), + "load_s": load_s, + "load_source": load_source, + "denoise_s": denoise_s, + # The first step carries one-time costs (mx.compile tracing, kernel + # warm-up); steady-state throughput is the median of the rest. + "denoise_first_step_s": step_times[0] if step_times else None, + "denoise_steady_step_s": (float(np.median(step_times[1:])) if len(step_times) > 1 else None), + "decode_s": decode_s, + "total_s": load_s + denoise_s + decode_s, + "load_peak_gib": load_peak / (1024**3), + "peak_gib": max(load_peak, denoise_peak) / (1024**3), + "quantization": quantization or "none", + "compute_dtype": base_dtype, + "compile": os.environ.get("FASTVIDEO_MLX_COMPILE", "0") == "1", + "fast_norm": os.environ.get("FASTVIDEO_MLX_FAST_NORM", "0") == "1", + "mlx_memory_limit_gib": args.mlx_memory_limit_gib, + "mlx_cache_limit_gib": args.mlx_cache_limit_gib, + "mlx_disable_cache": args.mlx_disable_cache, + "mlx_wired_limit_gib": args.mlx_wired_limit_gib, + "torch_mps_high_watermark_ratio": args.torch_mps_high_watermark_ratio, + "torch_mps_low_watermark_ratio": args.torch_mps_low_watermark_ratio, + } + return Cell( + prompt_id=args.current_prompt_id, + prompt=args.current_prompt, + mode=mode, + decoder=decoder, + video_path=video_path, + latents=latents_np, + metrics=metrics, + ) + + +def main() -> None: + preset_parser = argparse.ArgumentParser(add_help=False) + preset_parser.add_argument("--benchmark-preset", choices=tuple(BENCHMARK_PRESETS), default="default") + preset_args, _ = preset_parser.parse_known_args() + preset = BENCHMARK_PRESETS[preset_args.benchmark_preset] + + parser = argparse.ArgumentParser(description="MLX FastWan prove-out benchmark (latency + quality).") + parser.add_argument("--benchmark-preset", + choices=tuple(BENCHMARK_PRESETS), + default=preset_args.benchmark_preset, + help="Memory-tier benchmark defaults. Explicit CLI flags override preset values.") + parser.add_argument("--model-root", type=Path, default=DEFAULT_MODEL_ROOT) + parser.add_argument("--prompt", default="A paper boat sails through a shallow stream in a mossy forest.") + parser.add_argument( + "--prompt-file", + type=Path, + default=None, + help= + "Optional prompt set. Plain text uses one prompt per non-empty line; .jsonl accepts prompt/text/caption plus optional id/name.", + ) + parser.add_argument( + "--prompt-set", + choices=("single", *PROMPT_SETS.keys()), + default="single", + help="Built-in standard prompt set. Ignored when --prompt-file is supplied.", + ) + parser.add_argument("--height", type=int, default=preset.height) + parser.add_argument("--width", type=int, default=preset.width) + parser.add_argument("--num-frames", type=int, default=preset.num_frames) + parser.add_argument("--dmd-denoising-steps", default="1000,757,522") + parser.add_argument("--flow-shift", type=float, default=8.0) + parser.add_argument("--max-sequence-length", type=int, default=512) + parser.add_argument("--seed", type=int, default=1024) + parser.add_argument("--fps", type=int, default=16) + parser.add_argument("--modes", default=preset.modes) + parser.add_argument("--decoders", default=preset.decoders) + parser.add_argument("--output-dir", type=Path, default=Path("video_samples/mlx_fastwan_bench")) + parser.add_argument("--torch-device", default="auto") + parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16") + parser.add_argument( + "--reference", + type=Path, + default=None, + help="External reference mp4 to score every cell against. Defaults to the fp16+wan-vae cell.", + ) + parser.add_argument("--assert-min-ssim", + type=float, + default=None, + help="Fail if any cell's MS-SSIM vs the reference falls below this value.") + parser.add_argument("--compile", + action="store_true", + help="Enable mx.compile on the DiT forward (sets FASTVIDEO_MLX_COMPILE=1).") + parser.add_argument("--lpips", action="store_true", help="Also compute LPIPS (needs the `lpips` package).") + parser.add_argument("--taehv-source-path", type=Path, default=None) + parser.add_argument("--taehv-checkpoint-path", type=Path, default=None) + parser.add_argument("--taehv-parallel", action="store_true") + parser.add_argument( + "--mlx-checkpoint-cache", + type=Path, + default=None, + help="Directory of per-mode pre-quantized MLX checkpoints. The first cell of a mode " + "converts from Diffusers weights and saves here (load_source=diffusers_then_saved); " + "later cells and later runs reload without requantizing (load_source=mlx_checkpoint), " + "which is also how the checkpoint load-time win is measured.", + ) + add_memory_limit_args( + parser, + mlx_memory_limit_gib=preset.mlx_memory_limit_gib, + mlx_disable_cache=preset.mlx_disable_cache, + torch_mps_high_watermark_ratio=preset.torch_mps_high_watermark_ratio, + torch_mps_low_watermark_ratio=preset.torch_mps_low_watermark_ratio, + ) + args = parser.parse_args() + + if args.compile: + os.environ["FASTVIDEO_MLX_COMPILE"] = "1" + + import mlx.core as mx + + runtime_limits = apply_memory_limits( + mlx_memory_limit_gib=args.mlx_memory_limit_gib, + mlx_cache_limit_gib=args.mlx_cache_limit_gib, + mlx_disable_cache=args.mlx_disable_cache, + mlx_wired_limit_gib=args.mlx_wired_limit_gib, + torch_mps_high_watermark_ratio=args.torch_mps_high_watermark_ratio, + torch_mps_low_watermark_ratio=args.torch_mps_low_watermark_ratio, + mx_module=mx, + ).as_metrics() + import torch + + mx.random.seed(args.seed) + torch.manual_seed(args.seed) + args.output_dir.mkdir(parents=True, exist_ok=True) + + modes = _parse_list(args.modes, ALLOWED_MODES, "modes") + decoders = _parse_list(args.decoders, ALLOWED_DECODERS, "decoders") + + config_path = args.model_root / "transformer/config.json" + checkpoint_path = args.model_root / "transformer/diffusion_pytorch_model.safetensors" + config = json.loads(config_path.read_text()) + latent_frames = (args.num_frames - 1) // 4 + 1 + latent_height = args.height // 8 + latent_width = args.width // 8 + freqs_cis = make_rotary_embeddings( + config, + latent_frames=latent_frames, + latent_height=latent_height, + latent_width=latent_width, + ) + generator = torch.Generator(device="cpu").manual_seed(args.seed) + latents_seed = torch.randn( + (1, int(config["in_channels"]), latent_frames, latent_height, latent_width), + generator=generator, + dtype=torch.float32, + ).numpy() + + timesteps = [int(step.strip()) for step in args.dmd_denoising_steps.split(",") if step.strip()] + # Keep DMD stochasticity identical across benchmark cells. Without this, + # FP16/INT8/decoder comparisons can accidentally measure different re-noise + # samples instead of only quantization or decoder differences. + renoise_by_step = [ + torch.randn(latents_seed.shape, generator=generator, dtype=torch.float32).numpy() + for _ in range(max(0, + len(timesteps) - 1)) + ] + + from fastvideo.mlx_runtime.fastwan import UnsupportedMLXQuantizationError + + prompt_cases = _load_prompt_cases(args.prompt, args.prompt_file, args.prompt_set) + cells: list[Cell] = [] + unsupported_rows: list[dict] = [] + for prompt_case in prompt_cases: + args.current_prompt_id = prompt_case.id + args.current_prompt = prompt_case.prompt + prompt_embeds = encode_prompt( + model_root=args.model_root, + prompt=prompt_case.prompt, + max_sequence_length=args.max_sequence_length, + device_arg=args.torch_device, + dtype_arg=args.torch_dtype, + ) + encoder_hidden_states = mx.array(prompt_embeds.numpy()) + + for mode in modes: + for decoder in decoders: + print(f"=== cell: prompt={prompt_case.id} mode={mode} decoder={decoder} ===") + try: + cells.append( + _generate_cell( + args=args, + mode=mode, + decoder=decoder, + checkpoint_path=checkpoint_path, + config_path=config_path, + encoder_hidden_states=encoder_hidden_states, + freqs_cis=freqs_cis, + timesteps=timesteps, + latents_seed=latents_seed, + renoise_by_step=renoise_by_step, + )) + except UnsupportedMLXQuantizationError as exc: + # Record the cell as unsupported and keep sweeping: a partial + # report on this MLX build beats crashing the whole run. + print(f"skipping cell (unsupported by this MLX build): {exc}") + unsupported_rows.append({ + "prompt_id": prompt_case.id, + "prompt": prompt_case.prompt, + "mode": mode, + "decoder": decoder, + "status": "unsupported_by_mlx", + "error": str(exc), + }) + + if not cells: + metrics_path = args.output_dir / "metrics.json" + metrics_path.write_text(json.dumps(unsupported_rows, indent=2)) + raise SystemExit(f"No benchmark cell could run: every requested mode is unsupported by this MLX build. " + f"Wrote {metrics_path}.") + + # Resolve one internal reference per prompt. A single external reference, if + # supplied, is used for every prompt and only video metrics are computed. + reference_by_prompt: dict[str, tuple[Path, np.ndarray | None]] = {} + if args.reference is not None: + for prompt_case in prompt_cases: + reference_by_prompt[prompt_case.id] = (args.reference, None) + else: + for prompt_case in prompt_cases: + prompt_cells = [c for c in cells if c.prompt_id == prompt_case.id] + if not prompt_cells: + continue + ref_cell = next( + (c for c in prompt_cells if c.mode == REFERENCE_MODE and c.decoder == REFERENCE_DECODER), + prompt_cells[0], + ) + reference_by_prompt[prompt_case.id] = (ref_cell.video_path, ref_cell.latents) + print(f"Using internal reference cell for {prompt_case.id}: " + f"mode={ref_cell.mode} decoder={ref_cell.decoder}") + + lpips_fn = _load_lpips() if args.lpips else None + + rows: list[dict] = [] + failures: list[str] = [] + for cell in cells: + reference_video, reference_latents = reference_by_prompt[cell.prompt_id] + ms_ssim = _ms_ssim(Path(reference_video), cell.video_path, required=args.assert_min_ssim is not None) + cell.metrics["ms_ssim_vs_ref"] = ms_ssim + cell.metrics.update(runtime_limits) + if reference_latents is not None: + cell.metrics.update(_latent_delta_metrics(cell.latents, reference_latents)) + cell.metrics["lpips_vs_ref"] = (_lpips_between(lpips_fn, Path(reference_video), cell.video_path) + if lpips_fn else None) + if args.assert_min_ssim is not None and ms_ssim is not None and ms_ssim < args.assert_min_ssim: + failures.append( + f"{cell.prompt_id}/{cell.mode}/{cell.decoder}: MS-SSIM {ms_ssim:.4f} < {args.assert_min_ssim}") + rows.append(dict(cell.metrics)) + print(json.dumps(cell.metrics, indent=2)) + rows.extend(unsupported_rows) + + metrics_path = args.output_dir / "metrics.json" + metrics_path.write_text(json.dumps(rows, indent=2)) + table_path = args.output_dir / "metrics.md" + table = _markdown_table(rows) + table_path.write_text(table + "\n") + html_path = _write_html_grid(rows, args.output_dir) + print("\n" + table) + print(f"\nWrote {metrics_path}, {table_path}, and {html_path}") + + if failures: + raise SystemExit("SSIM regression gate failed:\n " + "\n ".join(failures)) + + +def _load_lpips() -> object | None: + """Return an LPIPS model, or ``None`` if the optional dep is unavailable.""" + try: + import lpips # noqa: PLC0415 - optional dependency. + except ImportError: + print("LPIPS requested but the `lpips` package is not installed; skipping (install `.[eval]`).") + return None + return lpips.LPIPS(net="alex") + + +def _lpips_between(lpips_fn, reference_video: Path, candidate_video: Path) -> float | None: + if lpips_fn is None or not reference_video.exists() or not candidate_video.exists(): + return None + import torch + + ref = _read_video_frames(reference_video) + cand = _read_video_frames(candidate_video) + if ref is None or cand is None or ref.shape != cand.shape: + return None + # LPIPS expects NCHW in [-1, 1]. + ref_t = torch.from_numpy(ref).permute(0, 3, 1, 2).float() / 127.5 - 1.0 + cand_t = torch.from_numpy(cand).permute(0, 3, 1, 2).float() / 127.5 - 1.0 + with torch.no_grad(): + scores = lpips_fn(ref_t, cand_t) + return float(scores.mean().item()) + + +def _read_video_frames(path: Path) -> np.ndarray | None: + try: + import cv2 + except ImportError: + return None + cap = cv2.VideoCapture(str(path)) + frames = [] + try: + while True: + ok, frame_bgr = cap.read() + if not ok: + break + frames.append(cv2.cvtColor(frame_bgr, cv2.COLOR_BGR2RGB)) + finally: + cap.release() + if not frames: + return None + return np.stack(frames, axis=0) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/mlx_runtime/__init__.py b/fastvideo/mlx_runtime/__init__.py new file mode 100644 index 000000000..89fa279c2 --- /dev/null +++ b/fastvideo/mlx_runtime/__init__.py @@ -0,0 +1,115 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Experimental Apple MLX runtime helpers. + +This package is intentionally small for now. It exists to grow the Apple-native +FastWan path in measurable steps: shape planning, primitive benchmarks, then +Wan block parity, then full DiT/runtime support. +""" + +from fastvideo.mlx_runtime.fastwan import ( + FastWanShape, + MLXQuantizationSpec, + MLXWanDiT, + MLXWanTransformerBlock, + UnsupportedMLXQuantizationError, + ensure_quantization_supported, + fastwan_shape, + fastwan_shape_from_config, + mlx_dit_from_diffusers_safetensors, + mlx_block_weights_from_torch, + mlx_block_weights_from_diffusers_safetensors, + quantization_support_error, +) +from fastvideo.mlx_runtime.checkpoint import ( + load_mlx_dit_checkpoint, + save_mlx_dit_checkpoint, +) +from fastvideo.mlx_runtime.memory import ( + AppliedMemoryLimits, + add_memory_limit_args, + apply_memory_limits, + gib_to_bytes, +) +from fastvideo.mlx_runtime.refine import ( + DEFAULT_REFINE_SIGMA, + RefinePlan, + TwoPassResult, + default_refine_timesteps, + plan_refine_resolutions, + prepare_refine_latents, + refine_sigma_from_schedule, + run_dmd_loop, + run_two_pass_dmd, + upsample_latents_spatial, +) +from fastvideo.mlx_runtime.frame_upsample import ( + DEFAULT_PIXEL_UPSAMPLE_MODE, + PIXEL_UPSAMPLE_MODES, + unsharp, + upsample_frame, + upsample_frames, +) +from fastvideo.mlx_runtime.fast_spatial import ( + DEFAULT_FAST_SPATIAL_SHARPEN, + FastSpatialPlan, + apply_fast_spatial_upsample, + plan_fast_spatial, + resolve_spatial_mode, +) +from fastvideo.mlx_runtime.prompt_enhance import ( + DEFAULT_ENHANCE_SYSTEM_PROMPT, + DEFAULT_MLX_LM_MODEL, + EnhanceResult, + enhance_prompt, + enhance_prompt_template, + enhance_result_as_metrics, + load_or_enhance_prompt, +) + +__all__ = [ + "AppliedMemoryLimits", + "DEFAULT_ENHANCE_SYSTEM_PROMPT", + "DEFAULT_MLX_LM_MODEL", + "DEFAULT_FAST_SPATIAL_SHARPEN", + "DEFAULT_PIXEL_UPSAMPLE_MODE", + "DEFAULT_REFINE_SIGMA", + "EnhanceResult", + "FastSpatialPlan", + "FastWanShape", + "MLXQuantizationSpec", + "MLXWanDiT", + "MLXWanTransformerBlock", + "RefinePlan", + "TwoPassResult", + "UnsupportedMLXQuantizationError", + "add_memory_limit_args", + "apply_fast_spatial_upsample", + "apply_memory_limits", + "enhance_prompt", + "enhance_prompt_template", + "enhance_result_as_metrics", + "ensure_quantization_supported", + "fastwan_shape", + "fastwan_shape_from_config", + "gib_to_bytes", + "load_mlx_dit_checkpoint", + "load_or_enhance_prompt", + "mlx_dit_from_diffusers_safetensors", + "mlx_block_weights_from_diffusers_safetensors", + "mlx_block_weights_from_torch", + "PIXEL_UPSAMPLE_MODES", + "default_refine_timesteps", + "plan_fast_spatial", + "plan_refine_resolutions", + "prepare_refine_latents", + "quantization_support_error", + "refine_sigma_from_schedule", + "resolve_spatial_mode", + "run_dmd_loop", + "run_two_pass_dmd", + "save_mlx_dit_checkpoint", + "unsharp", + "upsample_frame", + "upsample_frames", + "upsample_latents_spatial", +] diff --git a/fastvideo/mlx_runtime/checkpoint.py b/fastvideo/mlx_runtime/checkpoint.py new file mode 100644 index 000000000..a68fafa7d --- /dev/null +++ b/fastvideo/mlx_runtime/checkpoint.py @@ -0,0 +1,273 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pre-quantized MLX checkpoint save/load for the FastWan DiT. + +Loading the Diffusers fp32/fp16 checkpoint and quantizing at startup costs +both download size and load time on every run. This module persists an already +cast (and optionally already quantized) ``MLXWanDiT`` so 16 GB users download +and load roughly half the bytes and skip requantization entirely: + + dit = mlx_dit_from_diffusers_safetensors(ckpt, cfg, quantization="int8") + save_mlx_dit_checkpoint(dit, "FastWan2.1-T2V-1.3B-mlx-int8") + ... + dit = load_mlx_dit_checkpoint("FastWan2.1-T2V-1.3B-mlx-int8") + +Format (one directory): + +- ``mlx_dit.safetensors`` — every array, saved with ``mx.save_safetensors``. + Plain weights keep their key; a quantized weight ``K`` is stored as the + packed ``K`` plus ``K.scales`` (and ``K.biases`` for affine modes). +- ``mlx_dit.json`` — format version, the model config, the quantization spec, + and which keys are quantized, so the loader can rebuild ``QuantizedMatrix`` + objects without guessing. +""" + +from __future__ import annotations + +import json +import shutil +import tempfile +from pathlib import Path +from typing import Any + +from fastvideo.logger import init_logger +from fastvideo.mlx_runtime.fastwan import ( + MLXQuantizationSpec, + MLXWanDiT, + MLXWanTransformerBlock, + QuantizedMatrix, + ensure_quantization_supported, +) + +logger = init_logger(__name__) + +FORMAT_VERSION = 1 +WEIGHTS_FILENAME = "mlx_dit.safetensors" +MANIFEST_FILENAME = "mlx_dit.json" + +_BLOCK_PREFIX = "blocks" + +_DTYPE_TO_NAME = {"float16": "fp16", "bfloat16": "bf16", "float32": "fp32"} + + +def _dtype_name(dtype) -> str: + """Return the manifest name for a supported MLX data type. + + Parameters: + dtype: The MLX data type to convert. + + Returns: + str: The manifest name corresponding to the data type. + + Raises: + ValueError: If the data type is not supported for checkpointing. + """ + import mlx.core as mx + + for raw, name in _DTYPE_TO_NAME.items(): + if dtype == getattr(mx, raw): + return name + raise ValueError(f"Unsupported MLX dtype for checkpointing: {dtype}") + + +def _name_to_dtype(name: str): + """Convert a manifest dtype name to its corresponding MLX dtype. + + Parameters: + name (str): Manifest name, such as ``"fp16"``, ``"bf16"``, or ``"fp32"``. + + Returns: + The corresponding MLX dtype. + """ + import mlx.core as mx + + return {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[name] + + +def _flatten_weights(dit: MLXWanDiT) -> dict[str, Any]: + """Combine model and transformer-block weights into a single flattened mapping. + + Parameters: + dit (MLXWanDiT): Model whose weights should be flattened. + + Returns: + dict[str, Any]: Mapping containing top-level weights and indexed transformer-block weights. + """ + flat: dict[str, Any] = dict(dit.weights) + for index, block in enumerate(dit.blocks): + for name, value in block.weights.items(): + flat[f"{_BLOCK_PREFIX}.{index}.{name}"] = value + return flat + + +def save_mlx_dit_checkpoint(dit: MLXWanDiT, checkpoint_dir: str | Path) -> Path: + """Save a plain or quantized MLX Wan DiT checkpoint to a directory. + + Parameters: + dit (MLXWanDiT): Model whose weights and configuration will be saved. + checkpoint_dir (str | Path): Destination directory for the checkpoint. + + Returns: + Path: Path to the checkpoint directory. + """ + import mlx.core as mx + + checkpoint_dir = Path(checkpoint_dir) + arrays: dict[str, Any] = {} + quantized: dict[str, dict[str, Any]] = {} + spec: MLXQuantizationSpec | None = None + for key, value in _flatten_weights(dit).items(): + if isinstance(value, QuantizedMatrix): + if spec is not None and value.spec != spec: + raise ValueError(f"Mixed quantization specs in one checkpoint ({spec} vs {value.spec} at '{key}') " + "are not supported.") + spec = value.spec + arrays[key] = value.weight + arrays[f"{key}.scales"] = value.scales + if value.biases is not None: + arrays[f"{key}.biases"] = value.biases + quantized[key] = { + "dequantized_dtype": _dtype_name(value.dequantized_dtype), + "has_biases": value.biases is not None, + } + else: + arrays[key] = value + + manifest = { + "format_version": FORMAT_VERSION, + "config": dit.config, + "num_blocks": len(dit.blocks), + "quantization": None if spec is None else { + "mode": spec.mode, + "bits": spec.bits, + "group_size": spec.group_size, + }, + "quantized_keys": quantized, + } + + manifest_json = json.dumps(manifest, indent=2) + checkpoint_dir.parent.mkdir(parents=True, exist_ok=True) + staging_dir = Path(tempfile.mkdtemp(dir=checkpoint_dir.parent, prefix=f".{checkpoint_dir.name}.staging-")) + backup_root: Path | None = None + try: + staged_weights = staging_dir / WEIGHTS_FILENAME + staged_manifest = staging_dir / MANIFEST_FILENAME + mx.save_safetensors(str(staged_weights), arrays) + staged_manifest.write_text(manifest_json) + if checkpoint_dir.exists(): + backup_root = Path(tempfile.mkdtemp(dir=checkpoint_dir.parent, prefix=f".{checkpoint_dir.name}.backup-")) + try: + checkpoint_dir.replace(backup_root / checkpoint_dir.name) + except Exception: + shutil.rmtree(backup_root, ignore_errors=True) + raise + try: + staging_dir.replace(checkpoint_dir) + except Exception: + if backup_root is not None: + (backup_root / checkpoint_dir.name).replace(checkpoint_dir) + shutil.rmtree(backup_root, ignore_errors=True) + raise + if backup_root is not None: + shutil.rmtree(backup_root, ignore_errors=True) + finally: + shutil.rmtree(staging_dir, ignore_errors=True) + logger.info("Saved MLX DiT checkpoint (%d arrays, quantization=%s) to %s", len(arrays), + spec.label if spec else "none", checkpoint_dir) + return checkpoint_dir + + +def load_mlx_dit_checkpoint(checkpoint_dir: str | Path, *, compile: bool = False) -> MLXWanDiT: + """ + Reconstruct an MLXWanDiT model from a versioned checkpoint. + + Parameters: + checkpoint_dir (str | Path): Directory containing the checkpoint manifest and weights. + compile (bool): Whether to configure the reconstructed model for compilation. + + Returns: + MLXWanDiT: The reconstructed model. + + Raises: + FileNotFoundError: If the checkpoint manifest or weights file is missing. + ValueError: If the checkpoint format is unsupported or block weights are incomplete. + """ + import mlx.core as mx + + checkpoint_dir = Path(checkpoint_dir) + manifest_path = checkpoint_dir / MANIFEST_FILENAME + weights_path = checkpoint_dir / WEIGHTS_FILENAME + if not manifest_path.exists() or not weights_path.exists(): + raise FileNotFoundError(f"Not an MLX DiT checkpoint directory: {checkpoint_dir} " + f"(expected {MANIFEST_FILENAME} and {WEIGHTS_FILENAME}).") + + manifest = json.loads(manifest_path.read_text()) + version = manifest.get("format_version") + if version != FORMAT_VERSION: + raise ValueError(f"MLX DiT checkpoint {checkpoint_dir} has format_version={version}; " + f"this FastVideo build reads version {FORMAT_VERSION}. Re-export the checkpoint.") + + spec = None + if manifest["quantization"] is not None: + spec = MLXQuantizationSpec(**manifest["quantization"]) + # The packed layout of mx.quantize output is mode-specific, so a build + # that cannot run the mode cannot use these arrays at all. + ensure_quantization_supported(spec) + + arrays = mx.load(str(weights_path)) + quantized_keys: dict[str, dict[str, Any]] = manifest["quantized_keys"] + + def rebuild(key: str): + """ + Reconstructs a weight array or quantized matrix from checkpoint data. + + Parameters: + key (str): The weight key to rebuild. + + Returns: + The stored array for an unquantized weight or a reconstructed quantized matrix. + """ + if key not in quantized_keys: + return arrays[key] + info = quantized_keys[key] + assert spec is not None, f"Quantized key '{key}' in a checkpoint without a quantization spec" + return QuantizedMatrix( + weight=arrays[key], + scales=arrays[f"{key}.scales"], + biases=arrays[f"{key}.biases"] if info["has_biases"] else None, + spec=spec, + dequantized_dtype=_name_to_dtype(info["dequantized_dtype"]), + ) + + config = manifest["config"] + block_keys: dict[int, list[str]] = {} + top_level_keys: list[str] = [] + for key in arrays: + if key.endswith(".scales") or key.endswith(".biases"): + continue + if key.startswith(f"{_BLOCK_PREFIX}."): + index_str, _, _ = key[len(_BLOCK_PREFIX) + 1:].partition(".") + block_keys.setdefault(int(index_str), []).append(key) + else: + top_level_keys.append(key) + + weights = {key: rebuild(key) for key in top_level_keys} + + num_blocks = int(manifest["num_blocks"]) + if sorted(block_keys) != list(range(num_blocks)): + raise ValueError(f"MLX DiT checkpoint {checkpoint_dir} is missing block weights: " + f"manifest says {num_blocks} blocks, found indices {sorted(block_keys)}.") + + inner_dim = int(config["num_attention_heads"]) * int(config["attention_head_dim"]) + blocks = [] + for index in range(num_blocks): + prefix = f"{_BLOCK_PREFIX}.{index}." + block_weights = {key[len(prefix):]: rebuild(key) for key in block_keys[index]} + blocks.append( + MLXWanTransformerBlock( + block_weights, + dim=inner_dim, + ffn_dim=int(config["ffn_dim"]), + num_heads=int(config["num_attention_heads"]), + eps=float(config["eps"]), + )) + return MLXWanDiT(weights, blocks, config, compile=compile) diff --git a/fastvideo/mlx_runtime/fast_spatial.py b/fastvideo/mlx_runtime/fast_spatial.py new file mode 100644 index 000000000..c424808d5 --- /dev/null +++ b/fastvideo/mlx_runtime/fast_spatial.py @@ -0,0 +1,228 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Spatial fast mode for the MLX Wan runtime (RIFE's spatial twin). + +RIFE ``--fast`` cuts *frames* (temporal). This module cuts *pixels* +(spatial): denoise at ``target // scale``, decode at that size, then +resample the decoded frames up to the target. No second denoise pass — +that is ``--refine`` (quality). The two compose: + +* ``--fast-spatial`` alone → speed (≈ scale² fewer tokens) +* ``--refine`` alone → quality two-pass (H3 / LTX-2) +* ``--fast`` + ``--refine`` → fewer frames at base res, full-res refine +* ``--fast`` + ``--fast-spatial`` → fewer frames *and* fewer pixels + +The upsample runs in **pixel** space, after the VAE decode. It used to run +in latent space (bilinear over the latent H/W plane, sharing the refine +hand-off primitive) and that is what made spatial fast mode incoherent: an +interpolated Wan latent is off the decoder's manifold, so decode returned +the right silhouette under a smeared veil. ``--refine`` can get away with +the latent-space upsample because a second DMD pass re-denoises the result; +spatial fast mode hands the latent straight to the decoder, so it cannot. +See :mod:`fastvideo.mlx_runtime.frame_upsample` for the full rationale. + +MetalFX is intentionally not used: it needs game-engine motion vectors +and depth that diffusion output lacks. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from dataclasses import dataclass + +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.refine import RefinePlan, plan_refine_resolutions + +logger = init_logger(__name__) + +# Resampling from a smaller decode loses high-frequency detail the same way +# RIFE's flow warp does, so spatial fast mode borrows ``--fast``'s remedy: a +# light unsharp pass. 0.4 recovers perceived crispness on Wan2.1 output at 2x +# without the halos that show up by ~0.8. +DEFAULT_FAST_SPATIAL_SHARPEN = 0.4 + + +@dataclass(frozen=True) +class FastSpatialPlan: + """Resolved geometry for a spatial-fast (upsample-only) run.""" + + plan: RefinePlan + upsample_mode: str + sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN + + @property + def enabled(self) -> bool: + """ + Determine whether spatial scaling is enabled. + + Returns: + `true` if the spatial scale is greater than one, `false` otherwise. + """ + return self.plan.spatial_scale > 1 + + @property + def scale(self) -> int: + """Provides the configured spatial scaling factor. + + Returns: + int: The spatial scaling factor. + """ + return self.plan.spatial_scale + + @property + def target_height(self) -> int: + """ + Return the target output height for the spatial plan. + + Returns: + int: Target output height in pixels. + """ + return self.plan.target_height + + @property + def target_width(self) -> int: + """Return the target image width in pixels. + + Returns: + int: The target image width. + """ + return self.plan.target_width + + @property + def stage1_height(self) -> int: + """ + Provide the stage-one latent height used for reduced-resolution processing. + + Returns: + int: The stage-one latent height. + """ + return self.plan.stage1_height + + @property + def stage1_width(self) -> int: + """Get the stage-one latent width. + + Returns: + int: The stage-one latent width. + """ + return self.plan.stage1_width + + +def plan_fast_spatial( + *, + height: int, + width: int, + num_frames: int, + spatial_scale: int = 2, + vae_spatial_compression: int = 8, + vae_temporal_compression: int = 4, + patch_size: tuple[int, int, int] = (1, 2, 2), + upsample_mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE, + sharpen: float = DEFAULT_FAST_SPATIAL_SHARPEN, + enabled: bool = True, +) -> FastSpatialPlan: + """ + Build a plan for reduced-resolution denoising followed by pixel-space upsampling. + + Parameters: + upsample_mode (str): Pixel interpolation kernel, one of + :data:`~fastvideo.mlx_runtime.frame_upsample.PIXEL_UPSAMPLE_MODES`. + sharpen (float): Unsharp strength applied after the resize. + + Returns: + FastSpatialPlan: The validated spatial-fast processing plan. + + Raises: + ValueError: If the upsample mode is unsupported or ``sharpen`` is negative. + """ + 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.0: + raise ValueError(f"sharpen must be >= 0, got {sharpen}") + plan = plan_refine_resolutions( + height=height, + width=width, + num_frames=num_frames, + spatial_scale=spatial_scale, + vae_spatial_compression=vae_spatial_compression, + vae_temporal_compression=vae_temporal_compression, + patch_size=patch_size, + enabled=enabled, + mode_label="fast-spatial", + ) + if plan.spatial_scale > 1: + logger.info( + "[MLX fast-spatial] denoise+decode %dx%d → upsample %dx to %dx%d (%s, sharpen=%.2f)", + plan.stage1_width, + plan.stage1_height, + plan.spatial_scale, + plan.target_width, + plan.target_height, + upsample_mode, + sharpen, + ) + return FastSpatialPlan(plan=plan, upsample_mode=upsample_mode, sharpen=sharpen) + + +def apply_fast_spatial_upsample( + frames: Iterable[np.ndarray], + spatial: FastSpatialPlan, +) -> list[np.ndarray]: + """Resample decoded stage-1 frames up to the target resolution. + + This runs on decoded RGB frames, *not* on latents: see the module + docstring for why the latent-space version produced a blurred veil. + + Parameters: + frames (Iterable[np.ndarray]): Decoded HxWx3 uint8 RGB frames, produced + by decoding at the stage-one resolution. + spatial (FastSpatialPlan): Plan defining the target size, interpolation + kernel, and unsharp strength. + + Returns: + list[np.ndarray]: Frames at the target resolution. When spatial scaling + is disabled the frames are returned unchanged, as a list. + """ + if not spatial.enabled: + return list(frames) + return upsample_frames( + frames, + width=spatial.target_width, + height=spatial.target_height, + mode=spatial.upsample_mode, + sharpen=spatial.sharpen, + ) + + +def resolve_spatial_mode( + *, + refine: bool, + fast_spatial: bool, +) -> str: + """Select the active spatial processing mode, with refinement taking precedence. + + Returns: + str: ``"refine"`` when refinement is enabled, ``"fast_spatial"`` when + spatial-fast processing is enabled, or ``"off"`` otherwise. + """ + if refine: + return "refine" + if fast_spatial: + return "fast_spatial" + return "off" + + +__all__ = [ + "DEFAULT_FAST_SPATIAL_SHARPEN", + "FastSpatialPlan", + "apply_fast_spatial_upsample", + "plan_fast_spatial", + "resolve_spatial_mode", +] diff --git a/fastvideo/mlx_runtime/fastwan.py b/fastvideo/mlx_runtime/fastwan.py new file mode 100644 index 000000000..89f8d2fa9 --- /dev/null +++ b/fastvideo/mlx_runtime/fastwan.py @@ -0,0 +1,980 @@ +# SPDX-License-Identifier: Apache-2.0 +# mypy: disable-error-code=no-untyped-call +"""FastWan-oriented helpers for the experimental MLX runtime path.""" + +from __future__ import annotations + +import json +import math +import os +import statistics +import time +from dataclasses import dataclass +from pathlib import Path +from typing import TYPE_CHECKING, Any +from collections.abc import Callable + +from fastvideo.logger import init_logger + +if TYPE_CHECKING: + import mlx.core as mx + import torch + +logger = init_logger(__name__) + + +@dataclass(frozen=True) +class FastWanShape: + height: int + width: int + num_frames: int + latent_frames: int + latent_height: int + latent_width: int + patch_frames: int + patch_height: int + patch_width: int + tokens: int + hidden_size: int + num_heads: int + head_dim: int + + +class UnsupportedMLXQuantizationError(ValueError): + """A quantization mode the installed MLX build cannot execute. + + Raised by :func:`ensure_quantization_supported` before any model weights + are loaded, so callers (CLI flags, benchmark sweeps) can fail fast with an + actionable message -- or skip the mode -- instead of crashing deep inside + ``mx.quantize`` mid-load. + """ + + +@dataclass(frozen=True) +class MLXQuantizationSpec: + """MLX quantized-matmul configuration for DiT linear weights.""" + + mode: str + bits: int | None = None + group_size: int | None = None + + @classmethod + def from_name(cls, name: str | None) -> MLXQuantizationSpec | None: + if name is None or name in {"", "none", "fp16", "fp32"}: + return None + if name == "int8": + return cls(mode="affine", bits=8, group_size=64) + if name == "int4": + return cls(mode="affine", bits=4, group_size=64) + if name == "mxfp8": + return cls(mode="mxfp8") + if name == "mxfp4": + return cls(mode="mxfp4") + if name == "nvfp4": + return cls(mode="nvfp4") + raise ValueError(f"Unsupported MLX quantization mode: {name}") + + @property + def label(self) -> str: + if self.mode == "affine": + return f"int{self.bits}" + return self.mode + + +@dataclass(frozen=True) +class QuantizedMatrix: + weight: mx.array + scales: mx.array + biases: mx.array | None + spec: MLXQuantizationSpec + dequantized_dtype: mx.Dtype + + +def fastwan_shape( + *, + height: int, + width: int, + num_frames: int, + vae_temporal_compression: int = 4, + vae_spatial_compression: int = 8, + patch_size: tuple[int, int, int] = (1, 2, 2), + num_heads: int = 12, + head_dim: int = 128, +) -> FastWanShape: + """Return the approximate DiT token shape for Wan/FastWan T2V inference.""" + latent_frames = (num_frames - 1) // vae_temporal_compression + 1 + latent_height = height // vae_spatial_compression + latent_width = width // vae_spatial_compression + patch_frames = latent_frames // patch_size[0] + patch_height = latent_height // patch_size[1] + patch_width = latent_width // patch_size[2] + tokens = patch_frames * patch_height * patch_width + return FastWanShape( + height=height, + width=width, + num_frames=num_frames, + latent_frames=latent_frames, + latent_height=latent_height, + latent_width=latent_width, + patch_frames=patch_frames, + patch_height=patch_height, + patch_width=patch_width, + tokens=tokens, + hidden_size=num_heads * head_dim, + num_heads=num_heads, + head_dim=head_dim, + ) + + +def fastwan_shape_from_config( + config_path: str | Path, + *, + height: int, + width: int, + num_frames: int, +) -> FastWanShape: + config = json.loads(Path(config_path).read_text()) + return fastwan_shape( + height=height, + width=width, + num_frames=num_frames, + patch_size=tuple(config["patch_size"]), + num_heads=int(config["num_attention_heads"]), + head_dim=int(config["attention_head_dim"]), + ) + + +def replace_tokens(shape: FastWanShape, tokens: int) -> FastWanShape: + return FastWanShape(**{**shape.__dict__, "tokens": tokens}) + + +def median_ms(samples: list[float]) -> float: + return statistics.median(samples) * 1000.0 + + +def benchmark_mlx_attention(shape: FastWanShape, warmup: int, iters: int) -> float: + import mlx.core as mx + + q = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16) + k = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16) + v = mx.random.normal((1, shape.num_heads, shape.tokens, shape.head_dim), dtype=mx.float16) + scale = shape.head_dim**-0.5 + + for _ in range(warmup): + y = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale) + mx.eval(y) + + samples = [] + for _ in range(iters): + start = time.perf_counter() + y = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale) + mx.eval(y) + samples.append(time.perf_counter() - start) + return median_ms(samples) + + +def benchmark_mlx_linear(shape: FastWanShape, warmup: int, iters: int) -> float: + import mlx.core as mx + + x = mx.random.normal((shape.tokens, shape.hidden_size), dtype=mx.float16) + w = mx.random.normal((shape.hidden_size, shape.hidden_size), dtype=mx.float16) + b = mx.zeros((shape.hidden_size, ), dtype=mx.float16) + + for _ in range(warmup): + y = x @ w + b + mx.eval(y) + + samples = [] + for _ in range(iters): + start = time.perf_counter() + y = x @ w + b + mx.eval(y) + samples.append(time.perf_counter() - start) + return median_ms(samples) + + +def benchmark_torch_mps_attention(shape: FastWanShape, warmup: int, iters: int) -> float | None: + try: + import torch + import torch.nn.functional as F + except ImportError: + return None + + if not torch.backends.mps.is_available(): + return None + + device = torch.device("mps") + q = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16) + k = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16) + v = torch.randn((1, shape.num_heads, shape.tokens, shape.head_dim), device=device, dtype=torch.float16) + + for _ in range(warmup): + y = F.scaled_dot_product_attention(q, k, v) + torch.mps.synchronize() + _ = y + + samples = [] + for _ in range(iters): + start = time.perf_counter() + y = F.scaled_dot_product_attention(q, k, v) + torch.mps.synchronize() + _ = y + samples.append(time.perf_counter() - start) + return median_ms(samples) + + +def torch_to_mx(tensor) -> mx.array: + import mlx.core as mx + + return mx.array(tensor.detach().cpu().float().numpy()) + + +def weight_dtype(weight): + if isinstance(weight, QuantizedMatrix): + return weight.dequantized_dtype + return weight.dtype + + +_QUANT_SUPPORT_CACHE: dict[tuple[str, int | None, int | None], str | None] = {} + + +def quantization_support_error(spec: MLXQuantizationSpec) -> str | None: + """Probe whether the installed MLX build supports ``spec``. + + Runs a tiny ``mx.quantize`` + ``mx.quantized_matmul`` with exactly the + arguments :func:`quantize_matrix` / :func:`linear` use, so the result + reflects the real runtime path. The affine (int8/int4) modes are stable + across MLX releases, but the ``mxfp8``/``mxfp4``/``nvfp4`` mode strings + require newer MLX builds and raise otherwise. Returns ``None`` when the + mode works, else the underlying error message. Cached per spec. + """ + key = (spec.mode, spec.bits, spec.group_size) + if key not in _QUANT_SUPPORT_CACHE: + import mlx.core as mx + + try: + probe_dim = max(spec.group_size or 0, 64) + weight = mx.zeros((probe_dim, probe_dim), dtype=mx.float16) + quantized = quantize_matrix(weight, spec) + y = linear(mx.zeros((1, probe_dim), dtype=mx.float16), quantized) + mx.eval(y) + _QUANT_SUPPORT_CACHE[key] = None + except Exception as exc: # noqa: BLE001 - MLX raises varied error types per backend/version. + _QUANT_SUPPORT_CACHE[key] = f"{type(exc).__name__}: {exc}" + return _QUANT_SUPPORT_CACHE[key] + + +def ensure_quantization_supported(spec: MLXQuantizationSpec | None) -> None: + """Raise :class:`UnsupportedMLXQuantizationError` if ``spec`` cannot run here.""" + if spec is None: + return + error = quantization_support_error(spec) + if error is None: + return + import mlx.core as mx + + mlx_version = getattr(mx, "__version__", "unknown") + raise UnsupportedMLXQuantizationError(f"MLX quantization mode '{spec.label}' is not supported by the installed mlx " + f"({mlx_version}): {error}. Upgrade mlx or pick a supported mode " + f"(int8 is currently the most reliable quality/memory target).") + + +def quantize_matrix(weight, spec: MLXQuantizationSpec | None): + if spec is None: + return weight + import mlx.core as mx + + if len(weight.shape) < 2: + return weight + q = mx.quantize(weight, group_size=spec.group_size, bits=spec.bits, mode=spec.mode) + biases = q[2] if len(q) == 3 else None + eval_args = [q[0], q[1]] + if biases is not None: + eval_args.append(biases) + mx.eval(*eval_args) + return QuantizedMatrix( + weight=q[0], + scales=q[1], + biases=biases, + spec=spec, + dequantized_dtype=weight.dtype, + ) + + +def linear(x, weight, bias=None): + import mlx.core as mx + + if isinstance(weight, QuantizedMatrix): + y = mx.quantized_matmul( + x, + weight.weight, + weight.scales, + weight.biases, + transpose=True, + group_size=weight.spec.group_size, + bits=weight.spec.bits, + mode=weight.spec.mode, + ).astype(x.dtype) + else: + y = x @ weight.T + if bias is not None: + y = y + bias + return y + + +def _use_fast_norm() -> bool: + """Opt-in to MLX's fused ``mx.fast`` normalization kernels. + + Off by default so the numerically-explicit reference path stays the + baseline. Set ``FASTVIDEO_MLX_FAST_NORM=1`` to route LayerNorm/RMSNorm + through single fused Metal kernels (fewer intermediates, less memory + traffic) and benchmark the speedup. + """ + import os + + return os.environ.get("FASTVIDEO_MLX_FAST_NORM", "0") == "1" + + +def layer_norm(x, weight=None, bias=None, eps: float = 1e-6): + import mlx.core as mx + + if _use_fast_norm(): + # Compute in fp32 (matching the reference below) so downstream dtype + # and precision are identical across call sites. + w = weight.astype(mx.float32) if weight is not None else None + b = bias.astype(mx.float32) if bias is not None else None + return mx.fast.layer_norm(x.astype(mx.float32), w, b, eps) + + x_float = x.astype(mx.float32) + mean = mx.mean(x_float, axis=-1, keepdims=True) + var = mx.mean(mx.square(x_float - mean), axis=-1, keepdims=True) + y = (x_float - mean) * mx.rsqrt(var + eps) + if weight is not None: + y = y * weight + if bias is not None: + y = y + bias + return y + + +def rms_norm(x, weight, eps: float = 1e-6): + import mlx.core as mx + + if _use_fast_norm(): + return mx.fast.rms_norm(x, weight, eps) + + orig_dtype = x.dtype + x_float = x.astype(mx.float32) + variance = mx.mean(mx.square(x_float), axis=-1, keepdims=True) + y = x_float * mx.rsqrt(variance + eps) + return y.astype(orig_dtype) * weight + + +def apply_rotary_emb(x, cos, sin, *, is_neox_style: bool = False): + """Apply FastVideo's rotary convention to MLX tensors. + + Args: + x: [batch, seq, heads, head_dim] + cos/sin: [seq, head_dim] for Wan's full-dimension rotate-pair style, + or [seq, head_dim // 2] for traditional RoPE. + """ + import mlx.core as mx + + head_size = x.shape[-1] + rope_dim = cos.shape[-1] + cos = cos[None, :, None, :] + sin = sin[None, :, None, :] + x_float = x.astype(mx.float32) + + if rope_dim == head_size: + x_pairs = x_float.reshape(*x.shape[:-1], -1, 2) + x_real = x_pairs[..., 0] + x_imag = x_pairs[..., 1] + x_rotated = mx.stack([-x_imag, x_real], axis=-1).reshape(*x.shape) + return (x_float * cos + x_rotated * sin).astype(x.dtype) + + if is_neox_style: + x1, x2 = mx.split(x_float, 2, axis=-1) + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + return mx.concatenate([o1, o2], axis=-1).astype(x.dtype) + + x1 = x_float[..., ::2] + x2 = x_float[..., 1::2] + o1 = x1 * cos - x2 * sin + o2 = x2 * cos + x1 * sin + return mx.stack([o1, o2], axis=-1).reshape(*x.shape).astype(x.dtype) + + +_WINDOWED_ATTENTION_WARNED = False + + +def _warn_windowed_attention_once(window: int) -> None: + """Warn that FASTVIDEO_MLX_WINDOW degrades output on a dense-trained DiT. + + Sliding-window self-attention is fast (6.6x at a +-3-frame window on 1.3B) + but these checkpoints were trained with dense attention, and restricting it + at inference produces heavy colour-block noise: structural agreement with + the dense baseline drops to 0.25 at +-3 frames and 0.03 at +-5. Sparsity of + this kind is a training-time method. Kept as a research knob, but it should + never be on by accident. + """ + global _WINDOWED_ATTENTION_WARNED + if _WINDOWED_ATTENTION_WARNED: + return + _WINDOWED_ATTENTION_WARNED = True + logger.warning( + "FASTVIDEO_MLX_WINDOW=%d enables sliding-window self-attention. These " + "checkpoints are trained dense; expect severely degraded output. This is " + "a research knob, not a speed setting — use --fast-spatial for real " + "denoise savings.", + window, + ) + + +def gelu_tanh(x): + """tanh-approximate GELU, as used by Wan's FFN. + + ``mlx.nn.gelu_approx`` is the same tanh approximation behind a fused + kernel. On the 1.3B FFN shape (32760x8960) it is bit-identical to the + expanded expression below and 3.3x faster — 28.9ms -> 8.7ms per layer, + which is 0.6s per denoise step across 30 layers. + """ + import mlx.nn as nn + + return nn.gelu_approx(x) + + +def silu(x): + import mlx.core as mx + + return x * mx.sigmoid(x) + + +def timestep_embedding(t, dim: int, max_period: int = 10000): + import mlx.core as mx + + half = dim // 2 + freqs = mx.exp(-math.log(max_period) * mx.arange(0, half, dtype=mx.float32) / half) + args = t[:, None].astype(mx.float32) * freqs[None] + embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1) + if dim % 2: + embedding = mx.concatenate([embedding, mx.zeros_like(embedding[:, :1])], axis=-1) + return embedding + + +def scale_residual(residual, x, gate): + return residual + x * gate + + +def scale_residual_layer_norm_scale_shift(residual, x, gate, shift, scale, weight=None, bias=None, eps: float = 1e-6): + if isinstance(gate, int): + assert gate == 1 + residual_output = residual + x + else: + residual_output = residual + x * gate + normalized = layer_norm(residual_output, weight=weight, bias=bias, eps=eps) + modulated = normalized * (1.0 + scale) + shift + return modulated, residual_output + + +class MLXWanT2VCrossAttention: + + def __init__(self, weights: dict[str, mx.array], *, dim: int, num_heads: int, eps: float = 1e-6) -> None: + self.weights = weights + self.dim = dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.eps = eps + + def __call__(self, x, context): + import mlx.core as mx + + batch = x.shape[0] + q = linear(x, self.weights["attn2.to_q.weight"], self.weights.get("attn2.to_q.bias")) + q = rms_norm(q, self.weights["attn2.norm_q.weight"], eps=self.eps).reshape(batch, -1, self.num_heads, + self.head_dim) + + if context.shape[1] == 0: + attended = mx.zeros_like(q) + else: + k = linear(context, self.weights["attn2.to_k.weight"], self.weights.get("attn2.to_k.bias")) + k = rms_norm(k, self.weights["attn2.norm_k.weight"], + eps=self.eps).reshape(batch, -1, self.num_heads, self.head_dim) + v = linear(context, self.weights["attn2.to_v.weight"], + self.weights.get("attn2.to_v.bias")).reshape(batch, -1, self.num_heads, self.head_dim) + attended = mx.fast.scaled_dot_product_attention( + q.transpose(0, 2, 1, 3), + k.transpose(0, 2, 1, 3), + v.transpose(0, 2, 1, 3), + scale=self.head_dim**-0.5, + ).transpose(0, 2, 1, 3) + + attended = attended.reshape(batch, -1, self.dim) + return linear(attended, self.weights["attn2.to_out.weight"], self.weights.get("attn2.to_out.bias")) + + +class MLXWanTransformerBlock: + """Dense T2V Wan transformer block for the experimental MLX runtime. + + This mirrors the non-VSA PyTorch block for single-process dense attention. + Rotary embeddings and sequence-parallel paths are intentionally left out of + this first parity target. + """ + + def __init__(self, weights: dict[str, mx.array], *, dim: int, ffn_dim: int, num_heads: int, eps: float = 1e-6): + self.weights = weights + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.eps = eps + self.attn2 = MLXWanT2VCrossAttention(weights, dim=dim, num_heads=num_heads, eps=eps) + + def __call__(self, hidden_states, encoder_hidden_states, temb, freqs_cis=None): + import mlx.core as mx + + orig_dtype = hidden_states.dtype + e = self.weights["scale_shift_table"] + temb.astype(mx.float32) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = mx.split(e, 6, axis=1) + + norm_hidden_states = layer_norm(hidden_states.astype(mx.float32), eps=self.eps) + norm_hidden_states = (norm_hidden_states * (1.0 + scale_msa) + shift_msa).astype(orig_dtype) + + query = linear(norm_hidden_states, self.weights["to_q.weight"], self.weights.get("to_q.bias")) + key = linear(norm_hidden_states, self.weights["to_k.weight"], self.weights.get("to_k.bias")) + value = linear(norm_hidden_states, self.weights["to_v.weight"], self.weights.get("to_v.bias")) + + query = rms_norm(query, self.weights["norm_q.weight"], + eps=self.eps).reshape(hidden_states.shape[0], -1, self.num_heads, self.head_dim) + key = rms_norm(key, self.weights["norm_k.weight"], eps=self.eps).reshape(hidden_states.shape[0], -1, + self.num_heads, self.head_dim) + value = value.reshape(hidden_states.shape[0], -1, self.num_heads, self.head_dim) + + if freqs_cis is not None: + cos, sin = freqs_cis + query = apply_rotary_emb(query, cos, sin, is_neox_style=False) + key = apply_rotary_emb(key, cos, sin, is_neox_style=False) + + # Self-attention only. FASTVIDEO_MLX_WINDOW=0/unset → full SDPA (byte-identical + # to the historical path). When >0, use chunked symmetric sliding-window + # attention (see windowed_attention.py). Cross-attn (attn2) stays full. + # Optional FASTVIDEO_MLX_WINDOW_SINK (default 0) adds global sink tokens. + q_bh = query.transpose(0, 2, 1, 3) # (B, H, S, D) + k_bh = key.transpose(0, 2, 1, 3) + v_bh = value.transpose(0, 2, 1, 3) + scale = self.head_dim**-0.5 + window = int(os.environ.get("FASTVIDEO_MLX_WINDOW", "0") or "0") + if window > 0: + from fastvideo.mlx_runtime.windowed_attention import windowed_attention + + _warn_windowed_attention_once(window) + sink = int(os.environ.get("FASTVIDEO_MLX_WINDOW_SINK", "0") or "0") + attn_output = windowed_attention(q_bh, k_bh, v_bh, window=window, sink=sink, scale=scale) + else: + attn_output = mx.fast.scaled_dot_product_attention(q_bh, k_bh, v_bh, scale=scale) + attn_output = attn_output.transpose(0, 2, 1, 3) + attn_output = attn_output.reshape(hidden_states.shape[0], -1, self.dim) + attn_output = linear(attn_output, self.weights["to_out.weight"], self.weights.get("to_out.bias")) + + norm_hidden_states, hidden_states = scale_residual_layer_norm_scale_shift( + hidden_states, + attn_output, + gate_msa, + 0.0, + 0.0, + weight=self.weights["self_attn_residual_norm.norm.weight"], + bias=self.weights["self_attn_residual_norm.norm.bias"], + eps=self.eps, + ) + norm_hidden_states = norm_hidden_states.astype(orig_dtype) + hidden_states = hidden_states.astype(orig_dtype) + + attn_output = self.attn2(norm_hidden_states, encoder_hidden_states) + norm_hidden_states, hidden_states = scale_residual_layer_norm_scale_shift( + hidden_states, + attn_output, + 1, + c_shift_msa, + c_scale_msa, + eps=self.eps, + ) + norm_hidden_states = norm_hidden_states.astype(orig_dtype) + hidden_states = hidden_states.astype(orig_dtype) + + ff_output = linear(norm_hidden_states, self.weights["ffn.fc_in.weight"], self.weights.get("ffn.fc_in.bias")) + ff_output = gelu_tanh(ff_output) + ff_output = linear(ff_output, self.weights["ffn.fc_out.weight"], self.weights.get("ffn.fc_out.bias")) + hidden_states = scale_residual(hidden_states, ff_output, c_gate_msa) + return hidden_states.astype(orig_dtype) + + +def mlx_block_weights_from_torch(torch_block) -> dict[str, mx.array]: + return {name: torch_to_mx(value) for name, value in torch_block.state_dict().items()} + + +class MLXWanDiT: + """Experimental FP16 Wan/FastWan DiT forward path in MLX.""" + + def __init__( + self, + weights: dict[str, mx.array], + blocks: list[MLXWanTransformerBlock], + config: dict, + *, + compile: bool = False, + ) -> None: + import os + + self.weights = weights + self.blocks = blocks + self.config = config + self.num_heads = int(config["num_attention_heads"]) + self.head_dim = int(config["attention_head_dim"]) + self.hidden_size = self.num_heads * self.head_dim + self.ffn_dim = int(config["ffn_dim"]) + self.in_channels = int(config["in_channels"]) + self.out_channels = int(config["out_channels"]) + self.patch_size = tuple(config["patch_size"]) + self.freq_dim = int(config["freq_dim"]) + # Opt-in graph fusion. With fixed weights and static shapes, the whole + # denoise-step forward is a pure function of (latents, timestep) -- a + # good mx.compile target. Off by default so the eager path stays the + # baseline; enable via constructor or FASTVIDEO_MLX_COMPILE=1 and verify + # with the benchmark's SSIM ~= 1.0 check. + self._enable_compile = compile or os.environ.get("FASTVIDEO_MLX_COMPILE", "0") == "1" + self._compiled_forward: Callable[..., Any] | None = None + self._compiled_signature: tuple | None = None + + def patch_embed(self, hidden_states): + batch, channels, frames, height, width = hidden_states.shape + pt, ph, pw = self.patch_size + patch_dim = channels * pt * ph * pw + x = hidden_states.reshape(batch, channels, frames // pt, pt, height // ph, ph, width // pw, pw) + x = x.transpose(0, 2, 4, 6, 1, 3, 5, 7).reshape(batch, -1, patch_dim) + return linear(x, self.weights["patch_embedding.weight"], self.weights.get("patch_embedding.bias")) + + def condition(self, timestep, encoder_hidden_states): + t_freq = timestep_embedding(timestep, self.freq_dim).astype( + weight_dtype(self.weights["condition_embedder.time_embedder.linear_1.weight"])) + temb = linear( + t_freq, + self.weights["condition_embedder.time_embedder.linear_1.weight"], + self.weights["condition_embedder.time_embedder.linear_1.bias"], + ) + temb = silu(temb) + temb = linear( + temb, + self.weights["condition_embedder.time_embedder.linear_2.weight"], + self.weights["condition_embedder.time_embedder.linear_2.bias"], + ) + timestep_proj = silu(temb) + timestep_proj = linear( + timestep_proj, + self.weights["condition_embedder.time_proj.weight"], + self.weights["condition_embedder.time_proj.bias"], + ).reshape(timestep.shape[0], 6, self.hidden_size) + + encoder_hidden_states = linear( + encoder_hidden_states, + self.weights["condition_embedder.text_embedder.linear_1.weight"], + self.weights["condition_embedder.text_embedder.linear_1.bias"], + ) + encoder_hidden_states = gelu_tanh(encoder_hidden_states) + encoder_hidden_states = linear( + encoder_hidden_states, + self.weights["condition_embedder.text_embedder.linear_2.weight"], + self.weights["condition_embedder.text_embedder.linear_2.bias"], + ) + return temb, timestep_proj, encoder_hidden_states + + def output(self, hidden_states, temb, *, batch: int, frames: int, height: int, width: int): + pt, ph, pw = self.patch_size + post_patch_frames = frames // pt + post_patch_height = height // ph + post_patch_width = width // pw + shift, scale = mx_split_two(self.weights["scale_shift_table"] + temb[:, None, :], axis=1) + hidden_states = layer_norm(hidden_states, eps=float(self.config["eps"])) * (1.0 + scale) + shift + hidden_states = hidden_states.astype(weight_dtype(self.weights["proj_out.weight"])) + hidden_states = linear(hidden_states, self.weights["proj_out.weight"], self.weights["proj_out.bias"]) + hidden_states = hidden_states.reshape( + batch, + post_patch_frames, + post_patch_height, + post_patch_width, + pt, + ph, + pw, + self.out_channels, + ) + hidden_states = hidden_states.transpose(0, 7, 1, 4, 2, 5, 3, 6) + return hidden_states.reshape(batch, self.out_channels, frames, height, width) + + def _forward(self, hidden_states, encoder_hidden_states, timestep, cos, sin): + """Pure forward used both eagerly and as the mx.compile target. + + ``cos``/``sin`` are passed as separate array args (rather than a tuple) + so the function traces cleanly under mx.compile. + """ + batch, _, frames, height, width = hidden_states.shape + freqs_cis = (cos, sin) if cos is not None else None + hidden_states = self.patch_embed(hidden_states) + temb, timestep_proj, encoder_hidden_states = self.condition(timestep, encoder_hidden_states) + for block in self.blocks: + hidden_states = block(hidden_states, encoder_hidden_states, timestep_proj, freqs_cis=freqs_cis) + return self.output(hidden_states, temb, batch=batch, frames=frames, height=height, width=width) + + def __call__(self, hidden_states, encoder_hidden_states, timestep, freqs_cis): + cos, sin = freqs_cis if freqs_cis is not None else (None, None) + if self._enable_compile and cos is not None: + import mlx.core as mx + + # mx.compile keeps one traced graph per input signature, and each + # graph pins its own materialization of the quantized weights. The + # two-pass modes (--refine) call the DiT at a second resolution, so + # keeping both graphs alive doubles resident DiT memory: 14B refine + # peaked at 34.7 GiB instead of 20.8 GiB. Retire the previous graph + # when the signature changes; the retrace costs far less than a + # second copy of the weights. + signature = (hidden_states.shape, encoder_hidden_states.shape, timestep.shape) + if self._compiled_forward is not None and signature != self._compiled_signature: + self._compiled_forward = None + self._compiled_signature = None + mx.clear_cache() + if self._compiled_forward is None: + self._compiled_forward = mx.compile(self._forward) + self._compiled_signature = signature + compiled_forward = self._compiled_forward + try: + return compiled_forward(hidden_states, encoder_hidden_states, timestep, cos, sin) + except Exception as exc: # noqa: BLE001 - some quant graphs may not trace; fall back to eager. + logger.warning("mx.compile forward failed (%s); falling back to eager execution.", exc) + self._enable_compile = False + self._compiled_forward = None + self._compiled_signature = None + return self._forward(hidden_states, encoder_hidden_states, timestep, cos, sin) + + +def mx_split_two(x, *, axis: int): + import mlx.core as mx + + left, right = mx.split(x, 2, axis=axis) + return left, right + + +def _load_safetensor_value(handle, name: str): + return handle.get_tensor(name) + + +def _load_mx_array_from_safetensor(handle, name: str, dtype): + """Load a safetensors value and cast before creating the MLX array. + + The FastWan Diffusers checkpoint is fp32. Creating an MLX array first and + then casting it to fp16 briefly materializes a large fp32 MLX allocation. + Casting the CPU tensor before crossing into MLX keeps the transient GPU-side + footprint lower. + """ + import mlx.core as mx + import torch + + tensor = handle.get_tensor(name) + if dtype == mx.float16: + tensor = tensor.to(torch.float16) + elif dtype == mx.float32: + tensor = tensor.to(torch.float32) + elif dtype == mx.bfloat16: + # NumPy has no bfloat16, so bridge through fp32 and cast on-device below. + tensor = tensor.to(torch.float32) + array = mx.array(tensor.numpy()) + del tensor + if dtype is not None and array.dtype != dtype: + array = array.astype(dtype) + mx.eval(array) + return array + + +def _eval_loaded_weight(value) -> None: + import mlx.core as mx + + if isinstance(value, QuantizedMatrix): + eval_args = [value.weight, value.scales] + if value.biases is not None: + eval_args.append(value.biases) + mx.eval(*eval_args) + else: + mx.eval(value) + + +# Diffusers-to-FastVideo key mapping for WanTransformerBlock weights. +# Shared by both MLX and torch block loaders to keep mappings synchronized. +_WAN_BLOCK_KEY_MAP = { + "scale_shift_table": "scale_shift_table", + "attn1.to_q.weight": "to_q.weight", + "attn1.to_q.bias": "to_q.bias", + "attn1.to_k.weight": "to_k.weight", + "attn1.to_k.bias": "to_k.bias", + "attn1.to_v.weight": "to_v.weight", + "attn1.to_v.bias": "to_v.bias", + "attn1.to_out.0.weight": "to_out.weight", + "attn1.to_out.0.bias": "to_out.bias", + "attn1.norm_q.weight": "norm_q.weight", + "attn1.norm_k.weight": "norm_k.weight", + "attn2.to_q.weight": "attn2.to_q.weight", + "attn2.to_q.bias": "attn2.to_q.bias", + "attn2.to_k.weight": "attn2.to_k.weight", + "attn2.to_k.bias": "attn2.to_k.bias", + "attn2.to_v.weight": "attn2.to_v.weight", + "attn2.to_v.bias": "attn2.to_v.bias", + "attn2.to_out.0.weight": "attn2.to_out.weight", + "attn2.to_out.0.bias": "attn2.to_out.bias", + "attn2.norm_q.weight": "attn2.norm_q.weight", + "attn2.norm_k.weight": "attn2.norm_k.weight", + "ffn.net.0.proj.weight": "ffn.fc_in.weight", + "ffn.net.0.proj.bias": "ffn.fc_in.bias", + "ffn.net.2.weight": "ffn.fc_out.weight", + "ffn.net.2.bias": "ffn.fc_out.bias", + "norm2.weight": "self_attn_residual_norm.norm.weight", + "norm2.bias": "self_attn_residual_norm.norm.bias", +} + + +def mlx_block_weights_from_diffusers_safetensors( + checkpoint_path: str | Path, + *, + block_index: int = 0, + quantization: str | MLXQuantizationSpec | None = None, + dtype=None, +) -> dict[str, mx.array]: + """Load one Diffusers-format Wan block into the MLX dense-block key layout.""" + from safetensors import safe_open + + prefix = f"blocks.{block_index}." + key_map = _WAN_BLOCK_KEY_MAP + + spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None + or isinstance(quantization, str)) else quantization + ensure_quantization_supported(spec) + matrix_targets = {target for target in key_map.values() if target.endswith(".weight") and "norm" not in target} + weights = {} + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle: + available = set(handle.keys()) + for source_name, target_name in key_map.items(): + full = prefix + source_name + if full not in available: + # Biases are optional: e.g. Wan2.1-14B has bias-free attention/FFN. + # The block forward already fetches biases via ``.get(...)``. + if source_name.endswith(".bias"): + continue + raise KeyError(f"missing required block weight: {full}") + array = _load_mx_array_from_safetensor(handle, full, dtype) + loaded = quantize_matrix(array, spec) if target_name in matrix_targets else array + _eval_loaded_weight(loaded) + weights[target_name] = loaded + del array + return weights + + +def mlx_dit_from_diffusers_safetensors( + checkpoint_path: str | Path, + config_path: str | Path, + *, + dtype: str = "fp16", + num_blocks: int | None = None, + quantization: str | MLXQuantizationSpec | None = None, + compile: bool = False, +) -> MLXWanDiT: + import mlx.core as mx + from safetensors import safe_open + + config = json.loads(Path(config_path).read_text()) + total_blocks = int(config["num_layers"]) + if num_blocks is None: + num_blocks = total_blocks + cast_dtype = {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[dtype] + spec = MLXQuantizationSpec.from_name(quantization) if (quantization is None + or isinstance(quantization, str)) else quantization + ensure_quantization_supported(spec) + + top_level_names = [ + "patch_embedding.weight", + "patch_embedding.bias", + "condition_embedder.time_embedder.linear_1.weight", + "condition_embedder.time_embedder.linear_1.bias", + "condition_embedder.time_embedder.linear_2.weight", + "condition_embedder.time_embedder.linear_2.bias", + "condition_embedder.time_proj.weight", + "condition_embedder.time_proj.bias", + "condition_embedder.text_embedder.linear_1.weight", + "condition_embedder.text_embedder.linear_1.bias", + "condition_embedder.text_embedder.linear_2.weight", + "condition_embedder.text_embedder.linear_2.bias", + "scale_shift_table", + "proj_out.weight", + "proj_out.bias", + ] + weights = {} + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle: + available = set(handle.keys()) + for name in top_level_names: + if name not in available: + if name.endswith(".bias"): + continue + raise KeyError(f"missing required weight: {name}") + array = _load_mx_array_from_safetensor(handle, name, cast_dtype) + if name == "patch_embedding.weight": + array = array.reshape(int(config["num_attention_heads"]) * int(config["attention_head_dim"]), -1) + if name.endswith(".weight") and name not in {"scale_shift_table"}: + loaded = quantize_matrix(array, spec) + else: + loaded = array + _eval_loaded_weight(loaded) + weights[name] = loaded + del array + + blocks = [] + for block_index in range(num_blocks): + block_weights = mlx_block_weights_from_diffusers_safetensors( + checkpoint_path, + block_index=block_index, + quantization=spec, + dtype=cast_dtype, + ) + block_weights = { + name: (value if isinstance(value, QuantizedMatrix) else value.astype(cast_dtype)) + for name, value in block_weights.items() + } + for value in block_weights.values(): + _eval_loaded_weight(value) + blocks.append( + MLXWanTransformerBlock( + block_weights, + dim=int(config["num_attention_heads"]) * int(config["attention_head_dim"]), + ffn_dim=int(config["ffn_dim"]), + num_heads=int(config["num_attention_heads"]), + eps=float(config["eps"]), + )) + return MLXWanDiT(weights, blocks, config, compile=compile) + + +def torch_block_state_from_diffusers_safetensors( + checkpoint_path: str | Path, + *, + block_index: int = 0, +) -> dict[str, torch.Tensor]: + """Load one Diffusers-format Wan block into FastVideo's dense block keys.""" + from safetensors import safe_open + + prefix = f"blocks.{block_index}." + key_map = _WAN_BLOCK_KEY_MAP + + state = {} + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle: + available = set(handle.keys()) + for source_name, target_name in key_map.items(): + full = prefix + source_name + if full not in available: + # Biases are optional: e.g. Wan2.1-14B has bias-free attention/FFN. + if source_name.endswith(".bias"): + continue + raise KeyError(f"missing required block weight: {full}") + state[target_name] = handle.get_tensor(full).float() + return state diff --git a/fastvideo/mlx_runtime/frame_upsample.py b/fastvideo/mlx_runtime/frame_upsample.py new file mode 100644 index 000000000..ac2e077b5 --- /dev/null +++ b/fastvideo/mlx_runtime/frame_upsample.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Pixel-space spatial resampling for decoded MLX Wan frames. + +Spatial fast mode denoises on a smaller latent grid and has to get back to +the requested output size. That resize belongs *here* — after the VAE +decode — and not in latent space. + +A Wan latent cell is a learned code for an 8x8 (Wan2.1) or 16x16 (Wan2.2) +pixel block, not a low-pass sample of the image. Linearly blending two +adjacent codes does not produce the code of the blended blocks; it produces +a vector the decoder was never trained on. The decoder answers with smeared, +ringing texture laid over otherwise-correct structure — the silhouette +survives, the detail turns to haze. Measured on Wan2.1-1.3B at 480x832, a +2x bilinear latent upsample destroys 62% of the latent's high-frequency +energy while leaving its overall magnitude intact, which is exactly the +signature of that veil. + +Resampling decoded RGB frames has none of that problem: an image *is* a +sampled 2-D signal, so Lanczos/cubic interpolation is the operation it was +defined for. The result is soft — it carries stage-1's real detail budget +and no more — but it is clean and coherent. +""" + +from __future__ import annotations + +from collections.abc import Iterable + +import numpy as np + +# Pixel-space interpolation kernels, best-quality first. ``lanczos`` is the +# default: it holds edges better than cubic at 2x with no visible ringing on +# decoder output, which is already band-limited. +PIXEL_UPSAMPLE_MODES = ("lanczos", "cubic", "bilinear", "nearest") + +DEFAULT_PIXEL_UPSAMPLE_MODE = "lanczos" + + +def _interpolation_flag(mode: str) -> int: + """ + Map a pixel upsample mode name onto its OpenCV interpolation flag. + + Parameters: + mode (str): One of :data:`PIXEL_UPSAMPLE_MODES`. + + Returns: + int: The matching ``cv2.INTER_*`` flag. + + Raises: + ValueError: If the mode is not a supported pixel upsample mode. + """ + import cv2 + + flags = { + "lanczos": cv2.INTER_LANCZOS4, + "cubic": cv2.INTER_CUBIC, + "bilinear": cv2.INTER_LINEAR, + "nearest": cv2.INTER_NEAREST, + } + try: + return flags[mode] + except KeyError: + raise ValueError(f"Unsupported pixel upsample mode: {mode!r} " + f"(expected one of {', '.join(PIXEL_UPSAMPLE_MODES)})") from None + + +def unsharp(frame: np.ndarray, amount: float) -> np.ndarray: + """Light unsharp mask, used to counter resampling / optical-flow softening. + + Parameters: + frame (np.ndarray): HxWx3 uint8 RGB frame. + amount (float): Strength; ``0`` returns the frame unchanged. + + Returns: + np.ndarray: A new frame; the input is never modified in place. + """ + if amount <= 0.0: + return frame + import cv2 + + blur = cv2.GaussianBlur(frame, (0, 0), 1.0) + return cv2.addWeighted(frame, 1.0 + amount, blur, -amount, 0) + + +def upsample_frame( + frame: np.ndarray, + *, + width: int, + height: int, + mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE, + sharpen: float = 0.0, +) -> np.ndarray: + """ + Resample one decoded RGB frame to the target pixel size. + + Parameters: + frame (np.ndarray): HxWx3 uint8 RGB frame. + width (int): Target width in pixels. + height (int): Target height in pixels. + mode (str): Interpolation kernel, one of :data:`PIXEL_UPSAMPLE_MODES`. + sharpen (float): Unsharp strength applied after the resize. + + Returns: + np.ndarray: A new frame at ``height x width``; already-correct sizes + are still passed through ``sharpen``. + + Raises: + ValueError: If the frame is not HxWx3, or the target size is not positive. + """ + import cv2 + + array = np.asarray(frame) + if array.ndim != 3 or array.shape[2] != 3: + raise ValueError(f"frame must have shape HxWx3, got {array.shape}") + if width <= 0 or height <= 0: + raise ValueError(f"target size must be positive, got {width}x{height}") + if array.dtype != np.uint8: + array = np.clip(array, 0, 255).astype(np.uint8) + + if (array.shape[0], array.shape[1]) != (height, width): + array = cv2.resize(array, (width, height), interpolation=_interpolation_flag(mode)) + return unsharp(array, sharpen) + + +def upsample_frames( + frames: Iterable[np.ndarray], + *, + width: int, + height: int, + mode: str = DEFAULT_PIXEL_UPSAMPLE_MODE, + sharpen: float = 0.0, +) -> list[np.ndarray]: + """ + Resample every decoded frame to the target pixel size. + + Parameters: + frames (Iterable[np.ndarray]): Decoded HxWx3 uint8 RGB frames. + width (int): Target width in pixels. + height (int): Target height in pixels. + mode (str): Interpolation kernel, one of :data:`PIXEL_UPSAMPLE_MODES`. + sharpen (float): Unsharp strength applied after each resize. + + Returns: + list[np.ndarray]: New frames at the target size, in input order. + """ + return [upsample_frame(frame, width=width, height=height, mode=mode, sharpen=sharpen) for frame in frames] + + +__all__ = [ + "DEFAULT_PIXEL_UPSAMPLE_MODE", + "PIXEL_UPSAMPLE_MODES", + "unsharp", + "upsample_frame", + "upsample_frames", +] diff --git a/fastvideo/mlx_runtime/memory.py b/fastvideo/mlx_runtime/memory.py new file mode 100644 index 000000000..416d7cf8e --- /dev/null +++ b/fastvideo/mlx_runtime/memory.py @@ -0,0 +1,243 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Memory-tier helpers for Apple Silicon MLX/MPS experiments. + +macOS does not expose a perfect "pretend this machine only has 16 GB unified +memory" switch. MLX can cap the allocator used by the Apple-native DiT path, +and PyTorch MPS exposes process-level watermark environment variables for the +hybrid prompt/decode stages. Applying both gives benchmark and generation +entrypoints a practical, explicit way to exercise memory-tier presets. +""" + +from __future__ import annotations + +import argparse +import gc +import os +from dataclasses import dataclass, field +from typing import Any + +GIB = 1024**3 + + +@dataclass(frozen=True) +class AppliedMemoryLimits: + """Memory limits applied for one Apple Silicon benchmark/generation process.""" + + mlx_memory_limit_gib: float | None = None + mlx_cache_limit_gib: float | None = None + mlx_disable_cache: bool = False + mlx_wired_limit_gib: float | None = None + torch_mps_high_watermark_ratio: float | None = None + torch_mps_low_watermark_ratio: float | None = None + applied_bytes: dict[str, int] = field(default_factory=dict) + previous_bytes: dict[str, int] = field(default_factory=dict) + errors: dict[str, str] = field(default_factory=dict) + + def as_metrics(self) -> dict[str, int | float | str | bool | None]: + """Flatten the configured memory limits, applied values, previous values, and errors into a metrics dictionary. + + Returns: + dict[str, int | float | str | bool | None]: Metrics keyed by limit names and their corresponding values. + """ + metrics: dict[str, int | float | str | bool | None] = { + "mlx_memory_limit_gib": self.mlx_memory_limit_gib, + "mlx_cache_limit_gib": self.mlx_cache_limit_gib, + "mlx_disable_cache": self.mlx_disable_cache, + "mlx_wired_limit_gib": self.mlx_wired_limit_gib, + "torch_mps_high_watermark_ratio": self.torch_mps_high_watermark_ratio, + "torch_mps_low_watermark_ratio": self.torch_mps_low_watermark_ratio, + } + for name, value in self.applied_bytes.items(): + metrics[f"{name}_bytes"] = value + for name, value in self.previous_bytes.items(): + metrics[f"previous_{name}_bytes"] = value + for name, error in self.errors.items(): + metrics[f"{name}_error"] = error + return metrics + + +def gib_to_bytes(value: float | None) -> int | None: + """ + Convert a positive memory limit from GiB to bytes. + + Parameters: + value (float | None): Memory limit in GiB, or `None` when unset. + + Returns: + int | None: The memory limit in bytes, or `None` when no limit is provided. + + Raises: + ValueError: If `value` is zero or negative. + """ + if value is None: + return None + if value <= 0: + raise ValueError(f"Memory limit must be positive GiB, got {value}") + return int(value * GIB) + + +def cleanup_mlx(mx_module: Any | None = None) -> None: + """Collect unreachable MLX objects, then release their allocator cache.""" + if mx_module is None: + import mlx.core as mx + + mx_module = mx + gc.collect() + mx_module.clear_cache() + + +def cleanup_torch_mps(torch_module: Any | None = None) -> None: + """Collect unreachable Torch objects, then release the MPS allocator cache.""" + if torch_module is None: + import torch + + torch_module = torch + gc.collect() + if torch_module.backends.mps.is_available(): + torch_module.mps.empty_cache() + + +def _set_mps_env(name: str, value: float | None) -> float | None: + """Set a PyTorch MPS watermark environment variable. + + Parameters: + name (str): Name of the environment variable to set. + value (float | None): Watermark ratio, or `None` to leave the variable unchanged. + + Returns: + float | None: The configured watermark ratio, or `None` when no value is provided. + + Raises: + ValueError: If `value` is negative. + """ + if value is None: + return None + if value < 0: + raise ValueError(f"{name} must be non-negative, got {value}") + os.environ[name] = str(value) + return value + + +def apply_memory_limits( + *, + mlx_memory_limit_gib: float | None = None, + mlx_cache_limit_gib: float | None = None, + mlx_disable_cache: bool = False, + mlx_wired_limit_gib: float | None = None, + torch_mps_high_watermark_ratio: float | None = None, + torch_mps_low_watermark_ratio: float | None = None, + mx_module: Any | None = None, +) -> AppliedMemoryLimits: + """Apply optional MLX allocator limits and PyTorch MPS watermarks. + + PyTorch reads MPS watermark variables when the MPS backend initializes, so + call this before importing PyTorch. Specifying only a high watermark sets the + low watermark to ``0.0``. MLX limit-setting failures are recorded in the + result and do not prevent other limits from being applied. + + Parameters: + mlx_memory_limit_gib (float | None): Maximum MLX memory in GiB. + mlx_cache_limit_gib (float | None): Maximum MLX cache size in GiB. + mlx_disable_cache (bool): Whether to disable the MLX cache. + mlx_wired_limit_gib (float | None): Maximum MLX wired memory in GiB. + torch_mps_high_watermark_ratio (float | None): PyTorch MPS high watermark + ratio. + torch_mps_low_watermark_ratio (float | None): PyTorch MPS low watermark + ratio. + + Returns: + AppliedMemoryLimits: Configured values, applied and previous MLX byte + limits, MPS watermark values, and per-limit errors. + """ + if torch_mps_high_watermark_ratio is not None and torch_mps_low_watermark_ratio is None: + torch_mps_low_watermark_ratio = 0.0 + + high = _set_mps_env("PYTORCH_MPS_HIGH_WATERMARK_RATIO", torch_mps_high_watermark_ratio) + low = _set_mps_env("PYTORCH_MPS_LOW_WATERMARK_RATIO", torch_mps_low_watermark_ratio) + + memory_bytes = gib_to_bytes(mlx_memory_limit_gib) + cache_bytes = 0 if mlx_disable_cache else gib_to_bytes(mlx_cache_limit_gib) + wired_bytes = gib_to_bytes(mlx_wired_limit_gib) + + applied: dict[str, int] = {} + previous: dict[str, int] = {} + errors: dict[str, str] = {} + if memory_bytes is not None or cache_bytes is not None or wired_bytes is not None: + if mx_module is None: + import mlx.core as mx + + mx_module = mx + + # Apply each limit independently; record failures without stopping. + limits = [ + ("mlx_memory_limit", memory_bytes, mx_module.set_memory_limit), + ("mlx_cache_limit", cache_bytes, mx_module.set_cache_limit), + ("mlx_wired_limit", wired_bytes, mx_module.set_wired_limit), + ] + for name, value, setter in limits: + if value is not None: + try: + previous[name] = int(setter(value)) + applied[name] = value + except Exception as exc: # noqa: BLE001 - macOS/system-limit dependent. + errors[name] = f"{type(exc).__name__}: {exc}" + + return AppliedMemoryLimits( + mlx_memory_limit_gib=mlx_memory_limit_gib, + mlx_cache_limit_gib=mlx_cache_limit_gib, + mlx_disable_cache=mlx_disable_cache, + mlx_wired_limit_gib=mlx_wired_limit_gib, + torch_mps_high_watermark_ratio=high, + torch_mps_low_watermark_ratio=low, + applied_bytes=applied, + previous_bytes=previous, + errors=errors, + ) + + +def add_memory_limit_args( + parser: argparse.ArgumentParser, + *, + mlx_memory_limit_gib: float | None = None, + mlx_cache_limit_gib: float | None = None, + mlx_disable_cache: bool = False, + mlx_wired_limit_gib: float | None = None, + torch_mps_high_watermark_ratio: float | None = None, + torch_mps_low_watermark_ratio: float | None = None, +) -> None: + """ + Add configurable Apple Silicon memory-limit options to an argument parser. + + Parameters: + parser (argparse.ArgumentParser): Parser to which the options are added. + mlx_memory_limit_gib (float | None): Default MLX memory limit in GiB. + mlx_cache_limit_gib (float | None): Default MLX cache limit in GiB. + mlx_disable_cache (bool): Whether the cache limit defaults to zero. + mlx_wired_limit_gib (float | None): Default MLX wired-memory limit in GiB. + torch_mps_high_watermark_ratio (float | None): Default PyTorch MPS high-watermark ratio. + torch_mps_low_watermark_ratio (float | None): Default PyTorch MPS low-watermark ratio. + """ + parser.add_argument("--mlx-memory-limit-gib", + type=float, + default=mlx_memory_limit_gib, + help="Set MLX memory limit in GiB for memory-tier testing (DiT path).") + parser.add_argument("--mlx-cache-limit-gib", + type=float, + default=mlx_cache_limit_gib, + help="Set MLX cache limit in GiB. Use --mlx-disable-cache to force 0.") + parser.add_argument("--mlx-disable-cache", + action="store_true", + default=mlx_disable_cache, + help="Set MLX cache limit to 0 for stricter memory-tier tests.") + parser.add_argument("--mlx-wired-limit-gib", + type=float, + default=mlx_wired_limit_gib, + help="Set MLX wired-memory limit in GiB where supported by macOS/MLX.") + parser.add_argument("--torch-mps-high-watermark-ratio", + type=float, + default=torch_mps_high_watermark_ratio, + help="Set PYTORCH_MPS_HIGH_WATERMARK_RATIO before importing torch.") + parser.add_argument("--torch-mps-low-watermark-ratio", + type=float, + default=torch_mps_low_watermark_ratio, + help="Set PYTORCH_MPS_LOW_WATERMARK_RATIO before importing torch.") diff --git a/fastvideo/mlx_runtime/prompt_cache.py b/fastvideo/mlx_runtime/prompt_cache.py new file mode 100644 index 000000000..32e440f52 --- /dev/null +++ b/fastvideo/mlx_runtime/prompt_cache.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Best-effort prompt-embedding cache shared by the MLX entrypoints.""" + +from __future__ import annotations + +import hashlib +import io +import json +import logging +import tempfile +from pathlib import Path + +import numpy as np + +logger = logging.getLogger(__name__) + + +def fingerprint_digest(fingerprint: dict[str, object]) -> str: + payload = json.dumps(fingerprint, sort_keys=True, separators=(",", ":")) + return hashlib.sha256(payload.encode("utf-8")).hexdigest() + + +def text_encoder_fingerprint(model_root: Path) -> dict[str, object]: + """Return a cheap identity for the tokenizer and text-encoder files.""" + root = model_root.resolve() + components = [path for name in ("tokenizer", "text_encoder") if (path := root / name).is_dir()] + scan_roots = components or [root] + files: list[list[object]] = [] + complete = True + try: + for scan_root in scan_roots: + for path in sorted(scan_root.rglob("*")): + try: + if not path.is_file(): + continue + stat = path.stat() + files.append([ + path.relative_to(root).as_posix(), + stat.st_size, + stat.st_mtime_ns, + stat.st_ctime_ns, + ]) + except OSError: + complete = False + except OSError: + complete = False + # ponytail: metadata avoids hashing multi-GB weights; use a model manifest + # if supported workflows ever preserve size, mtime, and ctime while mutating. + return {"root": str(root), "files": files, "complete": complete} + + +def prompt_cache_meta_path(cache_path: Path) -> Path: + return cache_path.with_suffix(cache_path.suffix + ".json") + + +def _fingerprint_is_complete(fingerprint: dict[str, object]) -> bool: + text_encoder = fingerprint.get("text_encoder") + return not isinstance(text_encoder, dict) or text_encoder.get("complete") is not False + + +def load_prompt_cache( + cache_path: Path | None, + fingerprint: dict[str, object], +) -> np.ndarray | None: + """Load a matching cache entry, treating every cache failure as a miss.""" + if cache_path is None or not _fingerprint_is_complete(fingerprint): + return None + try: + metadata = json.loads(prompt_cache_meta_path(cache_path).read_text()) + if not isinstance(metadata, dict): + return None + if metadata.get("fingerprint_sha256") != fingerprint_digest(fingerprint): + return None + payload = cache_path.read_bytes() + if metadata.get("data_sha256") != hashlib.sha256(payload).hexdigest(): + return None + array = np.load(io.BytesIO(payload), allow_pickle=False) + return array if isinstance(array, np.ndarray) else None + except (EOFError, OSError, UnicodeError, ValueError) as exc: + logger.info("Prompt cache read skipped for %s: %s", cache_path, exc) + return None + + +def _atomic_write(path: Path, payload: bytes) -> None: + temp_path: Path | None = None + try: + with tempfile.NamedTemporaryFile( + mode="wb", + dir=path.parent, + prefix=f".{path.name}.", + suffix=".tmp", + delete=False, + ) as handle: + temp_path = Path(handle.name) + handle.write(payload) + temp_path.replace(path) + finally: + if temp_path is not None: + temp_path.unlink(missing_ok=True) + + +def save_prompt_cache( + cache_path: Path | None, + embeds: np.ndarray, + fingerprint: dict[str, object], +) -> bool: + """Atomically publish an integrity-bound cache entry when possible.""" + if cache_path is None or not _fingerprint_is_complete(fingerprint): + return False + try: + cache_path.parent.mkdir(parents=True, exist_ok=True) + buffer = io.BytesIO() + np.save(buffer, np.asarray(embeds), allow_pickle=False) + payload = buffer.getvalue() + metadata = (json.dumps( + { + "fingerprint_sha256": fingerprint_digest(fingerprint), + "data_sha256": hashlib.sha256(payload).hexdigest(), + "fingerprint": fingerprint, + }, + indent=2) + "\n").encode("utf-8") + # Publish data first. Until metadata follows, old metadata's data digest + # makes the torn pair a harmless miss rather than a stale cache hit. + _atomic_write(cache_path, payload) + _atomic_write(prompt_cache_meta_path(cache_path), metadata) + return True + except (OSError, TypeError, ValueError) as exc: + logger.info("Prompt cache write skipped for %s: %s", cache_path, exc) + return False diff --git a/fastvideo/mlx_runtime/prompt_enhance.py b/fastvideo/mlx_runtime/prompt_enhance.py new file mode 100644 index 000000000..60a4ed4bb --- /dev/null +++ b/fastvideo/mlx_runtime/prompt_enhance.py @@ -0,0 +1,454 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Local prompt enrichment for the MLX Wan runtime (H3 Context-IR-style). + +Wan's training captions are long and cinematic; short user prompts leave +quality on the table. This module expands a raw prompt into Wan-style +shot language **on device** — no remote API, no training. + +Backends (first match wins): + +1. **mlx-lm** — optional local LLM (``--enhance-prompt-model``). +2. **template** — deterministic cinematic expansion (always available). + +System-prompt contract matches the streaming server's enhancer defaults +in ``fastvideo/entrypoints/streaming/prompt/enhancer.py`` so remote and +local paths stay interchangeable. +""" + +from __future__ import annotations + +import hashlib +import json +import re +import time +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +from fastvideo.logger import init_logger + +logger = init_logger(__name__) + +# Keep in lockstep with streaming PromptEnhancer defaults (enhance op). +DEFAULT_ENHANCE_SYSTEM_PROMPT = ("You are a prompt enhancer for cinematic video generation. Given " + "a user prompt, produce an enhanced prompt that is more vivid, " + "specific, and concrete. Keep the subject intact; add lighting, " + "camera, and motion detail. Reply with just the enhanced prompt.") + +# Small default that fits 16 GB Macs alongside the 1.3B DiT when the user +# opts into mlx-lm. Override with --enhance-prompt-model. +DEFAULT_MLX_LM_MODEL = "mlx-community/Qwen2.5-0.5B-Instruct-4bit" + +_CAMERA_CUES = ( + "cinematic", + "camera", + "lens", + "shot", + "bokeh", + "dolly", + "tracking", + "close-up", + "wide shot", + "handheld", + "steadicam", +) +_LIGHT_CUES = ( + "light", + "lighting", + "sun", + "golden hour", + "neon", + "rim light", + "softbox", + "overcast", + "moonlight", + "volumetric", +) +_MOTION_CUES = ( + "moving", + "motion", + "walk", + "run", + "flies", + "flying", + "drifts", + "sails", + "flows", + "pan", + "tilt", + "zoom", +) + + +@dataclass(frozen=True) +class EnhanceResult: + """Outcome of a prompt enrichment call.""" + + original: str + enhanced: str + backend: str + elapsed_s: float + model: str | None = None + + @property + def changed(self) -> bool: + """Indicates whether the enhanced prompt differs from the original after trimming surrounding whitespace. + + Returns: + bool: `True` if the prompts differ, `False` otherwise. + """ + return self.enhanced.strip() != self.original.strip() + + +def _normalize_user_prompt(prompt: str) -> str: + """ + Normalize a user prompt for enhancement. + + Parameters: + prompt (str): User-provided prompt text. + + Returns: + str: The prompt with leading and trailing whitespace removed and internal whitespace collapsed. + + Raises: + ValueError: If the prompt is empty after whitespace normalization. + """ + text = " ".join(prompt.strip().split()) + if not text: + raise ValueError("prompt must be non-empty") + return text + + +def _already_rich(prompt: str) -> bool: + """ + Determine whether a prompt already contains substantial camera and lighting detail. + + Returns: + bool: `true` if the prompt is at least 160 characters long and includes camera and lighting cues, `false` otherwise. + """ + lower = prompt.lower() + has_camera = any(c in lower for c in _CAMERA_CUES) + has_light = any(c in lower for c in _LIGHT_CUES) + return len(prompt) >= 160 and has_camera and has_light + + +def enhance_prompt_template(prompt: str) -> str: + """ + Expand a prompt with cinematic camera, lighting, motion, and visual-quality details. + + Rich prompts are preserved, while thinner prompts receive deterministic enhancements + without changing their subject. + + Returns: + str: The original or expanded prompt with normalized whitespace and punctuation. + """ + text = _normalize_user_prompt(prompt) + if _already_rich(text): + return text + + lower = text.lower() + parts = [text.rstrip(".")] + + if not any(c in lower for c in _CAMERA_CUES): + parts.append("shot on a 35mm anamorphic lens, gentle handheld micro-movement, " + "shallow depth of field") + if not any(c in lower for c in _LIGHT_CUES): + parts.append("natural cinematic lighting with soft volumetric haze and subtle " + "rim light separating subject from background") + if not any(c in lower for c in _MOTION_CUES): + parts.append("smooth continuous motion with grounded physics") + + parts.append("highly detailed, coherent temporal continuity, film grain, " + "color graded like a contemporary drama") + enhanced = ", ".join(parts) + # Single trailing period; collapse duplicate whitespace. + enhanced = re.sub(r"\s+", " ", enhanced).strip() + if not enhanced.endswith("."): + enhanced += "." + return enhanced + + +def enhance_prompt_mlx_lm( + prompt: str, + *, + model: str = DEFAULT_MLX_LM_MODEL, + system_prompt: str = DEFAULT_ENHANCE_SYSTEM_PROMPT, + max_tokens: int = 128, + temp: float = 0.6, +) -> str: + """ + Enhance a user prompt with a locally hosted mlx-lm instruction model. + + Parameters: + prompt (str): The prompt to enhance. + model (str): The mlx-lm model identifier or path. + system_prompt (str): Instructions that guide prompt enhancement. + max_tokens (int): Maximum number of tokens to generate. + temp (float): Sampling temperature for generation. + + Returns: + str: The enhanced prompt. + + Raises: + RuntimeError: If mlx-lm is unavailable or produces an empty result. + """ + try: + from mlx_lm import generate, load + except ImportError as exc: # pragma: no cover - optional dep + raise RuntimeError("mlx-lm is not installed. `uv pip install mlx-lm` or use " + "--enhance-prompt-backend template.") from exc + + text = _normalize_user_prompt(prompt) + logger.info("[MLX enhance] loading %s", model) + mlx_model, tokenizer = load(model) + + messages = [ + { + "role": "system", + "content": system_prompt + }, + { + "role": "user", + "content": text + }, + ] + if hasattr(tokenizer, "apply_chat_template"): + chat = tokenizer.apply_chat_template( + messages, + tokenize=False, + add_generation_prompt=True, + ) + else: # pragma: no cover - ancient tokenizers + chat = f"{system_prompt}\n\nUser: {text}\nAssistant:" + + raw = generate( + mlx_model, + tokenizer, + prompt=chat, + max_tokens=max_tokens, + temp=temp, + verbose=False, + ) + enhanced = _clean_llm_output(raw, original=text) + if not enhanced: + raise RuntimeError("mlx-lm returned an empty enhance result") + return enhanced + + +def _clean_llm_output(raw: str, *, original: str) -> str: + """ + Clean generated prompt text and fall back to the original when the result is too short. + + Parameters: + raw (str): Raw text produced by the language model. + original (str): Original prompt used as the fallback value. + + Returns: + str: Cleaned first paragraph of the generated text, or the original prompt when the generated text is too short. + """ + text = raw.strip() + # Drop common prefatory phrases. + for prefix in ( + "enhanced prompt:", + "here's the enhanced prompt:", + "here is the enhanced prompt:", + "sure:", + "sure,", + ): + if text.lower().startswith(prefix): + text = text[len(prefix):].strip() + # Keep first non-empty paragraph only. + para = text.split("\n\n")[0].strip() + para = " ".join(para.split()) + if len(para) < max(12, len(original) // 4): + return original + return para + + +def enhance_prompt( + prompt: str, + *, + backend: str = "auto", + model: str | None = None, + system_prompt: str = DEFAULT_ENHANCE_SYSTEM_PROMPT, + max_tokens: int = 128, +) -> EnhanceResult: + """Enhance a prompt using the selected backend, falling back to a deterministic template when configured for automatic selection. + + Parameters: + prompt (str): The prompt to enhance. + backend (str): The enhancement backend: ``"auto"``, ``"mlx-lm"``, or ``"template"``. + model (str | None): The MLX language model to use. + system_prompt (str): Instructions provided to the MLX language model. + max_tokens (int): Maximum number of tokens generated by the MLX language model. + + Returns: + EnhanceResult: The original and enhanced prompts, selected backend, timing information, and model metadata. + + Raises: + ValueError: If the prompt is empty or the backend is unsupported. + Exception: If the explicitly selected ``"mlx-lm"`` backend fails. + """ + text = _normalize_user_prompt(prompt) + backend_norm = (backend or "auto").lower() + if backend_norm not in {"auto", "mlx-lm", "template"}: + raise ValueError(f"Unknown enhance backend: {backend}") + + start = time.perf_counter() + used_model: str | None = None + + if backend_norm in {"auto", "mlx-lm"}: + try: + used_model = model or DEFAULT_MLX_LM_MODEL + enhanced = enhance_prompt_mlx_lm( + text, + model=used_model, + system_prompt=system_prompt, + max_tokens=max_tokens, + ) + return EnhanceResult( + original=text, + enhanced=enhanced, + backend="mlx-lm", + elapsed_s=time.perf_counter() - start, + model=used_model, + ) + except Exception as exc: + if backend_norm == "mlx-lm": + raise + logger.info( + "[MLX enhance] mlx-lm unavailable (%s); using template backend", + exc, + ) + + enhanced = enhance_prompt_template(text) + return EnhanceResult( + original=text, + enhanced=enhanced, + backend="template", + elapsed_s=time.perf_counter() - start, + model=None, + ) + + +def enhance_cache_path( + prompt: str, + *, + backend: str, + model: str | None, + cache_dir: Path | None = None, +) -> Path: + """Content-addressed cache file for an enhanced prompt string.""" + root = cache_dir or (Path.home() / ".cache" / "fastvideo" / "enhanced_prompts") + key = "\0".join([prompt, backend, model or ""]) + digest = hashlib.sha256(key.encode("utf-8")).hexdigest()[:24] + return root / f"{digest}.json" + + +def load_or_enhance_prompt( + prompt: str, + *, + backend: str = "auto", + model: str | None = None, + system_prompt: str = DEFAULT_ENHANCE_SYSTEM_PROMPT, + max_tokens: int = 128, + cache: bool = True, + cache_dir: Path | None = None, +) -> EnhanceResult: + """ + Enhance a prompt, reusing a cached result when available. + + Parameters: + prompt (str): The prompt to enhance. + backend (str): Enhancement backend to use. + model (str | None): Optional model identifier. + system_prompt (str): System prompt for model-based enhancement. + max_tokens (int): Maximum number of tokens generated by the model. + cache (bool): Whether to read and write the on-disk cache. + cache_dir (Path | None): Optional directory for cached results. + + Returns: + EnhanceResult: The enhanced prompt and backend metadata. Cached results are marked with the ``"cache"`` backend. + """ + text = _normalize_user_prompt(prompt) + path = enhance_cache_path(text, backend=backend, model=model, cache_dir=cache_dir) + if cache and path.is_file(): + try: + payload = json.loads(path.read_text()) + return EnhanceResult( + original=str(payload.get("original", text)), + enhanced=str(payload["enhanced"]), + # Mark cache hits explicitly so metrics/logs can distinguish + # a free replay from a fresh template/mlx-lm call. + backend="cache", + elapsed_s=0.0, + model=payload.get("model"), + ) + except (OSError, KeyError, json.JSONDecodeError): + pass + + result = enhance_prompt( + text, + backend=backend, + model=model, + system_prompt=system_prompt, + max_tokens=max_tokens, + ) + if cache: + try: + path.parent.mkdir(parents=True, exist_ok=True) + path.write_text( + json.dumps( + { + "original": result.original, + "enhanced": result.enhanced, + "backend": result.backend, + "model": result.model, + }, + indent=2, + )) + except OSError as exc: # pragma: no cover - cache is best-effort + logger.info("[MLX enhance] cache write skipped: %s", exc) + return result + + +def enhance_result_as_metrics(result: EnhanceResult | None) -> dict[str, Any]: + """ + Convert prompt enhancement results into metrics fields. + + Parameters: + result (EnhanceResult | None): The enhancement result, or `None` when no enhancement was performed. + + Returns: + dict[str, Any]: A metrics mapping containing enhancement status, backend metadata, timing, and original and enhanced prompts. + """ + if result is None: + return { + "enhance_prompt": False, + "enhance_backend": None, + "enhance_model": None, + "enhance_elapsed_s": None, + "prompt_original": None, + "prompt_enhanced": None, + } + return { + "enhance_prompt": True, + "enhance_backend": result.backend, + "enhance_model": result.model, + "enhance_elapsed_s": result.elapsed_s, + "prompt_original": result.original, + "prompt_enhanced": result.enhanced, + } + + +__all__ = [ + "DEFAULT_ENHANCE_SYSTEM_PROMPT", + "DEFAULT_MLX_LM_MODEL", + "EnhanceResult", + "enhance_cache_path", + "enhance_prompt", + "enhance_prompt_mlx_lm", + "enhance_prompt_template", + "enhance_result_as_metrics", + "load_or_enhance_prompt", +] diff --git a/fastvideo/mlx_runtime/quant_backends.py b/fastvideo/mlx_runtime/quant_backends.py new file mode 100644 index 000000000..b5d37709a --- /dev/null +++ b/fastvideo/mlx_runtime/quant_backends.py @@ -0,0 +1,284 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MLX block-scaled quantization backends (affine INT8, MXFP8/4, NVFP4). + +Isolated experiment module: probes which ``mx.quantize`` modes the installed +MLX build supports and exposes a thin wrapper around native quantized matmul. +Depends only on ``mlx.core`` and the standard library — do not import the rest +of FastVideo from here. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from enum import Enum +from typing import Final +from collections.abc import Mapping + +import mlx.core as mx + +# Probe matrix side length: must be divisible by every mode's group size +# (affine g64, mxfp* g32, nvfp4 g16). +_PROBE_DIM: Final[int] = 64 + + +class QuantBackend(str, Enum): + """Named MLX quantization backends evaluated for M5 Neural Accelerators.""" + + AFFINE_INT8_G64 = "affine_int8_g64" + MXFP8 = "mxfp8" + MXFP4 = "mxfp4" + NVFP4 = "nvfp4" + + +BACKENDS: Final[tuple[str, ...]] = tuple(b.value for b in QuantBackend) + +# Backend name -> kwargs for mx.quantize / mx.quantized_matmul. +# Affine baseline matches FastVideo DiT load path (INT8, group size 64). +# MX/NV block-scaled modes use MLX defaults (see mx.quantize docs). +_BACKEND_KWARGS: Final[Mapping[str, Mapping[str, object]]] = { + QuantBackend.AFFINE_INT8_G64.value: { + "mode": "affine", + "bits": 8, + "group_size": 64, + }, + QuantBackend.MXFP8.value: { + "mode": "mxfp8", + "bits": None, + "group_size": None, + }, + QuantBackend.MXFP4.value: { + "mode": "mxfp4", + "bits": None, + "group_size": None, + }, + QuantBackend.NVFP4.value: { + "mode": "nvfp4", + "bits": None, + "group_size": None, + }, +} + +_SUPPORT_CACHE: dict[str, bool] = {} +_SUPPORT_ERROR_CACHE: dict[str, str | None] = {} +_BYTES_CACHE: dict[str, float] = {} + + +@dataclass(frozen=True) +class QuantizedWeight: + """Packed quantized weight plus scales/biases for one backend.""" + + weight: mx.array + scales: mx.array + biases: mx.array | None + backend: str + mode: str + bits: int | None + group_size: int | None + # Original (rows, cols) of the fp weight, used for bytes-per-element. + orig_shape: tuple[int, int] + + +def _normalize_backend(backend: str) -> str: + """Normalize a quantization backend name and validate that it is supported. + + Parameters: + backend (str): Backend name to normalize. + + Returns: + str: The lowercase backend name without surrounding whitespace. + + Raises: + ValueError: If the backend name is unknown. + """ + name = backend.strip().lower() + if name not in _BACKEND_KWARGS: + known = ", ".join(BACKENDS) + raise ValueError(f"Unknown quant backend {backend!r}. Expected one of: {known}") + return name + + +def _kwargs_for(backend: str) -> dict[str, object]: + """Return the MLX quantization arguments configured for a backend. + + Parameters: + backend (str): Backend name to resolve. + + Returns: + dict[str, object]: Quantization arguments for the normalized backend. + """ + return dict(_BACKEND_KWARGS[_normalize_backend(backend)]) + + +def support_error(backend: str) -> str | None: + """ + Check whether a quantization backend is supported by the current MLX runtime. + + Parameters: + backend (str): Quantization backend name. + + Returns: + str | None: An error description when the backend is unsupported, or `None` when supported. + """ + name = _normalize_backend(backend) + if name in _SUPPORT_ERROR_CACHE: + return _SUPPORT_ERROR_CACHE[name] + + kwargs = _kwargs_for(name) + try: + w = mx.zeros((_PROBE_DIM, _PROBE_DIM), dtype=mx.float16) + quantized = mx.quantize( + w, + group_size=kwargs["group_size"], # type: ignore[arg-type] + bits=kwargs["bits"], # type: ignore[arg-type] + mode=str(kwargs["mode"]), + ) + w_q = quantized[0] + scales = quantized[1] + biases = quantized[2] if len(quantized) == 3 else None + x = mx.zeros((1, _PROBE_DIM), dtype=mx.float16) + y = mx.quantized_matmul( + x, + w_q, + scales, + biases, + transpose=True, + group_size=kwargs["group_size"], # type: ignore[arg-type] + bits=kwargs["bits"], # type: ignore[arg-type] + mode=str(kwargs["mode"]), + ) + mx.eval(y) + _SUPPORT_ERROR_CACHE[name] = None + _SUPPORT_CACHE[name] = True + except Exception as exc: # noqa: BLE001 - MLX raises varied types per mode/version. + msg = f"{type(exc).__name__}: {exc}" + _SUPPORT_ERROR_CACHE[name] = msg + _SUPPORT_CACHE[name] = False + return _SUPPORT_ERROR_CACHE[name] + + +def is_supported(backend: str) -> bool: + """Return True if the installed MLX build can quantize/matmul with ``backend``.""" + name = _normalize_backend(backend) + if name not in _SUPPORT_CACHE: + support_error(name) + return _SUPPORT_CACHE[name] + + +def quantize_weight(w: mx.array, backend: str) -> QuantizedWeight: + """ + Quantize a two-dimensional weight matrix using the specified native MLX backend. + + Parameters: + w (mx.array): The two-dimensional weight matrix to quantize. + backend (str): The quantization backend to use. + + Returns: + QuantizedWeight: The quantized weights and associated quantization metadata. + + Raises: + ValueError: If the backend is unknown, the weight is not two-dimensional, + or its last dimension is not divisible by the backend's group size. + RuntimeError: If the backend is unsupported by the installed MLX build. + """ + name = _normalize_backend(backend) + err = support_error(name) + if err is not None: + mlx_version = getattr(mx, "__version__", "unknown") + raise RuntimeError(f"Quant backend {name!r} is not supported by installed mlx " + f"({mlx_version}): {err}") + + if w.ndim != 2: + raise ValueError(f"quantize_weight expects a 2D weight, got shape {tuple(w.shape)}") + + rows, cols = int(w.shape[0]), int(w.shape[1]) + kwargs = _kwargs_for(name) + group_size = kwargs["group_size"] + # When group_size is None, MLX applies the mode default; only check when set. + if isinstance(group_size, int) and cols % group_size != 0: + raise ValueError(f"Weight last dim {cols} must be divisible by group_size={group_size} " + f"for backend {name!r}") + + quantized = mx.quantize( + w, + group_size=kwargs["group_size"], # type: ignore[arg-type] + bits=kwargs["bits"], # type: ignore[arg-type] + mode=str(kwargs["mode"]), + ) + w_q = quantized[0] + scales = quantized[1] + biases = quantized[2] if len(quantized) == 3 else None + eval_args = [w_q, scales] if biases is None else [w_q, scales, biases] + mx.eval(*eval_args) + + return QuantizedWeight( + weight=w_q, + scales=scales, + biases=biases, + backend=name, + mode=str(kwargs["mode"]), + bits=kwargs["bits"] if isinstance(kwargs["bits"], int) else None, + group_size=group_size if isinstance(group_size, int) else None, + orig_shape=(rows, cols), + ) + + +def quantized_matmul(x: mx.array, qw: QuantizedWeight) -> mx.array: + """Compute ``x @ w.T`` in the quantized domain via ``mx.quantized_matmul``.""" + return mx.quantized_matmul( + x, + qw.weight, + qw.scales, + qw.biases, + transpose=True, + group_size=qw.group_size, + bits=qw.bits, + mode=qw.mode, + ) + + +def _artifact_nbytes(qw: QuantizedWeight) -> int: + """ + Calculate the total storage size of a quantized weight artifact in bytes. + + Parameters: + qw (QuantizedWeight): Quantized weight artifact whose packed weights, scales, and optional biases are measured. + + Returns: + int: Total number of bytes used by the artifact's stored arrays. + """ + total = int(qw.weight.nbytes) + int(qw.scales.nbytes) + if qw.biases is not None: + total += int(qw.biases.nbytes) + return total + + +def bytes_per_weight(backend: str) -> float: + """ + Measure the effective storage cost of a quantized weight. + + Parameters: + backend (str): Quantization backend to measure. + + Returns: + float: Stored bytes per original weight element, including packed weights, + scales, and optional biases. + + Raises: + RuntimeError: If the backend is unsupported. + """ + name = _normalize_backend(backend) + if name in _BYTES_CACHE: + return _BYTES_CACHE[name] + + err = support_error(name) + if err is not None: + mlx_version = getattr(mx, "__version__", "unknown") + raise RuntimeError(f"Cannot measure bytes_per_weight for unsupported backend {name!r} " + f"(mlx {mlx_version}): {err}") + + probe = mx.zeros((_PROBE_DIM, _PROBE_DIM), dtype=mx.float16) + qw = quantize_weight(probe, name) + n_elem = qw.orig_shape[0] * qw.orig_shape[1] + value = _artifact_nbytes(qw) / float(n_elem) + _BYTES_CACHE[name] = value + return value diff --git a/fastvideo/mlx_runtime/refine.py b/fastvideo/mlx_runtime/refine.py new file mode 100644 index 000000000..81c627ed7 --- /dev/null +++ b/fastvideo/mlx_runtime/refine.py @@ -0,0 +1,690 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Two-pass spatial refine for the MLX Wan runtime (H3 / LTX-2 pattern). + +Biggest quality lever on Apple Silicon without a new model or training: +generate at base resolution, then run a second denoising pass with the +*same* DiT at a higher resolution. + +This is the MLX-side port of the CUDA refine template in +``fastvideo/pipelines/basic/ltx2/stages/ltx2_refine.py`` and the H3 +"base + regenerate" pattern documented in +``docs/design/mac_qad_two_product_strategy.md``: + +1. :func:`plan_refine_resolutions` — split the request into stage-1 + (base) and stage-2 (target) pixel sizes, validating VAE / patch + alignment the way :class:`LTX2RefineInitStage` does. +2. :func:`upsample_latents_spatial` — 2× (or N×) spatial upsample of + clean latents. Wan has no learned latent upsampler on Mac, so this + is bilinear over the H×W plane (temporal axis untouched) — same + role as LTX-2's ``upsample_video`` hand-off, without the learned + residual. +3. :func:`prepare_refine_latents` — upsample + re-noise the clean + stage-1 latents to the stage-2 sigma so the second denoise has + something to refine (mirrors :class:`LTX2UpsampleStage` + + ``apply_ltx2_gaussian_noiser``). +4. :func:`run_two_pass_dmd` — orchestrate stage-1 denoise → refine + hand-off → stage-2 denoise with the same model / prompt embeds. + +No LoRA swap, no dedicated SR weights, no new training — pure pipeline +work reusable by Wan2.1-14B and Wan2.2-5B on Apple Silicon. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any +from collections.abc import Callable, Sequence + +import numpy as np + +from fastvideo.logger import init_logger +from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, add_noise, dmd_step + +if TYPE_CHECKING: # pragma: no cover - typing only + import mlx.core as mx + +logger = init_logger(__name__) + +# Default stage-2 noise level when the caller does not supply a schedule. +# Matches the first entry of LTX-2's STAGE_2_DISTILLED_SIGMA_VALUES in spirit +# (start the refine denoise from a high-noise level) without hard-wiring the +# LTX-2 distilled grid onto Wan's flow-match schedule. +DEFAULT_REFINE_SIGMA = 0.909375 + + +@dataclass(frozen=True) +class RefinePlan: + """Resolved stage-1 / stage-2 geometry for a two-pass refine run.""" + + target_height: int + target_width: int + stage1_height: int + stage1_width: int + spatial_scale: int + vae_spatial_compression: int + vae_temporal_compression: int + num_frames: int + + @property + def stage1_latent_height(self) -> int: + """Return the stage-1 latent height after VAE spatial compression.""" + return self.stage1_height // self.vae_spatial_compression + + @property + def stage1_latent_width(self) -> int: + """Return the stage-one latent width after VAE spatial compression.""" + return self.stage1_width // self.vae_spatial_compression + + @property + def stage2_latent_height(self) -> int: + """Calculate the target-resolution latent height. + + Returns: + int: The target height divided by the VAE spatial compression factor. + """ + return self.target_height // self.vae_spatial_compression + + @property + def stage2_latent_width(self) -> int: + """Return the target image width in latent-space units.""" + return self.target_width // self.vae_spatial_compression + + @property + def latent_frames(self) -> int: + """Calculate the number of latent frames after VAE temporal compression. + + Returns: + int: The compressed latent frame count. + """ + return (self.num_frames - 1) // self.vae_temporal_compression + 1 + + +def plan_refine_resolutions( + *, + height: int, + width: int, + num_frames: int, + spatial_scale: int = 2, + vae_spatial_compression: int = 8, + vae_temporal_compression: int = 4, + patch_size: tuple[int, int, int] = (1, 2, 2), + enabled: bool = True, + mode_label: str = "Refine", +) -> RefinePlan: + """ + Validate the requested dimensions and create the stage-1 and target-resolution refinement plan. + + Parameters: + height (int): Target image height in pixels. + width (int): Target image width in pixels. + num_frames (int): Number of frames in the input sequence. + spatial_scale (int): Factor used to reduce spatial dimensions for stage 1. + vae_spatial_compression (int): Spatial compression factor of the VAE. + vae_temporal_compression (int): Temporal compression factor of the VAE. + patch_size (tuple[int, int, int]): Temporal and spatial patch dimensions used to validate latent-grid alignment. + enabled (bool): Whether to use two-pass refinement. + mode_label (str): Name of the calling mode, used to prefix validation + errors so ``--fast-spatial`` failures do not read as refine failures. + + Returns: + RefinePlan: The validated stage-1 and target-resolution plan. + """ + if height <= 0 or width <= 0: + raise ValueError(f"height/width must be positive, got {height}x{width}") + if spatial_scale < 1: + raise ValueError(f"spatial_scale must be >= 1, got {spatial_scale}") + if num_frames <= 0: + raise ValueError(f"num_frames must be positive, got {num_frames}") + if vae_spatial_compression < 1 or vae_temporal_compression < 1: + raise ValueError("VAE compression factors must be positive") + if height % vae_spatial_compression != 0 or width % vae_spatial_compression != 0: + raise ValueError(f"height/width must be divisible by vae_spatial_compression={vae_spatial_compression} " + f"(got {height}x{width}).") + if (num_frames - 1) % vae_temporal_compression != 0: + raise ValueError(f"num_frames must be 1 modulo vae_temporal_compression={vae_temporal_compression} " + f"(got {num_frames}).") + + if not enabled or spatial_scale == 1: + plan = RefinePlan( + target_height=height, + target_width=width, + stage1_height=height, + stage1_width=width, + spatial_scale=1, + vae_spatial_compression=vae_spatial_compression, + vae_temporal_compression=vae_temporal_compression, + num_frames=num_frames, + ) + _validate_plan(plan, patch_size=patch_size, mode_label=mode_label) + return plan + + if height % spatial_scale != 0 or width % spatial_scale != 0: + raise ValueError(f"{mode_label} requires height/width divisible by spatial_scale={spatial_scale} " + f"(got {height}x{width}).") + + stage1_height = height // spatial_scale + stage1_width = width // spatial_scale + # Stage-1 must land on a VAE-aligned grid so the first denoise produces + # valid latents; the LTX-2 init stage enforces the same constraint. + if (stage1_height % vae_spatial_compression != 0 or stage1_width % vae_spatial_compression != 0): + raise ValueError(f"{mode_label} requires height/width divisible by " + f"{spatial_scale * vae_spatial_compression} " + f"(got {height}x{width}, vae_spatial={vae_spatial_compression}).") + + plan = RefinePlan( + target_height=height, + target_width=width, + stage1_height=stage1_height, + stage1_width=stage1_width, + spatial_scale=spatial_scale, + vae_spatial_compression=vae_spatial_compression, + vae_temporal_compression=vae_temporal_compression, + num_frames=num_frames, + ) + _validate_plan(plan, patch_size=patch_size, mode_label=mode_label) + logger.info( + "[MLX refine] enabled: stage1=%dx%d stage2=%dx%d scale=%dx", + stage1_width, + stage1_height, + width, + height, + spatial_scale, + ) + return plan + + +def _validate_plan(plan: RefinePlan, *, patch_size: tuple[int, int, int], mode_label: str = "Refine") -> None: + """ + Validate that both refinement stages have latent dimensions aligned to the patch grid. + + Parameters: + patch_size (tuple[int, int, int]): Temporal, height, and width patch dimensions. + + Raises: + ValueError: If a stage's spatial latent dimensions or the temporal latent + dimension is not divisible by the corresponding patch dimension. + """ + pt, ph, pw = patch_size + for label, lh, lw in ( + ("stage1", plan.stage1_latent_height, plan.stage1_latent_width), + ("stage2", plan.stage2_latent_height, plan.stage2_latent_width), + ): + if lh % ph != 0 or lw % pw != 0: + raise ValueError(f"{mode_label} {label} latent grid {lh}x{lw} is not divisible by " + f"patch spatial size {ph}x{pw}.") + if plan.latent_frames % pt != 0: + raise ValueError(f"{mode_label} latent_frames={plan.latent_frames} is not divisible by " + f"patch temporal size {pt}.") + + +def upsample_latents_spatial( + latents: Any, + *, + scale: int = 2, + mode: str = "bilinear", +) -> Any: + """ + Upsample the spatial dimensions of 5-D latent arrays while preserving the batch, channel, and temporal dimensions. + + Parameters: + latents (Any): Latents with shape ``(B, C, T, H, W)``. + scale (int): Integer factor for enlarging the spatial dimensions. + mode (str): Interpolation mode, either ``"nearest"`` or ``"bilinear"``. + + Returns: + Any: Latents with shape ``(B, C, T, H * scale, W * scale)``. + """ + if scale < 1: + raise ValueError(f"scale must be >= 1, got {scale}") + if scale == 1: + return latents + + # Accept both mx.array and np.ndarray so unit tests can run without MLX. + is_mlx = hasattr(latents, "dtype") and type(latents).__module__.startswith("mlx") + if is_mlx: + return _upsample_latents_mlx(latents, scale=scale, mode=mode) + return _upsample_latents_numpy(np.asarray(latents), scale=scale, mode=mode) + + +def _upsample_latents_numpy( + latents: np.ndarray, + *, + scale: int, + mode: str, +) -> np.ndarray: + """Upsample 5-D latent arrays spatially using nearest-neighbor or bilinear interpolation. + + Parameters: + latents (np.ndarray): Latents with shape ``(B, C, T, H, W)``. + scale (int): Spatial upsampling factor. + mode (str): Interpolation mode, either ``"nearest"`` or ``"bilinear"``. + + Returns: + np.ndarray: Spatially upsampled latents with preserved batch, channel, and temporal dimensions. + + Raises: + ValueError: If the latents are not five-dimensional or the interpolation mode is unsupported. + """ + if latents.ndim != 5: + raise ValueError(f"Expected 5-D latents (B,C,T,H,W), got shape {latents.shape}") + b, c, t, h, w = latents.shape + if mode == "nearest": + # (B,C,T,H,1,W,1) -> broadcast to (B,C,T,H,scale,W,scale) -> merge. + out = np.repeat(np.repeat(latents, scale, axis=3), scale, axis=4) + return out + + if mode != "bilinear": + raise ValueError(f"Unsupported upsample mode: {mode}") + + # Bilinear over the spatial plane. Flatten (B,C,T) into a batch of 2-D + # maps so a single vectorized gather covers every frame/channel. + src = latents.reshape(b * c * t, h, w).astype(np.float32, copy=False) + out_h, out_w = h * scale, w * scale + # Map output pixel centers onto the input grid (align_corners=False). + ys = (np.arange(out_h, dtype=np.float32) + 0.5) * (h / out_h) - 0.5 + xs = (np.arange(out_w, dtype=np.float32) + 0.5) * (w / out_w) - 0.5 + ys = np.clip(ys, 0.0, h - 1.0) + xs = np.clip(xs, 0.0, w - 1.0) + y0 = np.floor(ys).astype(np.int64) + x0 = np.floor(xs).astype(np.int64) + y1 = np.minimum(y0 + 1, h - 1) + x1 = np.minimum(x0 + 1, w - 1) + wy = (ys - y0.astype(np.float32))[:, None] + wx = (xs - x0.astype(np.float32))[None, :] + # Gather the four corners: shape (N, out_h, out_w). + Ia = src[:, y0[:, None], x0[None, :]] + Ib = src[:, y0[:, None], x1[None, :]] + Ic = src[:, y1[:, None], x0[None, :]] + Id = src[:, y1[:, None], x1[None, :]] + wa = (1.0 - wy) * (1.0 - wx) + wb = (1.0 - wy) * wx + wc = wy * (1.0 - wx) + wd = wy * wx + out = wa * Ia + wb * Ib + wc * Ic + wd * Id + return out.reshape(b, c, t, out_h, out_w).astype(latents.dtype, copy=False) + + +def _upsample_latents_mlx(latents: mx.array, *, scale: int, mode: str) -> mx.array: + """ + Upsample MLX latent tensors along their spatial dimensions. + + Parameters: + latents (mx.array): A latent tensor with shape `(B, C, T, H, W)`. + scale (int): The integer spatial upsampling factor. + mode (str): The interpolation mode, such as `"nearest"` or `"bilinear"`. + + Returns: + mx.array: The spatially upsampled latent tensor with its original data type. + """ + import mlx.core as mx + + # Route through NumPy for the interpolation math. Latent tensors at Mac + # resolutions are small (e.g. 1×16×21×30×52 ≈ 1 MB) so the host hop is + # cheaper than carrying a bespoke Metal bilinear kernel, and it keeps + # the CPU-only unit tests and the MLX path on one implementation. + np_latents = np.array(latents.astype(mx.float32)) + up = _upsample_latents_numpy(np_latents, scale=scale, mode=mode) + return mx.array(up).astype(latents.dtype) + + +def prepare_refine_latents( + clean_latents: Any, + *, + scale: int = 2, + sigma: float = DEFAULT_REFINE_SIGMA, + noise: Any | None = None, + add_noise_flag: bool = True, + upsample_mode: str = "bilinear", + seed: int | None = None, +) -> Any: + """ + Upsample clean latents spatially and optionally mix them with Gaussian noise. + + Parameters: + clean_latents: The stage-1 latent tensor. + sigma: Noise mixing factor between 0 and 1. + noise: Optional noise tensor to mix with the upsampled latents. + add_noise_flag: Whether to apply noise mixing. + upsample_mode: Spatial interpolation mode. + seed: Optional seed for generated noise. + + Returns: + The upsampled latents, optionally mixed with noise. + + Raises: + ValueError: If sigma is outside the range from 0 to 1. + """ + if sigma < 0.0 or sigma > 1.0: + raise ValueError(f"sigma must be in [0, 1], got {sigma}") + + upsampled = upsample_latents_spatial(clean_latents, scale=scale, mode=upsample_mode) + if not add_noise_flag or sigma == 0.0: + return upsampled + + is_mlx = hasattr(upsampled, "dtype") and type(upsampled).__module__.startswith("mlx") + if noise is None: + noise = _draw_noise_like(upsampled, seed=seed, is_mlx=is_mlx) + return add_noise(upsampled, noise, float(sigma)) + + +def refine_sigma_from_schedule( + schedule: MLXDMDSchedule, + timesteps: Sequence[float | int], +) -> float: + """Derive the refinement noise level from the first refinement timestep. + + Parameters: + schedule (MLXDMDSchedule): Schedule used to map timesteps to noise levels. + timesteps (Sequence[float | int]): Refinement timesteps, whose first value determines the sigma. + + Returns: + float: Sigma corresponding to the first refinement timestep. + + Raises: + ValueError: If `timesteps` is empty. + """ + if not timesteps: + raise ValueError("timesteps must be non-empty to derive a refine sigma") + return float(schedule.sigma_for(float(timesteps[0]))) + + +def default_refine_timesteps( + schedule: MLXDMDSchedule, + timesteps: Sequence[float | int], +) -> list[float]: + """Derive stage-2 timesteps from the stage-1 DMD grid. + + The stage-2 pass must start *below* full noise, otherwise the hand-off + ``(1 - sigma) * upsampled + sigma * noise`` weights stage 1 at zero and + the refine pass silently becomes a plain full-resolution generation at + twice the cost. FastWan's stage-1 grid opens at ``t=1000`` (``sigma`` + exactly 1.0), so reusing it verbatim — which is what happens when + ``--refine-dmd-denoising-steps`` is left unset — discards stage 1. + + Dropping the leading full-noise entries keeps the pass on timesteps the + distilled student was actually trained on (no off-grid ``t`` the DiT has + never seen) while letting the stage-1 structure through. + + Parameters: + schedule (MLXDMDSchedule): Schedule used to map timesteps to noise levels. + timesteps (Sequence[float | int]): The stage-1 DMD timestep grid. + + Returns: + list[float]: The stage-1 grid with leading full-noise timesteps removed. + + Raises: + ValueError: If every timestep in the grid is at full noise, leaving no + usable refine step. + """ + steps = [float(step) for step in timesteps] + first = 0 + while first < len(steps) and schedule.sigma_for(steps[first]) >= 1.0: + first += 1 + if first == len(steps): + raise ValueError(f"No usable refine timesteps in {steps}: every entry is at sigma >= 1 " + "(full noise), which would discard the stage-1 result. Pass " + "explicit stage-2 timesteps below the full-noise step.") + return steps[first:] + + +def run_dmd_loop( + *, + dit: Any, + latents: Any, + encoder_hidden_states: Any, + freqs_cis: tuple[Any, Any], + timesteps: Sequence[float | int], + schedule: MLXDMDSchedule, + mx_dtype: Any, + seed: int | None = None, + step_callback: Callable[[int, int], None] | None = None, + label: str = "denoise", +) -> Any: + """ + Denoise latents over the supplied timesteps using the DMD schedule. + + Parameters: + timesteps (Sequence[float | int]): Denoising timesteps in execution order. + seed (int | None): Seed for reproducible intermediate noise generation. + step_callback (Callable[[int, int], None] | None): Callback receiving the + completed step number and total step count. + label (str): Label used for progress output when no callback is provided. + + Returns: + Any: The denoised latents. + """ + import mlx.core as mx + + renoise_rng = np.random.default_rng(seed) if seed is not None else None + latents_out = latents + n_steps = len(timesteps) + for step_index, timestep in enumerate(timesteps): + noise_input = latents_out + ts_val = float(timestep) + timestep_mx = mx.array([ts_val]).astype(mx.float32) + noise_pred = dit( + latents_out.astype(mx_dtype), + encoder_hidden_states, + timestep_mx, + freqs_cis, + ) + noise_input_f32 = noise_input.astype(mx.float32) + pred_noise_f32 = noise_pred.astype(mx.float32) + if step_index < n_steps - 1: + next_ts: float | None = float(timesteps[step_index + 1]) + if renoise_rng is not None: + renoise = mx.array(renoise_rng.standard_normal(tuple(noise_input_f32.shape)).astype(np.float32)) + else: + renoise = mx.random.normal(noise_input_f32.shape).astype(mx.float32) + else: + next_ts, renoise = None, None + latents_out = dmd_step( + latents=noise_input_f32, + noise_input_latent=noise_input_f32, + pred_noise=pred_noise_f32, + schedule=schedule, + timestep=ts_val, + next_timestep=next_ts, + noise=renoise, + ).astype(mx_dtype) + mx.eval(latents_out) + if step_callback is not None: + step_callback(step_index + 1, n_steps) + else: + print(f"{label} step {step_index + 1}/{n_steps} complete") + return latents_out + + +@dataclass(frozen=True) +class TwoPassResult: + """Outputs of :func:`run_two_pass_dmd`.""" + + latents: Any + stage1_latents: Any + plan: RefinePlan + refine_sigma: float + + +def run_two_pass_dmd( + *, + dit: Any, + encoder_hidden_states: Any, + noise_latents_stage1: Any, + freqs_cis_stage1: tuple[Any, Any], + freqs_cis_stage2: tuple[Any, Any] | None, + plan: RefinePlan, + schedule: MLXDMDSchedule, + timesteps: Sequence[float | int], + refine_timesteps: Sequence[float | int] | None = None, + mx_dtype: Any, + seed: int = 0, + add_noise_flag: bool = True, + upsample_mode: str = "bilinear", + refine_sigma: float | None = None, + step_callback: Callable[[str, int, int], None] | None = None, +) -> TwoPassResult: + """ + Run base denoising and, when enabled, spatial refinement denoising. + + Parameters: + dit: DiT callable used for both denoising passes. + encoder_hidden_states: Prompt embeddings shared across both passes. + noise_latents_stage1: Initial stage-1 noise latents. + freqs_cis_stage1: RoPE tables for the stage-1 resolution. + freqs_cis_stage2: RoPE tables for the stage-2 resolution, required when refinement is enabled. + plan: Refinement geometry and configuration. + schedule: Flow-matching schedule used by both passes. + timesteps: Stage-1 denoising timesteps. + refine_timesteps: Stage-2 denoising timesteps. Uses `timesteps` when omitted. + mx_dtype: MLX dtype used for DiT inputs and outputs. + seed: Base seed for reproducible noise generation. + add_noise_flag: Whether to add noise to the upsampled stage-1 latents. + upsample_mode: Spatial upsampling mode, either `"bilinear"` or `"nearest"`. + refine_sigma: Stage-2 starting noise level. Derived from the first refinement timestep when omitted. + step_callback: Optional callback receiving the phase name, step index, and total step count. + + Returns: + TwoPassResult containing the final latents, stage-1 latents, refinement plan, and applied refinement sigma. + + Raises: + ValueError: If refinement is enabled without stage-2 RoPE tables, without refinement timesteps, or if upsampled latents do not match the planned stage-2 dimensions. + """ + stage1_cb = None + stage2_cb = None + if step_callback is not None: + stage1_cb = lambda i, n: step_callback("stage1", i, n) # noqa: E731 + stage2_cb = lambda i, n: step_callback("stage2", i, n) # noqa: E731 + + stage1_latents = run_dmd_loop( + dit=dit, + latents=noise_latents_stage1, + encoder_hidden_states=encoder_hidden_states, + freqs_cis=freqs_cis_stage1, + timesteps=timesteps, + schedule=schedule, + mx_dtype=mx_dtype, + seed=seed, + step_callback=stage1_cb, + label="stage1 denoise", + ) + + if plan.spatial_scale == 1: + return TwoPassResult( + latents=stage1_latents, + stage1_latents=stage1_latents, + plan=plan, + refine_sigma=0.0, + ) + + if freqs_cis_stage2 is None: + raise ValueError("freqs_cis_stage2 is required when refine spatial_scale > 1") + + if refine_timesteps is not None: + stage2_timesteps = [float(step) for step in refine_timesteps] + if not stage2_timesteps: + raise ValueError("refine_timesteps must be non-empty when refine is enabled") + else: + # Not `list(timesteps)`: the stage-1 grid opens at full noise, which + # would weight the stage-1 result at zero. See default_refine_timesteps. + stage2_timesteps = default_refine_timesteps(schedule, timesteps) + grid_sigma = refine_sigma_from_schedule(schedule, stage2_timesteps) + sigma = float(refine_sigma) if refine_sigma is not None else grid_sigma + if refine_sigma is not None and abs(sigma - grid_sigma) > 1e-6: + # The loop tells the DiT `stage2_timesteps[0]`, which implies grid_sigma. + # Overriding the hand-off noise level breaks that correspondence, so the + # model is denoising from a level it was not told about. Useful for + # exploring schedules that bottom out too high, but say so out loud. + logger.warning( + "[MLX refine] refine_sigma=%.4f overrides the schedule's %.4f for timestep %g; " + "the DiT is told a timestep that no longer matches the noise it receives.", + sigma, + grid_sigma, + stage2_timesteps[0], + ) + + # A hand-off at sigma >= 1 is `0 * upsampled + 1 * noise`: stage 1 is + # thrown away and refine degrades to a plain full-res run at 2x the cost. + # Fail loudly rather than silently burning the first pass. + if add_noise_flag and sigma >= 1.0: + raise ValueError(f"Refine hand-off sigma={sigma:.4f} (from stage-2 timestep " + f"{stage2_timesteps[0]:g}) discards the stage-1 result entirely: " + "the upsampled latents are weighted (1 - sigma) = 0. Start the " + "stage-2 grid below the full-noise timestep, or pass " + "add_noise_flag=False to hand off the clean upsample.") + + stage2_input = prepare_refine_latents( + stage1_latents, + scale=plan.spatial_scale, + sigma=sigma, + add_noise_flag=add_noise_flag, + upsample_mode=upsample_mode, + seed=seed + 1, + ) + + # Shape guard: upsampled latents must match the stage-2 RoPE grid. + expected_h = plan.stage2_latent_height + expected_w = plan.stage2_latent_width + got_h, got_w = int(stage2_input.shape[-2]), int(stage2_input.shape[-1]) + if got_h != expected_h or got_w != expected_w: + raise ValueError(f"Refine upsample produced {got_h}x{got_w} latents, expected " + f"{expected_h}x{expected_w} for target " + f"{plan.target_height}x{plan.target_width}.") + + logger.info( + "[MLX refine] stage2 start: latent=%dx%d sigma=%.4f steps=%d", + expected_w, + expected_h, + sigma, + len(stage2_timesteps), + ) + + stage2_latents = run_dmd_loop( + dit=dit, + latents=stage2_input, + encoder_hidden_states=encoder_hidden_states, + freqs_cis=freqs_cis_stage2, + timesteps=stage2_timesteps, + schedule=schedule, + mx_dtype=mx_dtype, + seed=seed + 2, + step_callback=stage2_cb, + label="stage2 refine", + ) + return TwoPassResult( + latents=stage2_latents, + stage1_latents=stage1_latents, + plan=plan, + refine_sigma=sigma, + ) + + +__all__ = [ + "DEFAULT_REFINE_SIGMA", + "RefinePlan", + "TwoPassResult", + "default_refine_timesteps", + "plan_refine_resolutions", + "prepare_refine_latents", + "refine_sigma_from_schedule", + "run_dmd_loop", + "run_two_pass_dmd", + "upsample_latents_spatial", +] + + +def _draw_noise_like(like: Any, *, seed: int | None, is_mlx: bool) -> Any: + """Generate Gaussian noise with the shape and array type of the input.""" + shape = tuple(int(s) for s in like.shape) + if seed is not None: + rng = np.random.default_rng(seed) + noise_np = rng.standard_normal(shape).astype(np.float32) + if is_mlx: + import mlx.core as mx + + return mx.array(noise_np).astype(mx.float32) + return noise_np.astype(np.asarray(like).dtype, copy=False) + if is_mlx: + import mlx.core as mx + + return mx.random.normal(shape).astype(mx.float32) + return np.random.standard_normal(shape).astype(np.asarray(like).dtype, copy=False) diff --git a/fastvideo/mlx_runtime/rife_interp.py b/fastvideo/mlx_runtime/rife_interp.py new file mode 100644 index 000000000..f3cbff020 --- /dev/null +++ b/fastvideo/mlx_runtime/rife_interp.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Small MLX RIFE wrapper for frame interpolation experiments. + +The backend is the Apple-Silicon-native ``rife-mlx`` package, using the +``mlx-community/RIFE-4.25`` weights. Frames are HWC RGB ``uint8`` arrays. +""" + +from __future__ import annotations + +from collections.abc import Iterable +from functools import lru_cache + +import numpy as np +from huggingface_hub.utils import LocalEntryNotFoundError + + +class RIFEBackendError(RuntimeError): + """Raised when the MLX RIFE backend cannot be loaded or run.""" + + +class RIFEWeightsUnavailableError(RIFEBackendError): + """Raised when uncached RIFE weights cannot be downloaded.""" + + +def aligned_keyframe_count(target_frames: int, factor: int, temporal_compression: int = 4) -> int: + """Return the smallest VAE-aligned keyframe count that RIFE can expand to the target.""" + if target_frames < 1: + raise ValueError(f"target_frames must be >= 1, got {target_frames}") + if factor < 1: + raise ValueError(f"factor must be >= 1, got {factor}") + if temporal_compression < 1: + raise ValueError(f"temporal_compression must be >= 1, got {temporal_compression}") + required_intervals = (target_frames - 1 + factor - 1) // factor + aligned_intervals = ((required_intervals + temporal_compression - 1) // temporal_compression * temporal_compression) + return aligned_intervals + 1 + + +def _require_hwc_rgb(frame: np.ndarray, index: int) -> np.ndarray: + array = np.asarray(frame) + if array.ndim != 3 or array.shape[2] != 3: + raise ValueError(f"frame {index} must have shape HxWx3, got {array.shape}") + if array.dtype != np.uint8: + array = np.clip(array, 0, 255).astype(np.uint8) + return np.ascontiguousarray(array) + + +@lru_cache(maxsize=2) +def load_model(version: str = "4.25", weights_dir: str | None = None): + """Load the MLX-native RIFE model. + + ``weights_dir`` is passed through to ``build_model`` in the vendored ``rife_mlx``. + When it is ``None``, the package downloads/uses the Hugging Face + ``mlx-community/RIFE-4.25`` snapshot. + """ + try: + from fastvideo.third_party.rife_mlx.utils.weights import build_model + except ImportError: + # Fall back to a separately installed upstream package, for anyone who + # already has one in the environment. + try: + from rife_mlx.utils.weights import build_model + except ImportError as exc: + raise RIFEBackendError("MLX RIFE backend is unavailable. It ships vendored under " + "fastvideo/third_party/rife_mlx, so this usually means MLX " + "itself is missing: install with `uv pip install -e '.[mlx]'`.") from exc + + try: + return build_model(version, weights_dir=weights_dir) + except LocalEntryNotFoundError as exc: + raise RIFEWeightsUnavailableError(f"MLX RIFE {version} weights are unavailable: {exc}") from exc + except Exception as exc: # noqa: BLE001 - preserve exact backend failure. + raise RIFEBackendError(f"Failed to load MLX RIFE {version}: {exc}") from exc + + +def interpolate_pair( + frame_a: np.ndarray, + frame_b: np.ndarray, + timestep: float = 0.5, + *, + model=None, + scale: float = 1.0, +) -> np.ndarray: + """Interpolate one RGB frame between two input RGB frames.""" + if not 0.0 < timestep < 1.0: + raise ValueError(f"timestep must be inside (0, 1), got {timestep}") + img0 = _require_hwc_rgb(frame_a, 0) + img1 = _require_hwc_rgb(frame_b, 1) + if img0.shape != img1.shape: + raise ValueError(f"frame shapes must match, got {img0.shape} and {img1.shape}") + + if model is None: + model = load_model() + try: + try: + from fastvideo.third_party.rife_mlx.pipeline_mlx import interpolate_pair as _interpolate_pair + except ImportError: + from rife_mlx.pipeline_mlx import interpolate_pair as _interpolate_pair + + return _interpolate_pair(model, img0, img1, timestep=timestep, scale=scale) + except Exception as exc: # noqa: BLE001 - preserve exact backend failure. + raise RIFEBackendError(f"MLX RIFE interpolation failed at timestep={timestep}: {exc}") from exc + + +def interpolate( + frames: list[np.ndarray] | Iterable[np.ndarray], + factor: int = 2, + *, + model=None, + scale: float = 1.0, +) -> list[np.ndarray]: + """Return an Nx interpolated frame list. + + For ``len(frames)=41`` and ``factor=2``, the output length is 81: + ``(41 - 1) * 2 + 1``. Original keyframes are preserved in order and RIFE + fills ``factor - 1`` intermediate timesteps between each adjacent pair. + """ + frame_list = [_require_hwc_rgb(frame, idx) for idx, frame in enumerate(frames)] + if factor < 1: + raise ValueError(f"factor must be >= 1, got {factor}") + if len(frame_list) < 2 or factor == 1: + return [frame.copy() for frame in frame_list] + + first_shape = frame_list[0].shape + for idx, frame in enumerate(frame_list[1:], start=1): + if frame.shape != first_shape: + raise ValueError(f"all frames must have the same shape; frame 0={first_shape}, frame {idx}={frame.shape}") + + if model is None: + model = load_model() + + out: list[np.ndarray] = [] + for left, right in zip(frame_list[:-1], frame_list[1:], strict=True): + out.append(left) + for step in range(1, factor): + out.append(interpolate_pair(left, right, step / factor, model=model, scale=scale)) + out.append(frame_list[-1]) + return out diff --git a/fastvideo/mlx_runtime/sampling.py b/fastvideo/mlx_runtime/sampling.py new file mode 100644 index 000000000..2efecabea --- /dev/null +++ b/fastvideo/mlx_runtime/sampling.py @@ -0,0 +1,139 @@ +# SPDX-License-Identifier: Apache-2.0 +"""On-device (MLX) DMD sampling for the FastWan runtime. + +The hybrid proof-of-concept ran the FastWan DiT in MLX but bounced every +denoising step back through torch/NumPy to run the DMD scheduler math +(``MLX -> np.array -> torch (CPU) -> np.array -> MLX``). That host round-trip +forces a full device sync per step and defeats MLX's lazy graph execution. + +This module mirrors the exact DMD arithmetic from +``fastvideo/models/utils.py::pred_noise_to_pred_video`` and +``FlowMatchEulerDiscreteScheduler.add_noise`` while keeping every large tensor +on the MLX device. The schedule lookup (``argmin`` over the ~1000-entry +training schedule) is done once on the host in NumPy: it is tiny, it is the +same value torch would compute, and it sidesteps the reduction-index quirk that +affects ``argmin`` on the Metal/MPS backends (see the CPU fallbacks in +``fastvideo/models/utils.py`` and ``scheduling_flow_match_euler_discrete.py``). + +Because the DMD loop applies a single scalar timestep per step, ``sigma`` is a +scalar and the update is a plain elementwise affine combination — no +permute/flatten reshaping is required. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING, Any + +import numpy as np + +if TYPE_CHECKING: # pragma: no cover - typing only + import mlx.core as mx + + +@dataclass(frozen=True) +class MLXDMDSchedule: + """Host-side copy of a flow-match scheduler's ``(sigmas, timesteps)``. + + Holds the full training schedule so a DMD timestep (e.g. one of + ``1000, 757, 522``) can be mapped to its flow-match ``sigma`` with the same + nearest-timestep lookup the torch path uses. + """ + + sigmas: np.ndarray + timesteps: np.ndarray + + @classmethod + def from_torch_scheduler(cls, scheduler: Any) -> MLXDMDSchedule: + """Snapshot ``scheduler.sigmas`` / ``scheduler.timesteps`` to NumPy. + + Matches ``pred_noise_to_pred_video`` / ``add_noise``, which index the + scheduler's *full* training schedule (not the per-inference subset). + """ + sigmas = scheduler.sigmas.detach().to("cpu").double().numpy() + timesteps = scheduler.timesteps.detach().to("cpu").double().numpy() + return cls(sigmas=np.asarray(sigmas), timesteps=np.asarray(timesteps)) + + def sigma_for(self, timestep: float) -> float: + """ + Find the sigma associated with the scheduled timestep nearest to the given timestep. + + Parameters: + timestep (float): Timestep for which to find the nearest scheduled sigma. + + Returns: + float: Sigma associated with the nearest scheduled timestep. + """ + idx = int(np.argmin(np.abs(self.timesteps - float(timestep)))) + return float(self.sigmas[idx]) + + +def pred_noise_to_pred_video( + pred_noise: mx.array, + noise_input_latent: mx.array, + sigma: float, +) -> mx.array: + """ + Compute the clean latent prediction from a flow-matching noise prediction. + + Parameters: + pred_noise (mx.array): Predicted noise. + noise_input_latent (mx.array): Noised latent input. + sigma (float): Noise level used for the prediction. + + Returns: + mx.array: Predicted clean latent. + """ + return noise_input_latent - sigma * pred_noise + + +def add_noise( + clean_latent: mx.array, + noise: mx.array, + sigma: float, +) -> mx.array: + """Flow-match forward noising, mirroring the scheduler's ``add_noise``. + + ``sample = (1 - sigma) * clean_latent + sigma * noise``. + """ + return (1.0 - sigma) * clean_latent + sigma * noise + + +def dmd_step( + *, + latents: mx.array, + noise_input_latent: mx.array, + pred_noise: mx.array, + schedule: MLXDMDSchedule, + timestep: float, + next_timestep: float | None, + noise: mx.array | None = None, +) -> mx.array: + """ + Compute one DMD sampling update, optionally re-noising the clean latent prediction. + + Args: + latents: Retained for call-site compatibility and not used in the update. + noise_input_latent: Noisy latent used to compute the clean prediction. + pred_noise: Predicted noise or velocity. + schedule: Flow-matching schedule used to map timesteps to sigmas. + timestep: Current sampling timestep. + next_timestep: Timestep for the next update, or `None` for the final step. + noise: Fresh noise used for re-noising intermediate steps. + + Returns: + The re-noised latent for the next step or the clean latent prediction on + the final step. + + Raises: + ValueError: If `next_timestep` is provided without `noise`. + """ + del latents # symmetry with the torch loop; not needed for the math. + sigma = schedule.sigma_for(timestep) + pred_video = pred_noise_to_pred_video(pred_noise, noise_input_latent, sigma) + if next_timestep is None: + return pred_video + if noise is None: + raise ValueError("dmd_step requires `noise` when `next_timestep` is set (re-noise step).") + sigma_next = schedule.sigma_for(next_timestep) + return add_noise(pred_video, noise, sigma_next) diff --git a/fastvideo/mlx_runtime/taehv_decode.py b/fastvideo/mlx_runtime/taehv_decode.py new file mode 100644 index 000000000..6d68e1a30 --- /dev/null +++ b/fastvideo/mlx_runtime/taehv_decode.py @@ -0,0 +1,191 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Optional TAEHV decode helpers for Apple Silicon FastWan experiments. + +The TAEHV module itself is vendored at ``fastvideo/third_party/taehv`` (MIT, +madebyollin/taehv), so no source code is downloaded or executed at runtime. +Only the ``taew2_1.pth`` checkpoint is fetched on demand, and its sha256 is +verified before use. +""" + +from __future__ import annotations + +import hashlib +import importlib.util +import urllib.request +from pathlib import Path + +import numpy as np + +TAEW2_1_CHECKPOINT_URL = "https://raw.githubusercontent.com/madebyollin/taehv/main/taew2_1.pth" +# sha256 of the upstream taew2_1.pth this module was validated against +# (fetched 2026-07-02). If upstream publishes a new checkpoint, revalidate the +# decode path and update this pin. +TAEW2_1_CHECKPOINT_SHA256 = "d26151e76cdc2c9424bef988de874b33d9a53f30ef3060cd556c429c469c797e" +# Wan2.2 5B (z_dim=48) — see madebyollin/taehv taew2_2.pth; prefer +# ``fastvideo.mlx_runtime.wan_vae.ensure_taehv_checkpoint(z_dim=48)`` for new code. +TAEW2_2_CHECKPOINT_URL = ("https://raw.githubusercontent.com/madebyollin/taehv/" + "563f40bdc820ed86bcad72ea515ee48f06bd22ec/taew2_2.pth") + + +def _default_cache_dir() -> Path: + """Return the default directory used to cache TAEHV checkpoints. + + Returns: + Path: The TAEHV checkpoint cache directory under the user's home directory. + """ + return Path.home() / ".cache" / "fastvideo" / "taehv" + + +def _sha256(path: Path) -> str: + """ + Compute the SHA-256 digest of a file. + + Parameters: + path (Path): The file whose contents are hashed. + + Returns: + str: The file's SHA-256 digest in hexadecimal form. + """ + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _verify_checkpoint(path: Path) -> None: + """Verify that a TAEW2.1 checkpoint matches the expected SHA-256 digest. + + Parameters: + path (Path): Path to the checkpoint file. + + Raises: + RuntimeError: If the checkpoint digest does not match the expected value. + """ + actual = _sha256(path) + if actual != TAEW2_1_CHECKPOINT_SHA256: + raise RuntimeError(f"TAEHV checkpoint at {path} failed sha256 verification " + f"(expected {TAEW2_1_CHECKPOINT_SHA256}, got {actual}). " + "Delete the file to re-download it, or pass --taehv-checkpoint-path " + "pointing at a checkpoint you trust.") + + +def ensure_taew2_1_checkpoint(checkpoint_path: Path | None = None) -> Path: + """ + Ensure the TAEW2.1 checkpoint is available locally. + + A caller-provided path is treated as trusted and is only checked for existence. + The module-managed cached checkpoint is verified against the pinned SHA-256 digest + after downloading or before reuse. + + Parameters: + checkpoint_path (Path | None): Optional path to a caller-provided checkpoint. + + Returns: + Path: The available checkpoint path. + + Raises: + FileNotFoundError: If a caller-provided checkpoint does not exist. + RuntimeError: If a module-managed checkpoint fails verification. + """ + if checkpoint_path is not None: + if not checkpoint_path.exists(): + raise FileNotFoundError(f"TAEHV checkpoint not found: {checkpoint_path}") + return checkpoint_path + + checkpoint_path = _default_cache_dir() / "taew2_1.pth" + if not checkpoint_path.exists(): + checkpoint_path.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading {TAEW2_1_CHECKPOINT_URL} -> {checkpoint_path}") + import socket + import tempfile + # Download to a temporary file, verify, then atomically rename. + with tempfile.NamedTemporaryFile( + mode="wb", + dir=checkpoint_path.parent, + prefix=".tmp_taew2_1_", + suffix=".pth", + delete=False, + ) as tmp_file: + tmp_path = Path(tmp_file.name) + try: + old_timeout = socket.getdefaulttimeout() + socket.setdefaulttimeout(300) + try: + urllib.request.urlretrieve( + TAEW2_1_CHECKPOINT_URL, + tmp_path, # noqa: S310 - pinned public artifact, hash-verified below. + ) + finally: + socket.setdefaulttimeout(old_timeout) + _verify_checkpoint(tmp_path) + tmp_path.replace(checkpoint_path) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + else: + _verify_checkpoint(checkpoint_path) + return checkpoint_path + + +def _load_taehv_class(source_path: Path | None): + """Load the TAEHV class from the vendored implementation or a local source override. + + Parameters: + source_path (Path | None): Path to a local Python file defining `TAEHV`; `None` selects the vendored implementation. + + Returns: + The loaded `TAEHV` class. + + Raises: + RuntimeError: If the specified source cannot be loaded. + """ + if source_path is None: + from fastvideo.third_party.taehv import TAEHV + + return TAEHV + # Explicit local override for experimenting with a modified TAEHV; this is + # a user-supplied file on disk, never something this module downloads. + spec = importlib.util.spec_from_file_location("fastvideo_external_taehv", source_path) + if spec is None or spec.loader is None: + raise RuntimeError(f"Could not load TAEHV source from {source_path}") + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module.TAEHV + + +def decode_latents_to_video_taehv( + *, + latents_np: np.ndarray, + output_path: Path, + fps: int, + device, + dtype, + parallel: bool, + source_path: Path | None = None, + checkpoint_path: Path | None = None, +) -> None: + """Decode Wan/FastWan diffusion latents with TAEW2.1 and export MP4. + + TAEHV's Wan wrapper expects the diffusion latents directly, without applying + the standard Wan VAE's `latents_mean` / `latents_std` shift. + """ + import torch + from diffusers.utils import export_to_video + + checkpoint_path = ensure_taew2_1_checkpoint(checkpoint_path) + TAEHV = _load_taehv_class(source_path) + taehv = TAEHV(str(checkpoint_path)).to(device=device, dtype=dtype) + taehv.eval() + + latents = torch.from_numpy(latents_np).to(device=device, dtype=dtype) + with torch.no_grad(): + video_ntchw = taehv.decode_video( + latents.transpose(1, 2), + parallel=parallel, + show_progress_bar=False, + ) + video = video_ntchw.transpose(1, 2) + video_np = video[0].permute(1, 2, 3, 0).float().cpu().numpy() + output_path.parent.mkdir(parents=True, exist_ok=True) + export_to_video(video_np, str(output_path), fps=fps) diff --git a/fastvideo/mlx_runtime/wan22.py b/fastvideo/mlx_runtime/wan22.py new file mode 100644 index 000000000..6171481a6 --- /dev/null +++ b/fastvideo/mlx_runtime/wan22.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan2.2-TI2V-5B dense MLX runtime — Track D. + +The Wan2.2 TI2V-5B (FullAttn) differs from the ported Wan2.1-T2V only in: + +- **Scale** (24 heads x 128, hidden 3072, ffn 14336) — pure config, block math + identical, so the dense loader ``mlx_dit_from_diffusers_safetensors`` loads the + weights unchanged and we re-wrap the blocks here. +- **Per-token timestep conditioning** (``expand_timesteps=True``): the timestep is + ``[batch, seq_len]`` (a level per patch token — how TI2V keeps the conditioning + image frame at t=0 while the video frames are noised). ``timestep_proj`` becomes + ``[batch, seq_len, 6, dim]`` and the block/output modulation is per-token + (``[B, L, dim]``), a direct broadcast — this module implements exactly that. + +I2V rides on the same forward: encode the image, replace the first latent frame, +and set that frame's timestep to 0 (handled by the caller / sampler). See +``docs/design/ti2v_5b_port_guide.md``. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING, Any +from collections.abc import Callable + +from fastvideo.logger import init_logger +from fastvideo.mlx_runtime.fastwan import ( + MLXWanT2VCrossAttention, + gelu_tanh, + layer_norm, + linear, + mlx_dit_from_diffusers_safetensors, + rms_norm, + silu, + timestep_embedding, + weight_dtype, +) + +if TYPE_CHECKING: + import mlx.core as mx + +logger = init_logger(__name__) + + +class MLXWan22TransformerBlock: + """Dense Wan block with per-token (``[B, L, dim]``) timestep modulation.""" + + def __init__( + self, + weights: dict[str, mx.array], + *, + dim: int, + ffn_dim: int, + num_heads: int, + eps: float = 1e-6, + ): + self.weights = weights + self.dim = dim + self.ffn_dim = ffn_dim + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.eps = eps + self.attn2 = MLXWanT2VCrossAttention(weights, dim=dim, num_heads=num_heads, eps=eps) + + def __call__(self, hidden_states, encoder_hidden_states, timestep_proj, cos, sin) -> mx.array: + import mlx.core as mx + + orig_dtype = hidden_states.dtype + batch = hidden_states.shape[0] + + # timestep_proj: [B, L, 6, dim] -> six per-token [B, L, dim] modulations. + e = self.weights["scale_shift_table"][None] + timestep_proj.astype(mx.float32) + shift_msa, scale_msa, gate_msa, c_shift_msa, c_scale_msa, c_gate_msa = [ + part.squeeze(2) for part in mx.split(e, 6, axis=2) + ] + + # 1. Self-attention (dense, bidirectional) with per-token modulation. + norm_hidden = layer_norm(hidden_states.astype(mx.float32), eps=self.eps) + norm_hidden = (norm_hidden * (1.0 + scale_msa) + shift_msa).astype(orig_dtype) + + query = linear(norm_hidden, self.weights["to_q.weight"], self.weights.get("to_q.bias")) + key = linear(norm_hidden, self.weights["to_k.weight"], self.weights.get("to_k.bias")) + value = linear(norm_hidden, self.weights["to_v.weight"], self.weights.get("to_v.bias")) + + query = rms_norm(query, self.weights["norm_q.weight"], + eps=self.eps).reshape(batch, -1, self.num_heads, self.head_dim) + key = rms_norm(key, self.weights["norm_k.weight"], eps=self.eps).reshape(batch, -1, self.num_heads, + self.head_dim) + value = value.reshape(batch, -1, self.num_heads, self.head_dim) + + from fastvideo.mlx_runtime.fastwan import apply_rotary_emb + query = apply_rotary_emb(query, cos, sin, is_neox_style=False) + key = apply_rotary_emb(key, cos, sin, is_neox_style=False) + + attn = mx.fast.scaled_dot_product_attention( + query.transpose(0, 2, 1, 3), + key.transpose(0, 2, 1, 3), + value.transpose(0, 2, 1, 3), + scale=self.head_dim**-0.5, + ).transpose(0, 2, 1, 3) + attn = attn.reshape(batch, -1, self.dim) + attn = linear(attn, self.weights["to_out.weight"], self.weights.get("to_out.bias")) + + hidden_states = hidden_states + (attn * gate_msa).astype(orig_dtype) + norm_hidden = layer_norm(hidden_states.astype(mx.float32), + weight=self.weights["self_attn_residual_norm.norm.weight"], + bias=self.weights["self_attn_residual_norm.norm.bias"], + eps=self.eps).astype(orig_dtype) + + # 2. Cross-attention, then per-token shift/scale modulation. + cross = self.attn2(norm_hidden, encoder_hidden_states) + hidden_states = hidden_states + cross + norm_hidden = layer_norm(hidden_states.astype(mx.float32), eps=self.eps) + norm_hidden = (norm_hidden * (1.0 + c_scale_msa) + c_shift_msa).astype(orig_dtype) + + # 3. Feed-forward with per-token gate. + ff = linear(norm_hidden, self.weights["ffn.fc_in.weight"], self.weights.get("ffn.fc_in.bias")) + ff = gelu_tanh(ff) + ff = linear(ff, self.weights["ffn.fc_out.weight"], self.weights.get("ffn.fc_out.bias")) + hidden_states = hidden_states + (ff * c_gate_msa).astype(orig_dtype) + return hidden_states.astype(orig_dtype) + + +class MLXWan22DiT: + """Wan2.2-TI2V-5B dense DiT with per-token timestep conditioning.""" + + def __init__( + self, + weights: dict[str, mx.array], + blocks: list[MLXWan22TransformerBlock], + config: dict, + *, + compile: bool = False, + ) -> None: + import os + + self.weights = weights + self.blocks = blocks + self.config = config + self.num_heads = int(config["num_attention_heads"]) + self.head_dim = int(config["attention_head_dim"]) + self.hidden_size = self.num_heads * self.head_dim + self.freq_dim = int(config["freq_dim"]) + self.patch_size = tuple(config["patch_size"]) + self.out_channels = int(config["out_channels"]) + self.eps = float(config.get("eps", 1e-6)) + self._enable_compile = compile or os.environ.get("FASTVIDEO_MLX_COMPILE", "0") == "1" + self._compiled_forward: Callable[..., Any] | None = None + self._compiled_signature: tuple | None = None + + def _patch_embed(self, hidden_states) -> mx.array: + batch, channels, frames, height, width = hidden_states.shape + pt, ph, pw = self.patch_size + patch_dim = channels * pt * ph * pw + x = hidden_states.reshape(batch, channels, frames // pt, pt, height // ph, ph, width // pw, pw) + x = x.transpose(0, 2, 4, 6, 1, 3, 5, 7).reshape(batch, -1, patch_dim) + return linear(x, self.weights["patch_embedding.weight"], self.weights.get("patch_embedding.bias")) + + def _condition(self, timestep, encoder_hidden_states) -> tuple: + """Per-token conditioning. ``timestep`` is ``[B, L]`` (one level per token).""" + batch, seq = timestep.shape + t_freq = timestep_embedding(timestep.reshape(-1), self.freq_dim).astype( + weight_dtype(self.weights["condition_embedder.time_embedder.linear_1.weight"])) + temb = linear(t_freq, self.weights["condition_embedder.time_embedder.linear_1.weight"], + self.weights["condition_embedder.time_embedder.linear_1.bias"]) + temb = silu(temb) + temb = linear(temb, self.weights["condition_embedder.time_embedder.linear_2.weight"], + self.weights["condition_embedder.time_embedder.linear_2.bias"]) + timestep_proj = linear(silu(temb), self.weights["condition_embedder.time_proj.weight"], + self.weights["condition_embedder.time_proj.bias"]) + timestep_proj = timestep_proj.reshape(batch, seq, 6, self.hidden_size) + + ehs = linear(encoder_hidden_states, self.weights["condition_embedder.text_embedder.linear_1.weight"], + self.weights["condition_embedder.text_embedder.linear_1.bias"]) + ehs = gelu_tanh(ehs) + ehs = linear(ehs, self.weights["condition_embedder.text_embedder.linear_2.weight"], + self.weights["condition_embedder.text_embedder.linear_2.bias"]) + temb_out = temb.reshape(batch, seq, self.hidden_size) + return temb_out, timestep_proj, ehs + + def _output(self, hidden_states, temb_out, *, batch, frames, height, width) -> mx.array: + import mlx.core as mx + + pt, ph, pw = self.patch_size + post_pt, post_ph, post_pw = frames // pt, height // ph, width // pw + # Per-token output modulation: scale_shift_table[1,2,dim] + temb[B,L,1,dim]. + e = self.weights["scale_shift_table"][None] + temb_out[:, :, None, :].astype(mx.float32) + shift, scale = [part.squeeze(2) for part in mx.split(e, 2, axis=2)] + norm = layer_norm(hidden_states.astype(mx.float32), eps=self.eps) + norm = (norm * (1.0 + scale) + shift).astype(weight_dtype(self.weights["proj_out.weight"])) + out = linear(norm, self.weights["proj_out.weight"], self.weights["proj_out.bias"]) + out = out.reshape(batch, post_pt, post_ph, post_pw, pt, ph, pw, self.out_channels) + out = out.transpose(0, 7, 1, 4, 2, 5, 3, 6) + return out.reshape(batch, self.out_channels, frames, height, width) + + def _forward(self, hidden_states, encoder_hidden_states, timestep, cos, sin) -> mx.array: + batch, _, frames, height, width = hidden_states.shape + hidden = self._patch_embed(hidden_states) + temb_out, timestep_proj, ehs = self._condition(timestep, encoder_hidden_states) + for block in self.blocks: + hidden = block(hidden, ehs, timestep_proj, cos, sin) + return self._output(hidden, temb_out, batch=batch, frames=frames, height=height, width=width) + + def __call__(self, hidden_states, encoder_hidden_states, timestep, freqs_cis) -> mx.array: + cos, sin = freqs_cis + if self._enable_compile and cos is not None: + import mlx.core as mx + + # One traced graph per input signature, and each pins its own copy of + # the quantized weights. --refine denoises at two resolutions, so + # keeping both alive doubles resident DiT memory. Retire the previous + # graph when the signature changes. + signature = (hidden_states.shape, encoder_hidden_states.shape, timestep.shape) + if self._compiled_forward is not None and signature != self._compiled_signature: + self._compiled_forward = None + self._compiled_signature = None + mx.clear_cache() + if self._compiled_forward is None: + self._compiled_forward = mx.compile(self._forward) + self._compiled_signature = signature + compiled_forward = self._compiled_forward + try: + return compiled_forward(hidden_states, encoder_hidden_states, timestep, cos, sin) + except Exception as exc: # noqa: BLE001 - some graphs may not trace; fall back to eager. + logger.warning( + "Wan2.2 mx.compile forward failed (%s); falling back to eager execution.", + exc, + ) + self._enable_compile = False + self._compiled_forward = None + self._compiled_signature = None + return self._forward(hidden_states, encoder_hidden_states, timestep, cos, sin) + + +def mlx_wan22_dit_from_diffusers_safetensors( + checkpoint_path: str | Path, + config_path: str | Path, + *, + dtype: str = "fp16", + num_blocks: int | None = None, + quantization=None, + compile: bool = False, +) -> MLXWan22DiT: + """Load Wan2.2-TI2V-5B (FullAttn) into ``MLXWan22DiT`` via the dense loader.""" + dense = mlx_dit_from_diffusers_safetensors(checkpoint_path, + config_path, + dtype=dtype, + num_blocks=num_blocks, + quantization=quantization) + inner_dim = int(dense.config["num_attention_heads"]) * int(dense.config["attention_head_dim"]) + blocks = [ + MLXWan22TransformerBlock(block.weights, + dim=inner_dim, + ffn_dim=int(dense.config["ffn_dim"]), + num_heads=int(dense.config["num_attention_heads"]), + eps=float(dense.config.get("eps", 1e-6))) for block in dense.blocks + ] + return MLXWan22DiT(dense.weights, blocks, dense.config, compile=compile) + + +def mlx_wan22_dit_from_mlx_checkpoint( + checkpoint_dir: str | Path, + *, + compile: bool = False, +) -> MLXWan22DiT: + """Rewrap a persisted MLX DiT checkpoint with Wan2.2 conditioning. + + The generic checkpoint loader intentionally rebuilds ``MLXWanDiT`` because + it is also used by the Wan2.1 runtime. Wan2.2 TI2V has the same weight + layout but needs per-token timestep modulation, so callers must rewrap the + loaded weights and blocks as :class:`MLXWan22DiT` before sampling. + """ + from fastvideo.mlx_runtime.checkpoint import load_mlx_dit_checkpoint + + dense = load_mlx_dit_checkpoint(checkpoint_dir) + inner_dim = int(dense.config["num_attention_heads"]) * int(dense.config["attention_head_dim"]) + blocks = [ + MLXWan22TransformerBlock( + block.weights, + dim=inner_dim, + ffn_dim=int(dense.config["ffn_dim"]), + num_heads=int(dense.config["num_attention_heads"]), + eps=float(dense.config.get("eps", 1e-6)), + ) for block in dense.blocks + ] + return MLXWan22DiT(dense.weights, blocks, dense.config, compile=compile) diff --git a/fastvideo/mlx_runtime/wan22_sample.py b/fastvideo/mlx_runtime/wan22_sample.py new file mode 100644 index 000000000..543a12272 --- /dev/null +++ b/fastvideo/mlx_runtime/wan22_sample.py @@ -0,0 +1,113 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Dense DMD sampling for MLXWan22DiT (Wan2.2 per-token timestep). + +Matches the FastVideo pipeline's warped DMD schedule (``warp_denoising_step=True``, +``dmd_denoising_steps=[1000,757,522]``, ``flow_shift=5.0`` for TI2V-5B) rather +than treating raw step indices as continuous timesteps (a bug in early demos). +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from collections.abc import Sequence + +import numpy as np + +from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step, pred_noise_to_pred_video + +if TYPE_CHECKING: + import mlx.core as mx + + from fastvideo.mlx_runtime.wan22 import MLXWan22DiT + + +def build_wan22_dmd_schedule( + dmd_denoising_steps: Sequence[int] | None = None, + *, + flow_shift: float = 5.0, + warp_denoising_step: bool = True, +) -> tuple[MLXDMDSchedule, list[float]]: + """ + Build the flow-matching schedule and continuous timesteps used for Wan2.2 DMD sampling. + + Parameters: + dmd_denoising_steps (Sequence[int] | None): Denoising step values to use; defaults to 1000, 757, and 522. + flow_shift (float): Flow-matching shift applied when constructing the schedule. + warp_denoising_step (bool): Whether to convert denoising steps to scheduler-warped continuous timesteps. + + Returns: + tuple[MLXDMDSchedule, list[float]]: The DMD schedule and corresponding continuous timesteps. + """ + import torch + + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + + steps = list(dmd_denoising_steps or [1000, 757, 522]) + scheduler = FlowMatchEulerDiscreteScheduler(shift=flow_shift) + schedule = MLXDMDSchedule.from_torch_scheduler(scheduler) + step_idx = torch.tensor(steps, dtype=torch.long) + if warp_denoising_step: + warped = torch.cat((scheduler.timesteps.cpu(), torch.tensor([0.0], dtype=torch.float32))) + timesteps = [float(t) for t in warped[1000 - step_idx]] + else: + timesteps = [float(s) for s in steps] + return schedule, timesteps + + +def sample_wan22_dmd( + model: MLXWan22DiT, + encoder_hidden_states: mx.array, + noise_latents: mx.array, + freqs_cis: tuple, + *, + dmd_denoising_steps: Sequence[int] | None = None, + flow_shift: float = 5.0, + warp_denoising_step: bool = True, + seed: int = 0, +) -> mx.array: + """ + Generate clean video latents from noisy latents using iterative DMD denoising. + + Parameters: + noise_latents (mx.array): Initial noisy video latents. + freqs_cis (tuple): Rotary positional frequency tensors used by the model. + dmd_denoising_steps (Sequence[int] | None): DMD denoising steps, or the default schedule when omitted. + flow_shift (float): Flow-matching schedule shift. + warp_denoising_step (bool): Whether to warp the denoising timesteps. + seed (int): Seed for reproducible intermediate re-noising. + + Returns: + mx.array: Denoised video latents. + """ + import mlx.core as mx + + schedule, timesteps = build_wan22_dmd_schedule(dmd_denoising_steps, + flow_shift=flow_shift, + warp_denoising_step=warp_denoising_step) + # NumPy RNG so re-noise is bit-reproducible across MLX / torch A/B dumps. + renoise_rng = np.random.default_rng(seed) + latents = noise_latents + batch, _c, frames, height, width = latents.shape + pt, ph, pw = model.patch_size + tokens = (frames // pt) * (height // ph) * (width // pw) + last = len(timesteps) - 1 + for i, t in enumerate(timesteps): + ts = mx.full((batch, tokens), float(t), dtype=mx.float32) + pred = model(latents.astype(mx.float16), encoder_hidden_states, ts, freqs_cis) + ni = latents.astype(mx.float32) + pn = pred.astype(mx.float32) + if i < last: + renoise = mx.array(renoise_rng.standard_normal(tuple(latents.shape)).astype(np.float32)) + latents = dmd_step( + latents=ni, + noise_input_latent=ni, + pred_noise=pn, + schedule=schedule, + timestep=float(t), + next_timestep=float(timesteps[i + 1]), + noise=renoise, + ).astype(latents.dtype) + else: + latents = pred_noise_to_pred_video(pn, ni, schedule.sigma_for(float(t))).astype(latents.dtype) + mx.eval(latents) + return latents diff --git a/fastvideo/mlx_runtime/wan_vae.py b/fastvideo/mlx_runtime/wan_vae.py new file mode 100644 index 000000000..812e2c583 --- /dev/null +++ b/fastvideo/mlx_runtime/wan_vae.py @@ -0,0 +1,557 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan VAE decode helpers for Apple Silicon MLX inference. + +Two decode backends: + +1. **TAEHV (primary / fast)** — Tiny AutoEncoder (madebyollin/taehv). Fully + MLX-native Conv2d path. ``taew2_1.pth`` for Wan2.1 (z_dim=16), + ``taew2_2.pth`` for Wan2.2 5B (z_dim=48, patch_size=2). Expected decode + wall-clock ~seconds vs ~minutes for the full 3D VAE on MPS. + +2. **Full AutoencoderKLWan (reference / quality)** — denormalize with + ``latents_mean`` / ``latents_std`` then torch decode (MPS preferred). Used + for parity gates and when TAEHV is unavailable. A pure-MLX 3D-conv port of + the residual Wan2.2 decoder is left as follow-up (causal feat-cache + + residual up blocks are large); TAEHV covers the product latency path. + +Diffusion latents from the DiT are **not** mean/std-normalized for TAEHV +(matching ``taehv_decode.py``); full VAE decode **does** denormalize first +(matching ``mlx_wan_prompt_to_video.decode_latents_to_video``). +""" + +from __future__ import annotations + +import hashlib +import json +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import numpy as np + +from fastvideo.mlx_runtime.memory import cleanup_mlx, cleanup_torch_mps + +GIB = 1024**3 + +TAEW2_1_URL = "https://raw.githubusercontent.com/madebyollin/taehv/main/taew2_1.pth" +TAEW2_2_URL = ("https://raw.githubusercontent.com/madebyollin/taehv/" + "563f40bdc820ed86bcad72ea515ee48f06bd22ec/taew2_2.pth") +# Validated 2026-07-02 / 2026-07-09 against upstream madebyollin/taehv. +TAEW2_1_SHA256 = "d26151e76cdc2c9424bef988de874b33d9a53f30ef3060cd556c429c469c797e" +TAEW2_2_SHA256 = "d053e216ca50e2bb837bbcd79b85f0366bea00e5938025572382a773b74c559a" + +DecodeBackend = Literal["taehv", "taehv-torch", "wan-vae"] + + +def _cache_dir() -> Path: + """Return the local directory used to cache TAEHV files.""" + return Path.home() / ".cache" / "fastvideo" / "taehv" + + +def _sha256(path: Path) -> str: + """Compute the SHA-256 digest of a file. + + Parameters: + path (Path): Path to the file to hash. + + Returns: + str: Lowercase hexadecimal SHA-256 digest. + """ + digest = hashlib.sha256() + with path.open("rb") as handle: + for chunk in iter(lambda: handle.read(1 << 20), b""): + digest.update(chunk) + return digest.hexdigest() + + +def _verify_checkpoint(path: Path, expected_digest: str) -> None: + """ + Verify a checkpoint's SHA-256 digest against a required expected digest. + + Parameters: + path (Path): Path to the checkpoint file. + expected_digest (str): Expected lowercase SHA-256 digest. + + Raises: + RuntimeError: If verification is enabled and the checkpoint digest does not match. + """ + if len(expected_digest) != 64 or any(char not in "0123456789abcdef" for char in expected_digest): + raise ValueError("A valid lowercase SHA-256 digest is required for bundled TAEHV checkpoints") + actual = _sha256(path) + if actual != expected_digest: + raise RuntimeError(f"TAEHV checkpoint at {path} failed sha256 verification " + f"(expected {expected_digest}, got {actual}). " + "Delete the file to re-download it.") + + +def ensure_taehv_checkpoint(*, z_dim: int, checkpoint_path: Path | None = None) -> Path: + """ + Return a validated TAEHV checkpoint for the specified latent channel count. + + Parameters: + z_dim (int): Number of latent channels, supported values are 16 and 48. + checkpoint_path (Path | None): Optional existing checkpoint path to validate and use. + + Returns: + Path: Path to the validated TAEHV checkpoint. + + Raises: + FileNotFoundError: If the supplied checkpoint path does not exist. + ValueError: If no checkpoint is mapped to the specified latent channel count. + """ + if checkpoint_path is not None: + if not checkpoint_path.exists(): + raise FileNotFoundError(f"TAEHV checkpoint not found: {checkpoint_path}") + return checkpoint_path + if z_dim == 16: + name, url, expect = "taew2_1.pth", TAEW2_1_URL, TAEW2_1_SHA256 + elif z_dim == 48: + name, url, expect = "taew2_2.pth", TAEW2_2_URL, TAEW2_2_SHA256 + else: + raise ValueError(f"No TAEHV checkpoint mapped for z_dim={z_dim} (supported: 16, 48)") + path = _cache_dir() / name + if not path.exists(): + path.parent.mkdir(parents=True, exist_ok=True) + print(f"Downloading {url} -> {path}") + import socket + import tempfile + # Download to a temporary file, verify, then atomically rename. + with tempfile.NamedTemporaryFile( + mode="wb", + dir=path.parent, + prefix=f".tmp_{name}_", + suffix=".pth", + delete=False, + ) as tmp_file: + tmp_path = Path(tmp_file.name) + try: + old_timeout = socket.getdefaulttimeout() + socket.setdefaulttimeout(300) + try: + urllib.request.urlretrieve(url, tmp_path) # noqa: S310 - public pinned artifact. + finally: + socket.setdefaulttimeout(old_timeout) + _verify_checkpoint(tmp_path, expect) + tmp_path.replace(path) + except Exception: + tmp_path.unlink(missing_ok=True) + raise + else: + _verify_checkpoint(path, expect) + return path + + +@dataclass(frozen=True) +class WanVAEConfigView: + """Minimal config fields needed for denormalize + spatial scale.""" + + z_dim: int + latents_mean: tuple[float, ...] + latents_std: tuple[float, ...] + scale_factor_spatial: int = 8 + scale_factor_temporal: int = 4 + patch_size: int | None = None + vae_dir: Path | None = None + + @classmethod + def from_vae_dir(cls, vae_dir: Path) -> WanVAEConfigView: + """ + Load Wan VAE configuration values from a directory. + + Parameters: + vae_dir (Path): Directory containing the VAE ``config.json`` file. + + Returns: + WanVAEConfigView: Configuration loaded from the VAE directory. + """ + cfg = json.loads((vae_dir / "config.json").read_text()) + return cls( + z_dim=int(cfg["z_dim"]), + latents_mean=tuple(float(x) for x in cfg["latents_mean"]), + latents_std=tuple(float(x) for x in cfg["latents_std"]), + scale_factor_spatial=int(cfg.get("scale_factor_spatial", 8)), + scale_factor_temporal=int(cfg.get("scale_factor_temporal", 4)), + patch_size=cfg.get("patch_size"), + vae_dir=vae_dir, + ) + + +def denormalize_latents_np(latents: np.ndarray, config: WanVAEConfigView) -> np.ndarray: + """ + Denormalize Wan VAE latent values using the configured means and standard deviations. + + Parameters: + latents (np.ndarray): Latent values in normalized form. + config (WanVAEConfigView): Wan VAE latent statistics. + + Returns: + np.ndarray: Denormalized latent values as float32. + """ + mean = np.asarray(config.latents_mean, dtype=np.float32).reshape(1, -1, 1, 1, 1) + std = np.asarray(config.latents_std, dtype=np.float32).reshape(1, -1, 1, 1, 1) + return latents.astype(np.float32) * std + mean + + +# --------------------------------------------------------------------------- +# MLX TAEHV decoder (Conv2d stack — primary fully-MLX product path) +# --------------------------------------------------------------------------- + + +def _mlx_conv2d(x: Any, weight: Any, bias: Any, *, stride: int = 1) -> Any: + import mlx.core as mx + + # x: NCHW, weight: OIHW + y = mx.conv2d(x.transpose(0, 2, 3, 1), weight.transpose(0, 2, 3, 1), stride=stride, padding=1) + y = y.transpose(0, 3, 1, 2) + if bias is not None: + y = y + bias.reshape(1, -1, 1, 1) + return y + + +def _mlx_conv2d_1x1(x: Any, weight: Any, bias: Any = None) -> Any: + """ + Applies a 1×1 convolution to an MLX tensor in channel-first layout. + + Parameters: + x (Any): Input tensor with shape [batch, channels, height, width]. + weight (Any): Convolution weights. + bias (Any, optional): Optional output-channel bias. + + Returns: + Any: The convolved tensor with shape [batch, output_channels, height, width]. + """ + import mlx.core as mx + + y = mx.conv2d(x.transpose(0, 2, 3, 1), weight.transpose(0, 2, 3, 1), stride=1, padding=0) + y = y.transpose(0, 3, 1, 2) + if bias is not None: + y = y + bias.reshape(1, -1, 1, 1) + return y + + +def _load_torch_state(path: Path) -> dict[str, np.ndarray]: + """Load a PyTorch state dictionary as NumPy arrays. + + Parameters: + path (Path): Path to the PyTorch checkpoint. + + Returns: + dict[str, np.ndarray]: State dictionary with tensors converted to NumPy arrays. + """ + import torch + + sd = torch.load(path, map_location="cpu", weights_only=True) + return {k: v.detach().float().cpu().numpy() for k, v in sd.items()} + + +class MLXTAEHVDecoder: + """Minimal MLX port of TAEHV ``decoder`` (parallel-over-time MemBlocks).""" + + def __init__(self, checkpoint_path: Path, *, z_dim: int) -> None: + """Initialize the TAEHV decoder from a checkpoint for the specified latent dimensionality. + + Parameters: + checkpoint_path (Path): Path to the TAEHV checkpoint. + z_dim (int): Number of latent channels, determining the decoder patch size. + """ + import mlx.core as mx + + self.checkpoint_path = Path(checkpoint_path) + self.latent_channels = z_dim + # Derive patch_size from z_dim: 48 channels → patch_size=2, 16 → patch_size=1 + self.patch_size = 2 if z_dim == 48 else 1 + self.image_channels = 3 + self.frames_to_trim = 3 # TGrow strides (1,2,2) → 2**2 - 1 for w2.1/w2.2 defaults + + sd = _load_torch_state(self.checkpoint_path) + # Patch TGrow kernels like upstream TAEHV.patch_tgrow_layers. + self.weights = {k: mx.array(v) for k, v in sd.items()} + self._n_f = [256, 128, 64, 64] + + def decode_ntchw(self, latents_ntchw: Any) -> Any: + """ + Decode latent video batches into clipped RGB frames. + + Parameters: + latents_ntchw (Any): Latents with shape ``[N, T, C, H, W]`` and the + decoder's configured latent channel count. + + Returns: + Any: Decoded frames with shape ``[N, T_out, 3, H_out, W_out]`` and values + clipped to the range ``[0, 1]``. + """ + import mlx.core as mx + + x = latents_ntchw + n, t, c, h, w = x.shape + if c != self.latent_channels: + raise ValueError(f"expected C={self.latent_channels}, got {c}") + x = x.reshape(n * t, c, h, w) + x = self._run_decoder_parallel(x, n=n) + # Pixel-shuffle if patch_size > 1: (NT, 3*p*p, H, W) -> (NT, 3, H*p, W*p) + if self.patch_size > 1: + p = self.patch_size + nt, c_out, hh, ww = x.shape + x = x.reshape(nt, self.image_channels, p, p, hh, ww) + x = x.transpose(0, 1, 4, 2, 5, 3).reshape(nt, self.image_channels, hh * p, ww * p) + _, c_out, hh, ww = x.shape + t_out = x.shape[0] // n + x = x.reshape(n, t_out, c_out, hh, ww) + if self.frames_to_trim > 0 and t_out > self.frames_to_trim: + x = x[:, self.frames_to_trim:] + return mx.clip(x, 0.0, 1.0) + + def _run_decoder_parallel(self, x: Any, *, n: int) -> Any: + """ + Apply the TAEHV decoder stack to flattened batch and temporal frames while preserving temporal memory. + """ + import mlx.core as mx + + w = self.weights + + def memblock(base: int, xx: Any, past: Any) -> Any: + """ + Apply a temporal memory block to the current and past feature tensors. + + Parameters: + base (int): Decoder block index used to select the block weights. + xx (Any): Current feature tensor. + past (Any): Past feature tensor concatenated with the current features. + + Returns: + Any: Activated feature tensor produced by the memory block. + """ + cat = mx.concatenate([xx, past], axis=1) + h = _mlx_conv2d(cat, w[f"decoder.{base}.conv.0.weight"], w.get(f"decoder.{base}.conv.0.bias")) + h = mx.maximum(h, 0.0) + h = _mlx_conv2d(h, w[f"decoder.{base}.conv.2.weight"], w.get(f"decoder.{base}.conv.2.bias")) + h = mx.maximum(h, 0.0) + h = _mlx_conv2d(h, w[f"decoder.{base}.conv.4.weight"], w.get(f"decoder.{base}.conv.4.bias")) + skip_key = f"decoder.{base}.skip.weight" + skip = _mlx_conv2d_1x1(xx, w[skip_key], None) if skip_key in w else xx + return mx.maximum(h + skip, 0.0) + + def upsample2(xx: Any) -> Any: + """ + Upsample a four-dimensional tensor by a factor of two along its spatial dimensions. + + Parameters: + xx (Any): Tensor with shape `(N, C, H, W)`. + + Returns: + Any: Tensor with shape `(N, C, 2H, 2W)` containing replicated spatial values. + """ + nt, c, h, ww = xx.shape + xx = xx.reshape(nt, c, h, 1, ww, 1) + xx = mx.broadcast_to(xx, (nt, c, h, 2, ww, 2)) + return xx.reshape(nt, c, h * 2, ww * 2) + + def tgrow(base: int, xx: Any, stride: int) -> Any: + wt = w[f"decoder.{base}.conv.weight"] + out_ch = int(xx.shape[1]) * stride + if int(wt.shape[0]) > out_ch: + wt = wt[-out_ch:] + y = _mlx_conv2d_1x1(xx, wt, None) + if stride == 1: + return y + # TGrow.forward: (NT, C*stride, H, W) -> (NT*stride, C, H, W) + nt, c, h, ww = y.shape + c_in = c // stride + y = y.reshape(nt, stride, c_in, h, ww).transpose(0, 1, 2, 3, 4) + return y.reshape(nt * stride, c_in, h, ww) + + def mem_past(xx: Any) -> Any: + """ + Build a temporal memory tensor containing a zero frame followed by the preceding frame at each time step. + + Parameters: + xx (Any): Flattened batch and temporal tensor with shape ``(batch * time, channels, height, width)``. + + Returns: + Any: Tensor with the same shape as ``xx`` containing the preceding frame for each temporal position. + """ + nt, c, h, ww = xx.shape + t_cur = nt // n + x_ = xx.reshape(n, t_cur, c, h, ww) + # pad one zero frame at t=0, align past[t] = x[t-1] + past = mx.concatenate([mx.zeros_like(x_[:, :1]), x_[:, :-1]], axis=1) + return past.reshape(nt, c, h, ww) + + # 0 Clamp, 1 conv, 2 ReLU + x = mx.tanh(x / 3.0) * 3.0 + x = _mlx_conv2d(x, w["decoder.1.weight"], w.get("decoder.1.bias")) + x = mx.maximum(x, 0.0) + for mem_idx in (3, 4, 5): + x = memblock(mem_idx, x, mem_past(x)) + x = upsample2(x) + x = tgrow(7, x, 1) + x = _mlx_conv2d(x, w["decoder.8.weight"], w.get("decoder.8.bias")) + for mem_idx in (9, 10, 11): + x = memblock(mem_idx, x, mem_past(x)) + x = upsample2(x) + x = tgrow(13, x, 2) + x = _mlx_conv2d(x, w["decoder.14.weight"], w.get("decoder.14.bias")) + for mem_idx in (15, 16, 17): + x = memblock(mem_idx, x, mem_past(x)) + x = upsample2(x) + x = tgrow(19, x, 2) + x = _mlx_conv2d(x, w["decoder.20.weight"], w.get("decoder.20.bias")) + x = mx.maximum(x, 0.0) + x = _mlx_conv2d(x, w["decoder.22.weight"], w.get("decoder.22.bias")) + return x + + +def decode_latents_taehv_mlx( + latents_np: np.ndarray, + *, + z_dim: int | None = None, + checkpoint_path: Path | None = None, +) -> np.ndarray: + """ + Decode latent representations with the MLX TAEHV decoder. + + Parameters: + latents_np (np.ndarray): Latents arranged as [B, C, T, H, W]. + z_dim (int | None): Latent channel dimension used to select the decoder checkpoint. + checkpoint_path (Path | None): Optional path to a TAEHV checkpoint. + + Returns: + np.ndarray: Decoded pixels arranged as [B, T, H, W, 3] with values in [0, 1]. + + Raises: + ValueError: If `latents_np` does not have five dimensions. + """ + import mlx.core as mx + + if latents_np.ndim != 5: + raise ValueError(f"expected [B,C,T,H,W], got {latents_np.shape}") + c = latents_np.shape[1] + z = z_dim if z_dim is not None else c + ckpt = ensure_taehv_checkpoint(z_dim=z, checkpoint_path=checkpoint_path) + dec = MLXTAEHVDecoder(ckpt, z_dim=z) + # NTCHW + x = mx.array(latents_np.transpose(0, 2, 1, 3, 4).astype(np.float32)) + out = dec.decode_ntchw(x) # N T C H W + mx.eval(out) + arr = np.array(out) + # B T H W C + return arr.transpose(0, 1, 3, 4, 2) + + +def decode_latents_wan_vae_torch( + latents_np: np.ndarray, + *, + vae_dir: Path, + device: str = "auto", + dtype_name: str = "fp16", +) -> np.ndarray: + """Full AutoencoderKLWan decode on torch (MPS/CPU) with mean/std denormalize. + + Returns pixels ``[B, T, H, W, 3]`` float in [0, 1]. + """ + import torch + from diffusers import AutoencoderKLWan + from diffusers.video_processor import VideoProcessor + + if device == "auto": + device = "mps" if torch.backends.mps.is_available() else "cpu" + dtype = torch.float16 if dtype_name == "fp16" and device == "mps" else torch.float32 + config = WanVAEConfigView.from_vae_dir(vae_dir) + vae = AutoencoderKLWan.from_pretrained(vae_dir, torch_dtype=dtype, local_files_only=True).to(device) + vae.eval() + latents = torch.from_numpy(latents_np.astype(np.float32)).to(device=device, dtype=dtype) + mean = torch.tensor(config.latents_mean, device=device, dtype=dtype).view(1, -1, 1, 1, 1) + inv_std = (1.0 / torch.tensor(config.latents_std, device=device, dtype=dtype)).view(1, -1, 1, 1, 1) + latents = latents / inv_std + mean # matches prompt_to_video path + with torch.no_grad(): + video = vae.decode(latents, return_dict=False)[0] + video = VideoProcessor(vae_scale_factor=config.scale_factor_spatial).postprocess_video(video, output_type="np") + return video # [B, T, H, W, 3] + + +def decode_latents_to_video( + latents_np: np.ndarray, + output_path: Path, + *, + fps: int = 16, + backend: DecodeBackend = "taehv", + vae_dir: Path | None = None, + z_dim: int | None = None, + taehv_checkpoint: Path | None = None, + torch_device: str = "auto", +) -> dict[str, Any]: + """Decode latent video frames and export them as an MP4 file. + + Parameters: + latents_np (np.ndarray): Latent video representation to decode. + output_path (Path): Destination path for the MP4 file. + fps (int): Output video frame rate. + backend (DecodeBackend): Decoder backend to use. + vae_dir (Path | None): Directory containing the full Wan VAE when using + the ``wan-vae`` backend. + z_dim (int | None): Latent channel count for TAEHV decoding. + taehv_checkpoint (Path | None): Optional TAEHV checkpoint path. + torch_device (str): PyTorch device selection for PyTorch-based decoding. + + Returns: + dict[str, Any]: Decode time in seconds, backend name, output path, frame + count, and video resolution. + + Raises: + ValueError: If the full VAE backend lacks ``vae_dir`` or the backend is + unknown. + """ + import time + + from diffusers.utils import export_to_video + + t0 = time.perf_counter() + if backend in ("taehv", "taehv-torch"): + c = latents_np.shape[1] if z_dim is None else z_dim + if backend == "taehv": + video = decode_latents_taehv_mlx(latents_np, z_dim=c, checkpoint_path=taehv_checkpoint) + else: + # torch TAEHV (regression / parity reference) + import torch + from fastvideo.third_party.taehv import TAEHV + + ckpt = ensure_taehv_checkpoint(z_dim=c, checkpoint_path=taehv_checkpoint) + if torch_device == "auto": + torch_device = "mps" if torch.backends.mps.is_available() else "cpu" + dtype = torch.float16 if torch_device == "mps" else torch.float32 + model = TAEHV(str(ckpt)).to(device=torch_device, dtype=dtype).eval() + lat = torch.from_numpy(latents_np).to(device=torch_device, dtype=dtype) + with torch.no_grad(): + out = model.decode_video(lat.transpose(1, 2), parallel=True, show_progress_bar=False) + video = out[0].permute(0, 2, 3, 1).float().cpu().numpy()[None, ...] + # out is NTCHW -> need BTHWC; decode_video returns NTCHW for batch + if video.ndim == 4: + video = video[None] + elif backend == "wan-vae": + if vae_dir is None: + raise ValueError("vae_dir required for wan-vae backend") + video = decode_latents_wan_vae_torch(latents_np, vae_dir=vae_dir, device=torch_device) + else: + raise ValueError(f"unknown backend {backend}") + + decode_s = time.perf_counter() - t0 + output_path = Path(output_path) + output_path.parent.mkdir(parents=True, exist_ok=True) + # export_to_video expects list/array of frames HxWxC + frames = video[0] + frames = np.clip(frames, 0.0, 1.0) + export_to_video(frames, str(output_path), fps=fps) + if backend == "taehv": + cleanup_mlx() + else: + if backend == "taehv-torch": + del model, lat, out + cleanup_torch_mps() + return { + "decode_s": decode_s, + "backend": backend, + "output_path": str(output_path), + "num_frames": int(frames.shape[0]), + "resolution": f"{frames.shape[2]}x{frames.shape[1]}" if frames.ndim == 4 else None, + } diff --git a/fastvideo/mlx_runtime/windowed_attention.py b/fastvideo/mlx_runtime/windowed_attention.py new file mode 100644 index 000000000..2382f0583 --- /dev/null +++ b/fastvideo/mlx_runtime/windowed_attention.py @@ -0,0 +1,223 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Chunked non-causal sliding-window self-attention for MLX scaling studies. + +This module is intentionally standalone (``mlx.core`` + stdlib only) so it can +be micro-benchmarked without pulling in the DiT / FastVideo stack. + +Window policy +------------- +**Symmetric** sliding window (non-causal). For query index ``i`` the allowed +key indices are: + + sinks: ``j in [0, sink)`` (always visible to every query, if ``sink > 0``) + local: ``j in [max(0, i - half), min(S, i + half + 1))`` + where ``half = window // 2`` + +so each query sees roughly ``window + 1`` local keys (plus any sinks outside +that range). This is appropriate for a dense, bidirectional DiT denoise pass. + +Implementation note (FLOPs) +--------------------------- +A full-size additive attention mask still materialises an ``O(S^2)`` score +matrix inside SDPA and does **not** reduce work. Instead we tile the sequence +into query blocks and run ``mx.fast.scaled_dot_product_attention`` only against +the union of keys that block needs (local slice ± sinks). That makes +per-block work ``O(chunk * (window + sink) * D)`` and total work +``O(S * (window + sink) * D)``. +""" + +from __future__ import annotations + +import mlx.core as mx + + +def _default_scale(head_dim: int, scale: float | None) -> float: + """ + Determine the attention scaling factor from an explicit value or head dimension. + + Parameters: + head_dim (int): The attention head dimension used to derive the default scale. + scale (Optional[float]): An explicit scaling factor. + + Returns: + float: The explicit scale converted to a float, or the reciprocal square root of `head_dim`. + + Raises: + ValueError: If `scale` is not provided and `head_dim` is not positive. + """ + if scale is not None: + return float(scale) + if head_dim <= 0: + raise ValueError(f"head_dim must be positive, got {head_dim}") + return head_dim**-0.5 + + +def _validate_qkv(q: mx.array, k: mx.array, v: mx.array) -> tuple[int, int, int, int]: + """ + Validate compatible rank-4 query, key, and value tensors. + + Parameters: + q (mx.array): Query tensor shaped `(B, H, S, D)`. + k (mx.array): Key tensor with the same shape as `q`. + v (mx.array): Value tensor with the same shape as `q`. + + Returns: + tuple[int, int, int, int]: Batch size, head count, sequence length, and head dimension. + + Raises: + ValueError: If the tensors are not rank 4, do not have identical shapes, or have an empty sequence or head dimension. + """ + if q.ndim != 4 or k.ndim != 4 or v.ndim != 4: + raise ValueError(f"q/k/v must be rank-4 (B, H, S, D); got shapes " + f"{q.shape}, {k.shape}, {v.shape}") + b, h, s, d = q.shape + if k.shape != (b, h, s, d) or v.shape != (b, h, s, d): + raise ValueError(f"q/k/v shapes must match exactly; got q={q.shape}, k={k.shape}, v={v.shape}") + if s == 0: + raise ValueError("sequence length S must be > 0") + if d == 0: + raise ValueError("head dim D must be > 0") + return b, h, s, d + + +def full_attention( + q: mx.array, + k: mx.array, + v: mx.array, + scale: float | None = None, +) -> mx.array: + """ + Compute dense scaled dot-product attention over the full sequence. + + Parameters: + scale (float, optional): Attention scaling factor. If omitted, uses the + inverse square root of the head dimension. + + Returns: + mx.array: Attention output with shape ``(B, H, S, D)``. + """ + _, _, _, d = _validate_qkv(q, k, v) + sc = _default_scale(d, scale) + return mx.fast.scaled_dot_product_attention(q, k, v, scale=sc) + + +def _concat_kv_slices( + k: mx.array, + v: mx.array, + ranges: list[tuple[int, int]], +) -> tuple[mx.array, mx.array]: + """Concatenate non-overlapping ``[start, end)`` key/value slices along seq.""" + if not ranges: + raise ValueError("ranges must be non-empty") + if len(ranges) == 1: + s0, e0 = ranges[0] + return k[:, :, s0:e0, :], v[:, :, s0:e0, :] + k_parts = [k[:, :, s:e, :] for s, e in ranges] + v_parts = [v[:, :, s:e, :] for s, e in ranges] + return mx.concatenate(k_parts, axis=2), mx.concatenate(v_parts, axis=2) + + +def _key_ranges_for_block( + qs: int, + qe: int, + seq_len: int, + half: int, + sink: int, +) -> list[tuple[int, int]]: + """Return ordered, non-overlapping key ranges for a query block ``[qs, qe)``. + + Symmetric local window over every query in the block, plus global sinks + ``[0, sink)``. Overlap is merged into a single contiguous range when + possible so we avoid double-counting sink tokens. + """ + local_start = max(0, qs - half) + # Last query index is ``qe - 1``; its right edge is ``qe - 1 + half + 1 = qe + half``. + local_end = min(seq_len, qe + half) + if local_start >= local_end: + # Degenerate (should not happen for valid qs < qe); fall back to sinks only. + if sink > 0: + return [(0, min(sink, seq_len))] + raise ValueError(f"empty local key range for query block [{qs}, {qe})") + + if sink <= 0: + return [(local_start, local_end)] + + sink_end = min(sink, seq_len) + if local_start <= sink_end: + # Sinks abut or overlap the local window — one contiguous slice from 0. + return [(0, max(local_end, sink_end))] + # Gap between sinks and local window: two slices, concat at SDPA time. + return [(0, sink_end), (local_start, local_end)] + + +def windowed_attention( + q: mx.array, + k: mx.array, + v: mx.array, + window: int, + sink: int = 0, + scale: float | None = None, + *, + chunk_size: int | None = None, +) -> mx.array: + """ + Apply symmetric sliding-window self-attention with optional global sink positions. + + Parameters: + q (mx.array): Query tensor shaped `(B, H, S, D)`. + k (mx.array): Key tensor shaped `(B, H, S, D)`. + v (mx.array): Value tensor shaped `(B, H, S, D)`. + window (int): Symmetric attention window width in tokens; must be at least 1. + sink (int): Number of leading key positions available to every query; must + be between 0 and the sequence length. + scale (Optional[float]): Softmax scale. Defaults to `1 / sqrt(D)`. + chunk_size (Optional[int]): Query block length used for chunked processing. + Defaults to the smaller of `window` and 512. + + Returns: + mx.array: Attention output with the same shape as `q`. + + Raises: + ValueError: If the inputs or attention parameters are invalid. + RuntimeError: If a query block has no available keys. + """ + _, _, seq_len, d = _validate_qkv(q, k, v) + + if window < 1: + raise ValueError(f"window must be >= 1, got {window}") + if sink < 0: + raise ValueError(f"sink must be >= 0, got {sink}") + if sink > seq_len: + raise ValueError(f"sink ({sink}) cannot exceed sequence length ({seq_len})") + + sc = _default_scale(d, scale) + half = window // 2 + + # When the requested window is at least the sequence length, every query can + # see every key under a symmetric policy — fall back to one dense SDPA. + # (Sinks are redundant once the full key set is used.) + if window >= seq_len: + return mx.fast.scaled_dot_product_attention(q, k, v, scale=sc) + + chunk = min(window, 512) if chunk_size is None else int(chunk_size) + if chunk < 1: + raise ValueError(f"chunk_size must be >= 1, got {chunk}") + chunk = min(chunk, seq_len) + + outputs: list[mx.array] = [] + for qs in range(0, seq_len, chunk): + qe = min(seq_len, qs + chunk) + q_block = q[:, :, qs:qe, :] + ranges = _key_ranges_for_block(qs, qe, seq_len, half, sink) + k_block, v_block = _concat_kv_slices(k, v, ranges) + if k_block.shape[2] == 0: + raise RuntimeError(f"empty key set for query block [{qs}, {qe}) with window={window}, sink={sink}") + query_positions = mx.arange(qs, qe)[:, None] + key_positions = mx.array([position for start, end in ranges for position in range(start, end)])[None, :] + mask = mx.abs(query_positions - key_positions) <= half + if sink > 0: + mask = mask | (key_positions < sink) + out_block = mx.fast.scaled_dot_product_attention(q_block, k_block, v_block, scale=sc, mask=mask) + outputs.append(out_block) + + return mx.concatenate(outputs, axis=2) diff --git a/fastvideo/platforms/mps.py b/fastvideo/platforms/mps.py index 3367c109e..461360b80 100644 --- a/fastvideo/platforms/mps.py +++ b/fastvideo/platforms/mps.py @@ -48,7 +48,10 @@ class MpsPlatform(Platform): @classmethod def get_attn_backend_cls(cls, selected_backend: AttentionBackendEnum | None, head_size: int, dtype: torch.dtype) -> str: - # MPS supports SDPA (Scaled Dot-Product Attention) which is the most compatible + if selected_backend == AttentionBackendEnum.VIDEO_SPARSE_ATTN: + raise NotImplementedError("VIDEO_SPARSE_ATTN is not supported on MPS. Unset " + "FASTVIDEO_ATTENTION_BACKEND or set it to TORCH_SDPA.") + # MPS supports SDPA (Scaled Dot-Product Attention) which is the most compatible. logger.info("Using Torch SDPA backend for MPS.") return "fastvideo.attention.backends.sdpa.SDPABackend" diff --git a/fastvideo/tests/mlx/test_dmd_sampling.py b/fastvideo/tests/mlx/test_dmd_sampling.py new file mode 100644 index 000000000..2a14d1f39 --- /dev/null +++ b/fastvideo/tests/mlx/test_dmd_sampling.py @@ -0,0 +1,124 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only correctness contract for the on-device MLX DMD sampler. + +The sampler's arithmetic uses only ``+``/``-``/``*``, so NumPy arrays are a +valid stand-in for MLX arrays here: these tests validate the *math* against the +torch reference formulas without requiring an MLX/Metal device. +""" + +import numpy as np + +from fastvideo.mlx_runtime.sampling import ( + MLXDMDSchedule, + add_noise, + dmd_step, + pred_noise_to_pred_video, +) + + +def _reference_schedule() -> MLXDMDSchedule: + # A small monotonic flow-match-like schedule: timesteps 0..999, sigmas in + # (0, 1]. Exact values are irrelevant; the lookup must pick the nearest. + """Create the deterministic schedule used by the reference tests. + + Returns: + MLXDMDSchedule: A schedule with timesteps from 0 through 999 and linearly decreasing sigma values from 1.0 to 0.001. + """ + timesteps = np.arange(1000, dtype=np.float64) + sigmas = np.linspace(1.0, 1e-3, 1000, dtype=np.float64) + return MLXDMDSchedule(sigmas=sigmas, timesteps=timesteps) + + +def test_sigma_for_picks_nearest_timestep() -> None: + schedule = _reference_schedule() + # Exact hits. + assert schedule.sigma_for(0) == schedule.sigmas[0] + assert schedule.sigma_for(999) == schedule.sigmas[999] + # Nearest-neighbour rounding for a value between grid points. + assert schedule.sigma_for(522.4) == schedule.sigmas[522] + assert schedule.sigma_for(521.6) == schedule.sigmas[522] + + +def test_pred_noise_to_pred_video_matches_reference_formula() -> None: + rng = np.random.default_rng(0) + pred_noise = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32) + noise_input = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32) + sigma = 0.37 + + got = pred_noise_to_pred_video(pred_noise, noise_input, sigma) + expected = noise_input - sigma * pred_noise + np.testing.assert_allclose(got, expected, rtol=0, atol=0) + + +def test_add_noise_matches_reference_formula() -> None: + rng = np.random.default_rng(1) + clean = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32) + noise = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32) + sigma = 0.62 + + got = add_noise(clean, noise, sigma) + expected = (1.0 - sigma) * clean + sigma * noise + np.testing.assert_allclose(got, expected, rtol=0, atol=0) + + +def test_dmd_step_final_returns_clean_prediction() -> None: + schedule = _reference_schedule() + rng = np.random.default_rng(2) + pred_noise = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32) + noise_input = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32) + + out = dmd_step( + latents=noise_input, + noise_input_latent=noise_input, + pred_noise=pred_noise, + schedule=schedule, + timestep=522, + next_timestep=None, + ) + sigma = schedule.sigma_for(522) + np.testing.assert_allclose(out, noise_input - sigma * pred_noise, rtol=0, atol=0) + + +def test_dmd_step_intermediate_renoises_to_next_level() -> None: + """ + Verifies that an intermediate DMD step re-noises the denoised prediction at the next schedule level. + """ + schedule = _reference_schedule() + rng = np.random.default_rng(3) + pred_noise = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32) + noise_input = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32) + renoise = rng.standard_normal((1, 16, 3, 8, 8)).astype(np.float32) + + out = dmd_step( + latents=noise_input, + noise_input_latent=noise_input, + pred_noise=pred_noise, + schedule=schedule, + timestep=1000, + next_timestep=757, + noise=renoise, + ) + sigma = schedule.sigma_for(1000) + sigma_next = schedule.sigma_for(757) + pred_video = noise_input - sigma * pred_noise + expected = (1.0 - sigma_next) * pred_video + sigma_next * renoise + np.testing.assert_allclose(out, expected, rtol=0, atol=0) + + +def test_dmd_step_requires_noise_for_intermediate_step() -> None: + """Verify that an intermediate DMD step requires a noise tensor for re-noising.""" + schedule = _reference_schedule() + zeros = np.zeros((1, 16, 3, 8, 8), dtype=np.float32) + try: + dmd_step( + latents=zeros, + noise_input_latent=zeros, + pred_noise=zeros, + schedule=schedule, + timestep=1000, + next_timestep=757, + noise=None, + ) + except ValueError: + return + raise AssertionError("dmd_step must raise when re-noise is requested without noise") diff --git a/fastvideo/tests/mlx/test_frame_upsample.py b/fastvideo/tests/mlx/test_frame_upsample.py new file mode 100644 index 000000000..644fa0cc6 --- /dev/null +++ b/fastvideo/tests/mlx/test_frame_upsample.py @@ -0,0 +1,99 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only contracts for pixel-space frame resampling.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from fastvideo.mlx_runtime.frame_upsample import ( + PIXEL_UPSAMPLE_MODES, + unsharp, + upsample_frame, + upsample_frames, +) +from fastvideo.mlx_runtime.rife_interp import aligned_keyframe_count + + +def _frame(height: int = 24, width: int = 40, seed: int = 0) -> np.ndarray: + rng = np.random.default_rng(seed) + return rng.integers(0, 256, (height, width, 3), dtype=np.uint8) + + +@pytest.mark.parametrize("mode", PIXEL_UPSAMPLE_MODES) +def test_upsample_frame_hits_target_size_in_every_mode(mode: str) -> None: + out = upsample_frame(_frame(), width=80, height=48, mode=mode) + assert out.shape == (48, 80, 3) + assert out.dtype == np.uint8 + + +def test_upsample_frame_rejects_unknown_mode() -> None: + with pytest.raises(ValueError, match="Unsupported pixel upsample mode"): + upsample_frame(_frame(), width=80, height=48, mode="mitchell") + + +def test_upsample_frame_rejects_bad_shape() -> None: + with pytest.raises(ValueError, match="HxWx3"): + upsample_frame(np.zeros((24, 40), dtype=np.uint8), width=80, height=48) + + +def test_upsample_frame_rejects_nonpositive_size() -> None: + with pytest.raises(ValueError, match="must be positive"): + upsample_frame(_frame(), width=0, height=48) + + +def test_upsample_frame_at_target_size_is_a_passthrough() -> None: + frame = _frame() + np.testing.assert_array_equal(upsample_frame(frame, width=40, height=24), frame) + + +def test_nearest_upsample_is_exact_block_replication() -> None: + """A 2x nearest resize must duplicate pixels, not blend them.""" + frame = _frame(height=4, width=6) + out = upsample_frame(frame, width=12, height=8, mode="nearest") + np.testing.assert_array_equal(out, np.repeat(np.repeat(frame, 2, axis=0), 2, axis=1)) + + +def test_unsharp_zero_amount_is_identity() -> None: + frame = _frame() + assert unsharp(frame, 0.0) is frame + + +def test_unsharp_steepens_a_soft_edge() -> None: + # A blurred step edge: unsharp masking should raise its peak gradient. + import cv2 + + frame = np.zeros((24, 40, 3), dtype=np.uint8) + frame[:, 20:] = 200 + frame = cv2.GaussianBlur(frame, (0, 0), 2.0) + + def peak_gradient(image: np.ndarray) -> float: + return float(np.abs(np.diff(image[:, :, 0].astype(np.int32), axis=1)).max()) + + assert peak_gradient(unsharp(frame, 0.8)) > peak_gradient(frame) + + +def test_upsample_frames_preserves_order_and_leaves_input_alone() -> None: + frames = [_frame(seed=index) for index in range(3)] + originals = [frame.copy() for frame in frames] + out = upsample_frames(frames, width=80, height=48) + + assert len(out) == 3 + assert all(frame.shape == (48, 80, 3) for frame in out) + # Distinct inputs stay distinct and in order. + assert not np.array_equal(out[0], out[1]) + for frame, original in zip(frames, originals, strict=True): + np.testing.assert_array_equal(frame, original) + + +@pytest.mark.parametrize( + ("target_frames", "factor", "keyframes"), + [(1, 999, 1), (81, 2, 41), (81, 3, 29), (82, 2, 45), (121, 4, 33)], +) +def test_aligned_keyframes_expand_to_exact_target(target_frames: int, factor: int, keyframes: int) -> None: + assert aligned_keyframe_count(target_frames, factor) == keyframes + interpolated_frames = (keyframes - 1) * factor + 1 + assert interpolated_frames >= target_frames + assert keyframes % 4 == 1 + if keyframes > 1: + assert (keyframes - 5) * factor + 1 < target_frames diff --git a/fastvideo/tests/mlx/test_import_laziness.py b/fastvideo/tests/mlx/test_import_laziness.py new file mode 100644 index 000000000..3dabd08b2 --- /dev/null +++ b/fastvideo/tests/mlx/test_import_laziness.py @@ -0,0 +1,30 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Import hygiene for the lightweight Apple Silicon MLX test path.""" + +from __future__ import annotations + +import subprocess +import sys + +_HEAVY = "fastvideo.entrypoints.video_generator" + + +def _modules_after(import_line: str) -> set[str]: + """Return sys.modules after running one import in a fresh interpreter.""" + code = f"import sys; {import_line}; print('\\n'.join(sorted(sys.modules)))" + result = subprocess.run([sys.executable, "-c", code], check=True, capture_output=True, text=True) + return set(result.stdout.split()) + + +def test_mlx_runtime_import_adds_no_heavy_entrypoints() -> None: + """Importing the MLX runtime must not drag in the generator stack itself. + + ``fastvideo/__init__.py`` imports ``VideoGenerator`` eagerly, so every + ``fastvideo.*`` import loads it transitively. That is a property of the + top-level package rather than of this runtime, so compare against a plain + ``import fastvideo`` baseline and assert the MLX runtime adds nothing heavy + on top of it. + """ + baseline = _modules_after("import fastvideo") + with_runtime = _modules_after("import fastvideo.mlx_runtime.memory") + assert _HEAVY not in with_runtime - baseline diff --git a/fastvideo/tests/mlx/test_memory_limits.py b/fastvideo/tests/mlx/test_memory_limits.py new file mode 100644 index 000000000..0f6a02935 --- /dev/null +++ b/fastvideo/tests/mlx/test_memory_limits.py @@ -0,0 +1,100 @@ +# SPDX-License-Identifier: Apache-2.0 + +import os + +import pytest + +from fastvideo.mlx_runtime.memory import add_memory_limit_args, apply_memory_limits, gib_to_bytes + + +class _FakeMLX: + def __init__(self) -> None: + self.calls: list[tuple[str, int]] = [] + + def set_memory_limit(self, value: int) -> int: + """ + Record a requested memory limit and provide the prior limit. + + Parameters: + value (int): The memory limit to record. + + Returns: + int: The previously configured memory limit. + """ + self.calls.append(("memory", value)) + return 111 + + def set_cache_limit(self, value: int) -> int: + """Record and apply a cache memory limit. + + Parameters: + value (int): Cache memory limit in bytes. + + Returns: + int: The previously configured cache memory limit. + """ + self.calls.append(("cache", value)) + return 222 + + def set_wired_limit(self, value: int) -> int: + self.calls.append(("wired", value)) + return 333 + + +def test_gib_to_bytes_rejects_non_positive_values() -> None: + assert gib_to_bytes(None) is None + assert gib_to_bytes(1.5) == int(1.5 * 1024**3) + with pytest.raises(ValueError, match="positive"): + gib_to_bytes(0) + + +def test_apply_memory_limits_sets_mlx_limits_and_metrics(monkeypatch) -> None: + monkeypatch.delenv("PYTORCH_MPS_HIGH_WATERMARK_RATIO", raising=False) + monkeypatch.delenv("PYTORCH_MPS_LOW_WATERMARK_RATIO", raising=False) + fake_mlx = _FakeMLX() + + applied = apply_memory_limits( + mlx_memory_limit_gib=16, + mlx_disable_cache=True, + mlx_wired_limit_gib=12, + torch_mps_high_watermark_ratio=0.57, + mx_module=fake_mlx, + ) + + assert os.environ["PYTORCH_MPS_HIGH_WATERMARK_RATIO"] == "0.57" + assert os.environ["PYTORCH_MPS_LOW_WATERMARK_RATIO"] == "0.0" + assert fake_mlx.calls == [ + ("memory", 16 * 1024**3), + ("cache", 0), + ("wired", 12 * 1024**3), + ] + metrics = applied.as_metrics() + assert metrics["mlx_memory_limit_bytes"] == 16 * 1024**3 + assert metrics["mlx_cache_limit_bytes"] == 0 + assert metrics["previous_mlx_memory_limit_bytes"] == 111 + assert metrics["previous_mlx_cache_limit_bytes"] == 222 + assert metrics["previous_mlx_wired_limit_bytes"] == 333 + + +def test_apply_memory_limits_validates_watermarks() -> None: + with pytest.raises(ValueError, match="non-negative"): + apply_memory_limits(torch_mps_high_watermark_ratio=-1) + + +def test_add_memory_limit_args_uses_defaults() -> None: + import argparse + + parser = argparse.ArgumentParser() + add_memory_limit_args( + parser, + mlx_memory_limit_gib=16, + mlx_disable_cache=True, + torch_mps_high_watermark_ratio=0.57, + torch_mps_low_watermark_ratio=0.0, + ) + + args = parser.parse_args([]) + assert args.mlx_memory_limit_gib == 16 + assert args.mlx_disable_cache is True + assert args.torch_mps_high_watermark_ratio == 0.57 + assert args.torch_mps_low_watermark_ratio == 0.0 diff --git a/fastvideo/tests/mlx/test_mlx_checkpoint.py b/fastvideo/tests/mlx/test_mlx_checkpoint.py new file mode 100644 index 000000000..55bc9bd6e --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_checkpoint.py @@ -0,0 +1,123 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Round-trip tests for pre-quantized MLX DiT checkpoints. + +Saving an already-quantized DiT and reloading it must reproduce the original +forward pass exactly (same packed weights, scales, and config), for both plain +fp32 and int8-quantized models. Runs on any MLX backend (Metal or CPU). +""" + +from __future__ import annotations + +import json + +import numpy as np +import pytest + +pytest.importorskip("mlx.core", reason="MLX is required for checkpoint tests") + +from fastvideo.mlx_runtime.checkpoint import ( # noqa: E402 + MANIFEST_FILENAME, + load_mlx_dit_checkpoint, + save_mlx_dit_checkpoint, +) +from fastvideo.mlx_runtime.fastwan import ( # noqa: E402 + MLXQuantizationSpec, + QuantizedMatrix, +) +from fastvideo.tests.mlx.tiny_wan import ( # noqa: E402 + build_hf_config, + build_inputs, + build_tiny_wan_config, + build_torch_model, + mlx_dit_from_torch_model, + mlx_output, + mlx_rotary_embeddings, +) + + +def _forward(dit) -> np.ndarray: + """Run the DiT model with standard inputs and return its output as a NumPy array.""" + hidden_states, encoder_hidden_states, timestep = build_inputs() + return mlx_output(dit, hidden_states, encoder_hidden_states, timestep, mlx_rotary_embeddings(hidden_states)) + + +def test_plain_checkpoint_round_trip(distributed_setup, tmp_path) -> None: + dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config())) + before = _forward(dit) + + save_mlx_dit_checkpoint(dit, tmp_path / "ckpt") + loaded = load_mlx_dit_checkpoint(tmp_path / "ckpt") + + # JSON canonicalizes tuples to lists (matching what the Diffusers + # config.json loader produces), so compare JSON-normalized configs. + assert json.loads(json.dumps(loaded.config)) == json.loads(json.dumps(dit.config)) + assert len(loaded.blocks) == len(dit.blocks) + after = _forward(loaded) + np.testing.assert_array_equal(before, after) + + +def test_int8_checkpoint_round_trip_skips_requantization(distributed_setup, tmp_path) -> None: + spec = MLXQuantizationSpec.from_name("int8") + dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config()), + quantization=spec) + before = _forward(dit) + + save_mlx_dit_checkpoint(dit, tmp_path / "ckpt-int8") + loaded = load_mlx_dit_checkpoint(tmp_path / "ckpt-int8") + + # The loaded model must carry the same packed quantized weights -- no + # requantization pass -- and produce a bitwise-identical forward. + reloaded_q = loaded.weights["proj_out.weight"] + original_q = dit.weights["proj_out.weight"] + assert isinstance(reloaded_q, QuantizedMatrix) + assert reloaded_q.spec == spec + np.testing.assert_array_equal(np.array(reloaded_q.weight), np.array(original_q.weight)) + np.testing.assert_array_equal(np.array(reloaded_q.scales), np.array(original_q.scales)) + + block_q = loaded.blocks[0].weights["to_q.weight"] + assert isinstance(block_q, QuantizedMatrix) + + after = _forward(loaded) + np.testing.assert_array_equal(before, after) + + +def test_missing_checkpoint_raises_clear_error(tmp_path) -> None: + with pytest.raises(FileNotFoundError, match="Not an MLX DiT checkpoint"): + load_mlx_dit_checkpoint(tmp_path / "does-not-exist") + + +def test_existing_checkpoint_is_replaced(distributed_setup, tmp_path) -> None: + dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config())) + checkpoint_dir = save_mlx_dit_checkpoint(dit, tmp_path / "ckpt") + dit.config["checkpoint_revision"] = "second" + + save_mlx_dit_checkpoint(dit, checkpoint_dir) + + assert load_mlx_dit_checkpoint(checkpoint_dir).config["checkpoint_revision"] == "second" + + +def test_failed_checkpoint_overwrite_preserves_existing(distributed_setup, tmp_path) -> None: + dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config())) + checkpoint_dir = save_mlx_dit_checkpoint(dit, tmp_path / "ckpt") + weights_before = (checkpoint_dir / "mlx_dit.safetensors").read_bytes() + manifest_before = (checkpoint_dir / MANIFEST_FILENAME).read_bytes() + dit.config["not_json_serializable"] = object() + + with pytest.raises(TypeError, match="not JSON serializable"): + save_mlx_dit_checkpoint(dit, checkpoint_dir) + + assert (checkpoint_dir / "mlx_dit.safetensors").read_bytes() == weights_before + assert (checkpoint_dir / MANIFEST_FILENAME).read_bytes() == manifest_before + + +def test_future_format_version_is_rejected(distributed_setup, tmp_path) -> None: + dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config())) + ckpt = save_mlx_dit_checkpoint(dit, tmp_path / "ckpt") + + manifest_path = ckpt / MANIFEST_FILENAME + manifest = json.loads(manifest_path.read_text()) + manifest["format_version"] = 99 + manifest_path.write_text(json.dumps(manifest)) + + with pytest.raises(ValueError, match="format_version=99"): + load_mlx_dit_checkpoint(ckpt) diff --git a/fastvideo/tests/mlx/test_mlx_compile_parity.py b/fastvideo/tests/mlx/test_mlx_compile_parity.py new file mode 100644 index 000000000..06714ee48 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_compile_parity.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Guard the ``mx.compile`` path of the MLX FastWan forward. + +`mx.compile` on the DiT forward was silently broken: a NumPy scalar multiplying +a traced array (``np.sqrt(2/pi) * x`` in ``gelu_tanh``) dispatched through NumPy, +which evals the traced array — illegal during compile. It raised "Attempting to +eval an array during function transformations" (caught → eager fallback, no +speedup) or segfaulted the process. These tests keep the compiled forward tracing +cleanly and numerically identical to eager, so the ~1.4x compile speedup cannot +silently regress. See ``fastvideo/mlx_runtime/fastwan.py`` and +``docs/design/apple_silicon_benchmark_baseline.md``. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core", reason="MLX is required for the compile-parity tests") + +from fastvideo.mlx_runtime.fastwan import ( # noqa: E402 + MLXWanTransformerBlock, + gelu_tanh, + timestep_embedding, +) + +# One module-level generator so successive _rand calls of the same shape get +# distinct values (re-seeding per call made every (dim, dim) weight identical). +_RNG = np.random.default_rng(0) + +# Through MLX 0.31 the compiler emitted the same kernels as eager for these +# graphs, so bit-identity was the contract. MLX 0.32 fuses more aggressively and +# reassociates float ops, which moves a handful of elements by a few ULP (max +# observed: 2.4e-07 absolute, 6.4e-06 relative, on 41 of 7680 elements). That is +# reassociation, not a regression, so assert bit-identity where it is guaranteed +# and a tight tolerance otherwise. Either way a silent fallback to eager or a +# broken trace still fails the test. +_MLX_VERSION = tuple(int(part) for part in mx.__version__.split(".")[:2]) +_COMPILE_IS_BITWISE = _MLX_VERSION < (0, 32) + + +def _assert_compile_parity(eager: "mx.array", compiled: "mx.array") -> None: + """Compare an eager result against its compiled counterpart.""" + if _COMPILE_IS_BITWISE: + np.testing.assert_array_equal(np.array(eager), np.array(compiled)) + else: + np.testing.assert_allclose(np.array(eager), np.array(compiled), rtol=1e-5, atol=1e-6) + + +def _rand(*shape: int, scale: float = 1.0) -> "mx.array": + """ + Generate a float32 MLX array of normally distributed random values. + + Parameters: + shape (int): Dimensions of the generated array. + scale (float): Factor applied to the sampled values. + + Returns: + mx.array: The generated random array. + """ + return mx.array((_RNG.standard_normal(shape) * scale).astype(np.float32)) + + +def test_gelu_tanh_compiles_and_matches_eager() -> None: + """The tanh-GELU (the op that broke compile) traces and matches eager.""" + x = _rand(1, 120, 64) + eager = gelu_tanh(x) + mx.eval(eager) + compiled = mx.compile(gelu_tanh)(x) + mx.eval(compiled) + _assert_compile_parity(eager, compiled) + + +def test_timestep_embedding_compiles_and_matches_eager() -> None: + """``timestep_embedding`` (math.log fix sibling of gelu_tanh) compiles cleanly. + + Unlike gelu_tanh, the exp/sin/cos chain can reassociate under compile, so + we allow a tight tolerance rather than bit-identity. + """ + t = mx.array(np.array([0.0, 250.0, 500.0, 1000.0], dtype=np.float32)) + dim = 64 + eager = timestep_embedding(t, dim) + mx.eval(eager) + compiled = mx.compile(lambda steps: timestep_embedding(steps, dim))(t) + mx.eval(compiled) + np.testing.assert_allclose(np.array(eager), np.array(compiled), rtol=1e-4, atol=1e-4) + + +def _tiny_block_weights(dim: int, ffn: int) -> dict: + """ + Create randomized weights for a small transformer block test fixture. + + Parameters: + dim (int): Hidden dimension of the transformer block. + ffn (int): Intermediate dimension of the feed-forward network. + + Returns: + dict: Randomized attention, normalization, modulation, and feed-forward weights and biases. + """ + square = ["to_q", "to_k", "to_v", "to_out", "attn2.to_q", "attn2.to_k", "attn2.to_v", "attn2.to_out"] + weights = {f"{k}.weight": _rand(dim, dim, scale=0.05) for k in square} + weights.update({f"{k}.bias": _rand(dim, scale=0.05) for k in ["to_q", "to_k", "to_v", "to_out"]}) + weights.update({ + "scale_shift_table": _rand(1, 6, dim, scale=0.05), + "norm_q.weight": _rand(dim, scale=0.05), + "norm_k.weight": _rand(dim, scale=0.05), + "self_attn_residual_norm.norm.weight": _rand(dim, scale=0.05), + "self_attn_residual_norm.norm.bias": _rand(dim, scale=0.05), + "attn2.norm_q.weight": _rand(dim, scale=0.05), + "attn2.norm_k.weight": _rand(dim, scale=0.05), + "ffn.fc_in.weight": _rand(ffn, dim, scale=0.05), + "ffn.fc_in.bias": _rand(ffn, scale=0.05), + "ffn.fc_out.weight": _rand(dim, ffn, scale=0.05), + "ffn.fc_out.bias": _rand(dim, scale=0.05), + }) + return weights + + +def test_transformer_block_compiles_and_matches_eager() -> None: + """The full dense block (the mx.compile target's body) traces and matches.""" + dim, num_heads, head_dim, ffn, seq, ctx = 64, 4, 16, 128, 120, 32 + block = MLXWanTransformerBlock(_tiny_block_weights(dim, ffn), dim=dim, ffn_dim=ffn, num_heads=num_heads, eps=1e-6) + hidden = _rand(1, seq, dim, scale=0.05) + context = _rand(1, ctx, dim, scale=0.05) + temb = _rand(1, 6, dim, scale=0.05) + cos = _rand(seq, head_dim) + sin = _rand(seq, head_dim) + + eager = block(hidden, context, temb, (cos, sin)) + mx.eval(eager) + compiled = mx.compile(lambda h, e, t, c, s: block(h, e, t, (c, s)))(hidden, context, temb, cos, sin) + mx.eval(compiled) + + _assert_compile_parity(eager, compiled) diff --git a/fastvideo/tests/mlx/test_mlx_diffusers_compile_dispatch.py b/fastvideo/tests/mlx/test_mlx_diffusers_compile_dispatch.py new file mode 100644 index 000000000..ab42f80f5 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_diffusers_compile_dispatch.py @@ -0,0 +1,77 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for preserving MLX compile flags on raw Diffusers loads.""" + +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +mx = pytest.importorskip("mlx.core", reason="MLX required for compile dispatch") + + +def _write_tiny_diffusers_checkpoint(tmp_path: Path) -> tuple[Path, Path]: + import torch + from safetensors.torch import save_file + + config = { + "num_layers": 0, + "num_attention_heads": 1, + "attention_head_dim": 2, + "ffn_dim": 4, + "in_channels": 16, + "out_channels": 16, + "patch_size": [1, 2, 2], + "freq_dim": 2, + "eps": 1e-6, + } + config_path = tmp_path / "config.json" + checkpoint_path = tmp_path / "diffusion_pytorch_model.safetensors" + config_path.write_text(json.dumps(config)) + + hidden = 2 + weights = { + "patch_embedding.weight": torch.zeros(hidden, 64), + "condition_embedder.time_embedder.linear_1.weight": torch.zeros(hidden, 2), + "condition_embedder.time_embedder.linear_2.weight": torch.zeros(hidden, hidden), + "condition_embedder.time_proj.weight": torch.zeros(hidden * 6, hidden), + "condition_embedder.text_embedder.linear_1.weight": torch.zeros(hidden, 4096), + "condition_embedder.text_embedder.linear_2.weight": torch.zeros(hidden, hidden), + "scale_shift_table": torch.zeros(2, hidden), + "proj_out.weight": torch.zeros(64, hidden), + } + save_file(weights, checkpoint_path) + return checkpoint_path, config_path + + +def test_diffusers_loader_preserves_compile_flag(tmp_path: Path) -> None: + from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors + + checkpoint_path, config_path = _write_tiny_diffusers_checkpoint(tmp_path) + + compiled = mlx_dit_from_diffusers_safetensors( + checkpoint_path, + config_path, + compile=True, + num_blocks=0, + quantization=None, + ) + eager = mlx_dit_from_diffusers_safetensors( + checkpoint_path, + config_path, + compile=False, + num_blocks=0, + quantization=None, + ) + + assert compiled._enable_compile is True + assert eager._enable_compile is False + + +def test_prompt_to_video_cli_passes_compile_to_diffusers_loader() -> None: + root = Path(__file__).resolve().parents[3] + script = root / "examples/inference/basic/mlx_wan_prompt_to_video.py" + source = script.read_text() + + assert "compile=args.mlx_compile" in source diff --git a/fastvideo/tests/mlx/test_mlx_dit_parity.py b/fastvideo/tests/mlx/test_mlx_dit_parity.py new file mode 100644 index 000000000..716a19400 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_dit_parity.py @@ -0,0 +1,93 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Full-DiT parity: the MLX Wan runtime vs the PyTorch reference model. + +This is the M1 "trustworthy baseline" gate for the Apple Silicon path: a tiny +random-weight ``WanTransformer3DModel`` is run end to end (patch embed -> +condition -> transformer blocks -> unpatchify) in PyTorch and in +``fastvideo.mlx_runtime.fastwan.MLXWanDiT`` with identical weights, and the +outputs must match within pinned fp32 tolerances. + +Runs anywhere MLX is installed: on Apple Silicon it exercises the Metal +device, on Linux/CI it runs on MLX's CPU backend (``pip install 'mlx[cpu]'``) +-- the graph is identical, so CPU parity is the CI-friendly golden variant. + + pytest fastvideo/tests/mlx/test_mlx_dit_parity.py -v +""" + +from __future__ import annotations + +import numpy as np +import pytest + +pytest.importorskip("mlx.core", reason="MLX is required for DiT parity tests") + +from fastvideo.mlx_runtime.fastwan import MLXQuantizationSpec # noqa: E402 +from fastvideo.tests.mlx.tiny_wan import ( # noqa: E402 + build_hf_config, + build_inputs, + build_tiny_wan_config, + build_torch_model, + mlx_dit_from_torch_model, + mlx_output, + mlx_rotary_embeddings, + torch_reference_output, +) + +# Pinned fp32 tolerances for the full forward (patch embed through unpatchify), +# matching the single-block parity example. Measured max_abs_diff on the MLX +# CPU backend is ~1.7e-6, so 2e-4 keeps ~100x headroom for accumulation-order +# differences across MLX backends while still catching real math/layout bugs. +FP32_ATOL = 2e-4 +FP32_RTOL = 2e-4 + +# Quantized inference is lossy by design; gate it on signal-to-noise vs the +# fp32 MLX output instead of elementwise closeness. int8 (group size 64) on +# this tiny model measures ~43 dB; 20 dB leaves headroom without letting a +# broken dequant path (which lands near 0 dB) slip through. +INT8_MIN_SNR_DB = 20.0 + + +def test_full_dit_forward_matches_torch_reference(distributed_setup) -> None: + model = build_torch_model() + hidden_states, encoder_hidden_states, timestep = build_inputs() + + torch_out = torch_reference_output(model, hidden_states, encoder_hidden_states, timestep) + + dit = mlx_dit_from_torch_model(model, build_hf_config(build_tiny_wan_config())) + mlx_out = mlx_output(dit, hidden_states, encoder_hidden_states, timestep, mlx_rotary_embeddings(hidden_states)) + + assert mlx_out.shape == torch_out.shape + assert np.isfinite(mlx_out).all() + max_abs = float(np.abs(torch_out - mlx_out).max()) + assert np.allclose(torch_out, mlx_out, atol=FP32_ATOL, rtol=FP32_RTOL), ( + f"MLX full-DiT forward diverged from the torch reference: max_abs_diff={max_abs:.3e} " + f"(atol={FP32_ATOL}, rtol={FP32_RTOL})") + + +def test_full_dit_forward_int8_stays_close_to_fp32(distributed_setup) -> None: + """ + Ensure int8 quantization preserves the MLX DiT output within the required signal-to-noise ratio. + + Parameters: + distributed_setup: Test fixture providing the distributed test environment. + + Raises: + AssertionError: If the int8 output contains non-finite values or its signal-to-noise ratio is below the configured minimum. + """ + model = build_torch_model() + hidden_states, encoder_hidden_states, timestep = build_inputs() + hf_config = build_hf_config(build_tiny_wan_config()) + freqs_cis = mlx_rotary_embeddings(hidden_states) + + fp32_out = mlx_output( + mlx_dit_from_torch_model(model, hf_config), hidden_states, encoder_hidden_states, timestep, freqs_cis) + int8_out = mlx_output( + mlx_dit_from_torch_model(model, hf_config, quantization=MLXQuantizationSpec.from_name("int8")), + hidden_states, encoder_hidden_states, timestep, freqs_cis) + + assert np.isfinite(int8_out).all() + noise = float(np.mean(np.square(int8_out - fp32_out))) + signal = float(np.mean(np.square(fp32_out))) + snr_db = 10.0 * np.log10(signal / noise) if noise > 0 else float("inf") + assert snr_db >= INT8_MIN_SNR_DB, ( + f"int8-quantized MLX DiT output is too far from fp32: SNR {snr_db:.1f} dB < {INT8_MIN_SNR_DB} dB") diff --git a/fastvideo/tests/mlx/test_mlx_fast_spatial.py b/fastvideo/tests/mlx/test_mlx_fast_spatial.py new file mode 100644 index 000000000..77a3d42ad --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_fast_spatial.py @@ -0,0 +1,134 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only contracts for spatial fast mode + flag composition.""" + +from __future__ import annotations + +import numpy as np +import pytest + +from examples.inference.basic.mlx_wan22_generate import DEFAULT_HEIGHT, DEFAULT_NUM_FRAMES, DEFAULT_WIDTH +from fastvideo.mlx_runtime.fast_spatial import ( + DEFAULT_FAST_SPATIAL_SHARPEN, + apply_fast_spatial_upsample, + plan_fast_spatial, + resolve_spatial_mode, +) +from fastvideo.mlx_runtime.frame_upsample import upsample_frames + + +def _frames(count: int, height: int, width: int, seed: int = 0) -> list[np.ndarray]: + rng = np.random.default_rng(seed) + return [rng.integers(0, 256, (height, width, 3), dtype=np.uint8) for _ in range(count)] + + +def test_resolve_spatial_mode_priority() -> None: + assert resolve_spatial_mode(refine=False, fast_spatial=False) == "off" + assert resolve_spatial_mode(refine=False, fast_spatial=True) == "fast_spatial" + # Refine is the quality path and wins when both are requested. + assert resolve_spatial_mode(refine=True, fast_spatial=True) == "refine" + assert resolve_spatial_mode(refine=True, fast_spatial=False) == "refine" + + +def test_plan_fast_spatial_matches_refine_geometry() -> None: + spatial = plan_fast_spatial( + height=480, + width=832, + num_frames=81, + spatial_scale=2, + vae_spatial_compression=8, + ) + assert spatial.enabled + assert spatial.stage1_height == 240 + assert spatial.stage1_width == 416 + assert spatial.target_height == 480 + assert spatial.target_width == 832 + assert spatial.plan.stage1_latent_height == 30 + assert spatial.plan.stage2_latent_width == 104 + assert spatial.sharpen == DEFAULT_FAST_SPATIAL_SHARPEN + + +def test_plan_fast_spatial_disabled() -> None: + spatial = plan_fast_spatial(height=480, width=832, num_frames=81, enabled=False) + assert not spatial.enabled + assert spatial.scale == 1 + + +def test_plan_fast_spatial_accepts_wan22_defaults() -> None: + assert (DEFAULT_HEIGHT, DEFAULT_WIDTH, DEFAULT_NUM_FRAMES) == (448, 832, 121) + spatial = plan_fast_spatial( + height=DEFAULT_HEIGHT, + width=DEFAULT_WIDTH, + num_frames=DEFAULT_NUM_FRAMES, + spatial_scale=2, + vae_spatial_compression=16, + vae_temporal_compression=4, + patch_size=(1, 2, 2), + ) + assert (spatial.plan.stage1_latent_height, spatial.plan.stage1_latent_width) == (14, 26) + assert (spatial.plan.stage2_latent_height, spatial.plan.stage2_latent_width) == (28, 52) + + +def test_apply_fast_spatial_upsample_resizes_decoded_frames() -> None: + """The upsample is pixel-space: it takes decoded frames, not latents. + + Interpolating latents instead is what made spatial fast mode incoherent — + a blended Wan latent is off the decoder's manifold, so the decode came back + as the right silhouette under a smeared veil. + """ + spatial = plan_fast_spatial( + height=480, + width=832, + num_frames=81, + spatial_scale=2, + vae_spatial_compression=8, + upsample_mode="lanczos", + sharpen=0.0, + ) + frames = _frames(3, 240, 416) + out = apply_fast_spatial_upsample(frames, spatial) + + assert len(out) == 3 + assert all(frame.shape == (480, 832, 3) for frame in out) + assert all(frame.dtype == np.uint8 for frame in out) + expected = upsample_frames(frames, width=832, height=480, mode="lanczos", sharpen=0.0) + for got, want in zip(out, expected, strict=True): + np.testing.assert_array_equal(got, want) + # Inputs are not mutated. + assert all(frame.shape == (240, 416, 3) for frame in frames) + + +def test_apply_fast_spatial_upsample_applies_sharpen() -> None: + frames = _frames(1, 240, 416) + soft = plan_fast_spatial(height=480, width=832, num_frames=81, spatial_scale=2, sharpen=0.0) + crisp = plan_fast_spatial(height=480, width=832, num_frames=81, spatial_scale=2, sharpen=0.8) + assert not np.array_equal( + apply_fast_spatial_upsample(frames, soft)[0], + apply_fast_spatial_upsample(frames, crisp)[0], + ) + + +def test_apply_fast_spatial_noop_when_disabled() -> None: + frames = _frames(2, 32, 32, seed=1) + spatial = plan_fast_spatial(height=32, width=32, num_frames=5, enabled=False, patch_size=(1, 1, 1)) + out = apply_fast_spatial_upsample(frames, spatial) + for got, want in zip(out, frames, strict=True): + np.testing.assert_array_equal(got, want) + + +def test_plan_fast_spatial_rejects_bad_mode() -> None: + with pytest.raises(ValueError, match="upsample mode"): + plan_fast_spatial(height=480, width=832, num_frames=81, upsample_mode="bicubic") + + +def test_plan_fast_spatial_rejects_negative_sharpen() -> None: + with pytest.raises(ValueError, match="sharpen"): + plan_fast_spatial(height=480, width=832, num_frames=81, sharpen=-0.1) + + +def test_geometry_errors_name_fast_spatial_not_refine() -> None: + """Both modes share the resolution splitter; the error must name the caller. + + 832 is not divisible by 3, so scale 3 is rejected at plan time. + """ + with pytest.raises(ValueError, match=r"^fast-spatial requires"): + plan_fast_spatial(height=480, width=832, num_frames=81, spatial_scale=3) diff --git a/fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py b/fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py new file mode 100644 index 000000000..4e23cd040 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py @@ -0,0 +1,141 @@ +# SPDX-License-Identifier: Apache-2.0 + +import pytest + +from fastvideo.benchmarks.mlx_fastwan_bench import ( + ALLOWED_MODES, + BENCHMARK_PRESETS, + _html_grid, + _load_prompt_cases, + _mode_to_dtype_quant, + _parse_list, +) + + +@pytest.mark.parametrize( + ("mode", "expected"), + [ + ("fp16", ("fp16", None)), + ("bf16", ("bf16", None)), + ("int8", ("fp16", "int8")), + ("int4", ("fp16", "int4")), + ("mxfp8", ("fp16", "mxfp8")), + ("mxfp4", ("fp16", "mxfp4")), + ("nvfp4", ("fp16", "nvfp4")), + ], +) +def test_benchmark_modes_map_to_runtime_quantization(mode: str, expected: tuple[str, str | None]) -> None: + assert mode in ALLOWED_MODES + assert _mode_to_dtype_quant(mode) == expected + + +def test_benchmark_rejects_unknown_modes() -> None: + with pytest.raises(ValueError, match="Unsupported modes"): + _parse_list("fp16,not_a_mode", ALLOWED_MODES, "modes") + + +def test_load_prompt_cases_from_plain_text(tmp_path) -> None: + prompt_file = tmp_path / "prompts.txt" + prompt_file.write_text("\n# comment\nA fox runs through a forest.\nA raccoon walks in sunflowers.\n") + + cases = _load_prompt_cases("unused", prompt_file) + + assert [case.id for case in cases] == ["prompt-001", "prompt-002"] + assert [case.prompt for case in cases] == [ + "A fox runs through a forest.", + "A raccoon walks in sunflowers.", + ] + + +def test_load_prompt_cases_from_jsonl(tmp_path) -> None: + prompt_file = tmp_path / "prompts.jsonl" + prompt_file.write_text('{"id": "Fox Forest", "prompt": "A fox runs."}\n{"name": "clock", "caption": "A clock burns."}\n') + + cases = _load_prompt_cases("unused", prompt_file) + + assert [case.id for case in cases] == ["fox-forest", "clock"] + assert [case.prompt for case in cases] == ["A fox runs.", "A clock burns."] + + +def test_load_builtin_prompt_set() -> None: + cases = _load_prompt_cases("unused", None, "motion7") + assert len(cases) == 7 + assert cases[0].id == "beach-sunset" + + +def test_benchmark_presets_include_memory_tiers() -> None: + assert BENCHMARK_PRESETS["mac-16gb"].modes == "int8" + assert BENCHMARK_PRESETS["mac-16gb"].decoders == "taehv" + assert BENCHMARK_PRESETS["mac-16gb"].mlx_memory_limit_gib == 16.0 + assert BENCHMARK_PRESETS["mac-64gb"].decoders == "taehv,wan-vae" + + +def test_html_grid_includes_video_and_sync_controls() -> None: + rendered = _html_grid([ + { + "prompt_id": "fox", + "prompt": "A fox runs.", + "mode": "int8", + "decoder": "taehv", + "status": "ok", + "video_path": "fox/video_int8_taehv.mp4", + "total_s": 12.3, + "denoise_s": 10.0, + "decode_s": 1.0, + "peak_gib": 4.5, + } + ]) + + assert "Restart + play all" in rendered + assert "fox/video_int8_taehv.mp4" in rendered + assert "A fox runs." in rendered + + +def test_denoise_dmd_on_device_runs_tiny_dit_and_reports_step_times(distributed_setup) -> None: + mx = pytest.importorskip("mlx.core", reason="MLX is required for the on-device denoise test") + import numpy as np + import torch + + from fastvideo.benchmarks.mlx_fastwan_bench import denoise_dmd_on_device + from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import ( + FlowMatchEulerDiscreteScheduler, + ) + from fastvideo.tests.mlx.tiny_wan import ( + build_hf_config, + build_inputs, + build_tiny_wan_config, + build_torch_model, + mlx_dit_from_torch_model, + mlx_rotary_embeddings, + ) + + dit = mlx_dit_from_torch_model(build_torch_model(), build_hf_config(build_tiny_wan_config())) + hidden_states, encoder_hidden_states, _ = build_inputs() + schedule = MLXDMDSchedule.from_torch_scheduler(FlowMatchEulerDiscreteScheduler(shift=8.0)) + + timesteps = [1000, 757, 522] + generator = torch.Generator(device="cpu").manual_seed(7) + latents_seed = torch.randn(hidden_states.shape, generator=generator, dtype=torch.float32).numpy() + renoise_by_step = [ + torch.randn(hidden_states.shape, generator=generator, dtype=torch.float32).numpy() + for _ in range(len(timesteps) - 1) + ] + + latents_np, step_times = denoise_dmd_on_device( + mx=mx, + dit=dit, + latents=mx.array(latents_seed), + encoder_hidden_states=mx.array(encoder_hidden_states.numpy()), + freqs_cis=mlx_rotary_embeddings(hidden_states), + timesteps=timesteps, + renoise_by_step=renoise_by_step, + schedule=schedule, + dmd_step=dmd_step, + mx_dtype=mx.float32, + ) + + assert latents_np.shape == tuple(hidden_states.shape) + assert np.isfinite(latents_np).all() + assert len(step_times) == len(timesteps) + assert all(t > 0 for t in step_times) diff --git a/fastvideo/tests/mlx/test_mlx_prompt_enhance.py b/fastvideo/tests/mlx/test_mlx_prompt_enhance.py new file mode 100644 index 000000000..51bcd5992 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_prompt_enhance.py @@ -0,0 +1,86 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only contracts for local prompt enrichment.""" + +from __future__ import annotations + +import json + +from fastvideo.mlx_runtime.prompt_enhance import ( + DEFAULT_ENHANCE_SYSTEM_PROMPT, + enhance_prompt, + enhance_prompt_template, + enhance_result_as_metrics, + load_or_enhance_prompt, +) + + +def test_template_adds_cinematic_cues_to_thin_prompt() -> None: + raw = "A red fox in the snow" + out = enhance_prompt_template(raw) + assert raw in out + lower = out.lower() + assert "lens" in lower or "camera" in lower or "cinematic" in lower + assert "light" in lower + assert out.endswith(".") + + +def test_template_is_idempotent_on_rich_prompts() -> None: + rich = ( + "A red fox trotting through a snowy pine forest at golden hour, " + "cinematic wide shot on a 35mm lens, soft volumetric lighting, " + "gentle camera dolly, highly detailed film grain and color graded." + ) + assert enhance_prompt_template(rich) == " ".join(rich.split()) + + +def test_enhance_prompt_template_backend() -> None: + result = enhance_prompt("a paper boat on a stream", backend="template") + assert result.backend == "template" + assert result.changed + assert result.original == "a paper boat on a stream" + assert "paper boat" in result.enhanced.lower() + + +def test_enhance_prompt_rejects_empty() -> None: + try: + enhance_prompt(" ", backend="template") + except ValueError as exc: + assert "non-empty" in str(exc) + else: # pragma: no cover + raise AssertionError("expected ValueError") + + +def test_load_or_enhance_prompt_caches(tmp_path) -> None: + result1 = load_or_enhance_prompt( + "a lantern in the fog", + backend="template", + cache=True, + cache_dir=tmp_path, + ) + assert result1.backend == "template" + cached_files = list(tmp_path.glob("*.json")) + assert len(cached_files) == 1 + payload = json.loads(cached_files[0].read_text()) + assert payload["enhanced"] == result1.enhanced + + result2 = load_or_enhance_prompt( + "a lantern in the fog", + backend="template", + cache=True, + cache_dir=tmp_path, + ) + assert result2.enhanced == result1.enhanced + assert result2.backend == "cache" + assert result2.elapsed_s == 0.0 + + +def test_enhance_result_as_metrics_none() -> None: + metrics = enhance_result_as_metrics(None) + assert metrics["enhance_prompt"] is False + assert metrics["prompt_enhanced"] is None + + +def test_default_system_prompt_matches_streaming_contract() -> None: + # Lockstep with fastvideo/entrypoints/streaming/prompt/enhancer.py + assert "cinematic video generation" in DEFAULT_ENHANCE_SYSTEM_PROMPT.lower() + assert "just the enhanced prompt" in DEFAULT_ENHANCE_SYSTEM_PROMPT.lower() diff --git a/fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py b/fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py new file mode 100644 index 000000000..9a4982a0e --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py @@ -0,0 +1,85 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for the MLX prompt-to-video decode dispatch.""" + +from __future__ import annotations + +import importlib.util +from pathlib import Path + +import numpy as np + + +def _load_prompt_to_video_module(): + root = Path(__file__).resolve().parents[3] + script = root / "examples/inference/basic/mlx_wan_prompt_to_video.py" + spec = importlib.util.spec_from_file_location( + "mlx_wan_prompt_to_video_for_test", script + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_default_taehv_decode_uses_mlx_backend(monkeypatch, tmp_path): + module = _load_prompt_to_video_module() + import fastvideo.mlx_runtime.wan_vae as wan_vae + + calls = [] + + def fake_decode_latents_to_video(latents_np, output_path, **kwargs): + calls.append((latents_np, output_path, kwargs)) + output_path.write_bytes(b"fake mp4") + return {"backend": kwargs["backend"]} + + monkeypatch.setattr( + wan_vae, "decode_latents_to_video", fake_decode_latents_to_video + ) + latents = np.zeros((1, 16, 1, 2, 2), dtype=np.float32) + output_path = tmp_path / "out.mp4" + + module.decode_latents_to_video( + model_root=tmp_path, + latents_np=latents, + output_path=output_path, + fps=16, + device_arg="cpu", + dtype_arg="fp32", + backend="taehv", + taehv_source_path=None, + taehv_checkpoint_path=None, + taehv_parallel=False, + ) + + assert output_path.read_bytes() == b"fake mp4" + assert len(calls) == 1 + _, called_output_path, kwargs = calls[0] + assert called_output_path == output_path + assert kwargs["backend"] == "taehv" + assert kwargs["z_dim"] == 16 + + +def test_vendored_taehv_import_available(): + from fastvideo.third_party.taehv import TAEHV + + assert TAEHV.__name__ == "TAEHV" + + +def test_default_model_resolution_downloads_only_needed_transformer(monkeypatch, tmp_path): + module = _load_prompt_to_video_module() + import huggingface_hub + + calls = [] + + def fake_snapshot_download(model_id, allow_patterns): + calls.append((model_id, allow_patterns)) + return str(tmp_path) + + monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) + + assert module.resolve_model_root(None, include_transformer=True) == tmp_path + assert "transformer/*" in calls[-1][1] + assert module.resolve_model_root(None, include_transformer=False) == tmp_path + assert "transformer/*" not in calls[-1][1] + assert module.resolve_model_root(tmp_path) == tmp_path + assert len(calls) == 2 diff --git a/fastvideo/tests/mlx/test_mlx_refine.py b/fastvideo/tests/mlx/test_mlx_refine.py new file mode 100644 index 000000000..e0a727c0b --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_refine.py @@ -0,0 +1,248 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CPU-only contracts for the MLX two-pass refine pipeline. + +Mirrors the LTX-2 refine stage math (init resolution split, spatial +upsample, Gaussian re-noise) without requiring Metal / a real DiT. The +orchestration helpers that need ``mx`` are exercised only when MLX is +importable. +""" + +from __future__ import annotations + +import math + +import numpy as np +import pytest + +from examples.inference.basic.mlx_wan22_generate import DEFAULT_HEIGHT, DEFAULT_NUM_FRAMES, DEFAULT_WIDTH +from fastvideo.mlx_runtime.refine import ( + DEFAULT_REFINE_SIGMA, + default_refine_timesteps, + plan_refine_resolutions, + prepare_refine_latents, + refine_sigma_from_schedule, + upsample_latents_spatial, +) +from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, add_noise + + +def _fastwan_schedule() -> MLXDMDSchedule: + """FastWan's flow-match schedule: ``sigma == timestep / 1000``.""" + timesteps = np.arange(1000, 0, -1, dtype=np.float64) + return MLXDMDSchedule(sigmas=timesteps / 1000.0, timesteps=timesteps) + + +def test_plan_refine_resolutions_splits_even_target() -> None: + plan = plan_refine_resolutions( + height=480, + width=832, + num_frames=81, + spatial_scale=2, + vae_spatial_compression=8, + vae_temporal_compression=4, + patch_size=(1, 2, 2), + ) + assert plan.stage1_height == 240 + assert plan.stage1_width == 416 + assert plan.target_height == 480 + assert plan.target_width == 832 + assert plan.stage1_latent_height == 30 + assert plan.stage1_latent_width == 52 + assert plan.stage2_latent_height == 60 + assert plan.stage2_latent_width == 104 + assert plan.latent_frames == 21 + assert plan.spatial_scale == 2 + + +def test_plan_refine_disabled_collapses_to_single_pass() -> None: + plan = plan_refine_resolutions( + height=480, + width=832, + num_frames=81, + spatial_scale=2, + enabled=False, + ) + assert plan.spatial_scale == 1 + assert plan.stage1_height == plan.target_height == 480 + assert plan.stage1_width == plan.target_width == 832 + + +@pytest.mark.parametrize( + ("height", "width", "num_frames", "message"), + [(481, 832, 81, "height/width"), (480, 833, 81, "height/width"), (480, 832, 82, "num_frames")], +) +def test_plan_refine_rejects_requests_the_vae_would_truncate( + height: int, + width: int, + num_frames: int, + message: str, +) -> None: + with pytest.raises(ValueError, match=message): + plan_refine_resolutions(height=height, width=width, num_frames=num_frames, enabled=False) + + +def test_plan_refine_rejects_odd_target() -> None: + with pytest.raises(ValueError, match="divisible by"): + plan_refine_resolutions(height=481, width=832, num_frames=81, spatial_scale=2) + + +def test_plan_refine_rejects_vae_misaligned_stage1() -> None: + # 500/2 = 250, 250 % 8 != 0 → rejected (same guard as LTX2RefineInitStage). + with pytest.raises(ValueError, match="divisible by"): + plan_refine_resolutions( + height=500, + width=800, + num_frames=81, + spatial_scale=2, + vae_spatial_compression=8, + ) + + +def test_plan_refine_wan22_5b_720p_geometry() -> None: + # Wan2.2 TI2V-5B native 704x1280 with 16x spatial VAE compression. + plan = plan_refine_resolutions( + height=704, + width=1280, + num_frames=121, + spatial_scale=2, + vae_spatial_compression=16, + vae_temporal_compression=4, + patch_size=(1, 2, 2), + ) + assert plan.stage1_height == 352 + assert plan.stage1_width == 640 + assert plan.stage1_latent_height == 22 + assert plan.stage1_latent_width == 40 + assert plan.stage2_latent_height == 44 + assert plan.stage2_latent_width == 80 + assert plan.latent_frames == 31 + + +def test_plan_refine_wan22_default_geometry() -> None: + plan = plan_refine_resolutions( + height=DEFAULT_HEIGHT, + width=DEFAULT_WIDTH, + num_frames=DEFAULT_NUM_FRAMES, + spatial_scale=2, + vae_spatial_compression=16, + vae_temporal_compression=4, + patch_size=(1, 2, 2), + ) + assert (plan.stage1_latent_height, plan.stage1_latent_width) == (14, 26) + assert (plan.stage2_latent_height, plan.stage2_latent_width) == (28, 52) + assert plan.latent_frames == 31 + + +def test_upsample_nearest_repeats_spatial_samples() -> None: + rng = np.random.default_rng(0) + latents = rng.standard_normal((1, 4, 2, 3, 5)).astype(np.float32) + up = upsample_latents_spatial(latents, scale=2, mode="nearest") + assert up.shape == (1, 4, 2, 6, 10) + # Every 2x2 block equals the source sample broadcast across the block. + for y in range(3): + for x in range(5): + block = up[0, :, :, 2 * y:2 * y + 2, 2 * x:2 * x + 2] + src = latents[0, :, :, y, x][:, :, None, None] + np.testing.assert_array_equal(block, np.broadcast_to(src, block.shape)) + + +def test_upsample_bilinear_preserves_constants() -> None: + latents = np.full((1, 2, 1, 4, 4), 3.5, dtype=np.float32) + up = upsample_latents_spatial(latents, scale=2, mode="bilinear") + assert up.shape == (1, 2, 1, 8, 8) + np.testing.assert_allclose(up, 3.5, rtol=0, atol=1e-6) + + +def test_upsample_scale_one_is_identity() -> None: + rng = np.random.default_rng(1) + latents = rng.standard_normal((1, 2, 1, 4, 4)).astype(np.float32) + out = upsample_latents_spatial(latents, scale=1) + np.testing.assert_array_equal(out, latents) + + +def test_prepare_refine_latents_matches_add_noise_formula() -> None: + rng = np.random.default_rng(2) + clean = rng.standard_normal((1, 4, 2, 3, 5)).astype(np.float32) + noise = rng.standard_normal((1, 4, 2, 6, 10)).astype(np.float32) + sigma = 0.75 + + got = prepare_refine_latents( + clean, + scale=2, + sigma=sigma, + noise=noise, + add_noise_flag=True, + upsample_mode="nearest", + ) + up = upsample_latents_spatial(clean, scale=2, mode="nearest") + expected = add_noise(up, noise, sigma) + np.testing.assert_allclose(got, expected, rtol=0, atol=0) + + +def test_prepare_refine_latents_no_noise_returns_upsample() -> None: + rng = np.random.default_rng(3) + clean = rng.standard_normal((1, 2, 1, 4, 4)).astype(np.float32) + got = prepare_refine_latents(clean, scale=2, add_noise_flag=False, upsample_mode="nearest") + expected = upsample_latents_spatial(clean, scale=2, mode="nearest") + np.testing.assert_array_equal(got, expected) + + +def test_prepare_refine_latents_seed_is_deterministic() -> None: + rng = np.random.default_rng(4) + clean = rng.standard_normal((1, 2, 1, 3, 3)).astype(np.float32) + a = prepare_refine_latents(clean, scale=2, sigma=0.5, seed=123, upsample_mode="nearest") + b = prepare_refine_latents(clean, scale=2, sigma=0.5, seed=123, upsample_mode="nearest") + np.testing.assert_array_equal(a, b) + + +def test_refine_sigma_from_schedule_uses_first_timestep() -> None: + timesteps = np.arange(1000, dtype=np.float64) + sigmas = np.linspace(1.0, 1e-3, 1000, dtype=np.float64) + schedule = MLXDMDSchedule(sigmas=sigmas, timesteps=timesteps) + # First refine timestep drives the stage-2 start sigma. + assert refine_sigma_from_schedule(schedule, [999, 757, 522]) == pytest.approx(float(sigmas[999])) + assert refine_sigma_from_schedule(schedule, [522]) == pytest.approx(float(sigmas[522])) + + +def test_default_refine_sigma_matches_ltx2_stage2_head() -> None: + # Documented contract: DEFAULT_REFINE_SIGMA tracks LTX-2's first + # STAGE_2_DISTILLED_SIGMA_VALUES entry so the fallback is intentional. + assert math.isclose(DEFAULT_REFINE_SIGMA, 0.909375, rel_tol=0, abs_tol=0) + + +def test_full_noise_handoff_would_discard_stage1() -> None: + """Why the stage-2 grid may not open at the stage-1 timestep. + + FastWan's DMD grid starts at ``t=1000``, i.e. ``sigma == 1``. The hand-off + is ``(1 - sigma) * upsampled + sigma * noise``, so starting stage 2 there + weights the stage-1 result at exactly zero: refine silently degrades to a + plain full-resolution run costing two passes. + """ + schedule = _fastwan_schedule() + assert refine_sigma_from_schedule(schedule, [1000, 757, 522]) == pytest.approx(1.0) + + clean = np.random.default_rng(0).standard_normal((1, 2, 1, 4, 4)).astype(np.float32) + handoff = prepare_refine_latents(clean, scale=2, sigma=1.0, seed=7, upsample_mode="nearest") + upsampled = upsample_latents_spatial(clean, scale=2, mode="nearest") + # Nothing of stage 1 survives. + assert not np.allclose(handoff, upsampled) + + +def test_default_refine_timesteps_drops_the_full_noise_step() -> None: + schedule = _fastwan_schedule() + steps = default_refine_timesteps(schedule, [1000, 757, 522]) + assert steps == [757.0, 522.0] + # The resulting hand-off keeps a real share of stage 1. + sigma = refine_sigma_from_schedule(schedule, steps) + assert 0.0 < sigma < 1.0 + + +def test_default_refine_timesteps_keeps_an_already_valid_grid() -> None: + schedule = _fastwan_schedule() + assert default_refine_timesteps(schedule, [757, 522]) == [757.0, 522.0] + + +def test_default_refine_timesteps_rejects_an_all_full_noise_grid() -> None: + schedule = _fastwan_schedule() + with pytest.raises(ValueError, match="No usable refine timesteps"): + default_refine_timesteps(schedule, [1000]) diff --git a/fastvideo/tests/mlx/test_mlx_rife_interpolation.py b/fastvideo/tests/mlx/test_mlx_rife_interpolation.py new file mode 100644 index 000000000..5fe815fc7 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_rife_interpolation.py @@ -0,0 +1,63 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression tests for the Apple-Silicon MLX RIFE interpolation backend.""" + +from __future__ import annotations + +import numpy as np +import pytest +from huggingface_hub.utils import LocalEntryNotFoundError + +mx = pytest.importorskip("mlx.core", reason="MLX required for RIFE tests") + + +@pytest.mark.skipif( + not bool(getattr(mx, "metal", None) and mx.metal.is_available()), + reason="RIFE MLX regression requires Apple Silicon Metal", +) +def test_rife_interpolation_preserves_keyframes_shape_and_count() -> None: + from fastvideo.mlx_runtime.rife_interp import RIFEWeightsUnavailableError, interpolate, load_model + + frame0 = np.zeros((64, 96, 3), dtype=np.uint8) + frame1 = np.zeros((64, 96, 3), dtype=np.uint8) + frame1[:, :, 0] = 255 + try: + # The backend is vendored, but its weights are fetched from Hugging Face + # on first use, so an offline machine should skip rather than fail. + model = load_model() + except RIFEWeightsUnavailableError as exc: + pytest.skip(f"RIFE weights unavailable: {exc}") + + frames = interpolate([frame0, frame1], factor=2, model=model) + + assert len(frames) == 3 + assert frames[1].shape == frame0.shape + assert frames[1].dtype == np.uint8 + np.testing.assert_array_equal(frames[0], frame0) + np.testing.assert_array_equal(frames[-1], frame1) + + +def test_rife_download_unavailable_has_specific_error(monkeypatch) -> None: + from fastvideo.mlx_runtime.rife_interp import RIFEWeightsUnavailableError, load_model + from fastvideo.third_party.rife_mlx.utils import weights + + def unavailable(*args, **kwargs): + raise LocalEntryNotFoundError("offline") + + monkeypatch.setattr(weights, "build_model", unavailable) + load_model.cache_clear() + with pytest.raises(RIFEWeightsUnavailableError): + load_model() + + +def test_rife_backend_regression_is_not_skip_eligible(monkeypatch) -> None: + from fastvideo.mlx_runtime.rife_interp import RIFEBackendError, RIFEWeightsUnavailableError, load_model + from fastvideo.third_party.rife_mlx.utils import weights + + def broken(*args, **kwargs): + raise AssertionError("backend regression") + + monkeypatch.setattr(weights, "build_model", broken) + load_model.cache_clear() + with pytest.raises(RIFEBackendError) as caught: + load_model() + assert not isinstance(caught.value, RIFEWeightsUnavailableError) diff --git a/fastvideo/tests/mlx/test_mlx_taehv_parity.py b/fastvideo/tests/mlx/test_mlx_taehv_parity.py new file mode 100644 index 000000000..cd9baaa0a --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_taehv_parity.py @@ -0,0 +1,40 @@ +# SPDX-License-Identifier: Apache-2.0 +"""TAEHV decode parity tests for the MLX implementation. + +TAEHV MLX vs torch is bit-close (atol 1e-5) for z_dim=16 and 48. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core", reason="MLX required for Wan VAE/TAEHV tests") + +from fastvideo.mlx_runtime.wan_vae import ( # noqa: E402 + decode_latents_taehv_mlx, + ensure_taehv_checkpoint, +) + + +@pytest.mark.parametrize("z_dim", [16, 48]) +def test_taehv_mlx_matches_torch(z_dim: int) -> None: + import torch + from fastvideo.third_party.taehv import TAEHV + + ckpt = ensure_taehv_checkpoint(z_dim=z_dim) + rng = np.random.default_rng(0) + h = w = 8 + t = 5 + lat = (rng.standard_normal((1, z_dim, t, h, w)) * 0.5).astype(np.float32) + + model = TAEHV(str(ckpt)).eval() + with torch.no_grad(): + out_t = model.decode_video(torch.from_numpy(lat).transpose(1, 2), parallel=True, show_progress_bar=False) + torch_np = out_t[0].permute(0, 2, 3, 1).float().numpy() + mlx_np = decode_latents_taehv_mlx(lat, z_dim=z_dim)[0] + # Validate complete temporal output shape before comparing values. + assert mlx_np.shape[0] == torch_np.shape[0], ( + f"MLX output frame count {mlx_np.shape[0]} != torch {torch_np.shape[0]}" + ) + np.testing.assert_allclose(mlx_np, torch_np, atol=1e-5, rtol=1e-5) diff --git a/fastvideo/tests/mlx/test_mlx_wan22_memory_handoff.py b/fastvideo/tests/mlx/test_mlx_wan22_memory_handoff.py new file mode 100644 index 000000000..73abeef83 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_wan22_memory_handoff.py @@ -0,0 +1,116 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression test for the Wan2.2 MLX/MPS memory handoff.""" + +from __future__ import annotations + +import gc +import importlib.util +import json +import sys +import weakref +from pathlib import Path + +import numpy as np + + +def _load_wan22_module(): + root = Path(__file__).resolve().parents[3] + script = root / "examples/inference/basic/mlx_wan22_generate.py" + spec = importlib.util.spec_from_file_location("mlx_wan22_generate_memory_test", script) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_wan22_releases_compiled_dit_before_torch_decode_and_mlx_rife(monkeypatch, tmp_path: Path) -> None: + import mlx.core as mx + import torch + from diffusers import utils as diffusers_utils + + from examples.inference.basic import mlx_wan_prompt_to_video as shared + from fastvideo.mlx_runtime import wan22, wan22_sample, wan_vae + + module = _load_wan22_module() + checkpoint = tmp_path / "mlx" + checkpoint.mkdir() + (checkpoint / "mlx_dit.json").write_text( + json.dumps({"config": {"patch_size": [1, 2, 2], "in_channels": 1}}) + ) + + events: list[str] = [] + dit_ref = None + + class CyclicCompiledDiT: + def __init__(self) -> None: + self.compiled = self.forward + + def forward(self) -> None: + pass + + def fake_load(*args, **kwargs): + nonlocal dit_ref + assert kwargs["compile"] is True + dit = CyclicCompiledDiT() + dit_ref = weakref.ref(dit) + return dit + + def fake_decode(*args, **kwargs): + events.append("decode") + assert dit_ref is not None and dit_ref() is None + assert events[-3:] == ["gc", "mlx-clear", "decode"] + return np.zeros((1, 1, 2, 2, 3), dtype=np.float32) + + def fake_postprocess(**kwargs) -> None: + events.append("rife") + assert events[-4:] == ["decode", "gc", "mps-clear", "rife"] + + real_collect = gc.collect + + def tracked_collect() -> int: + events.append("gc") + return real_collect() + + monkeypatch.setattr(shared, "encode_prompt", lambda **kwargs: torch.zeros((1, 1, 1))) + monkeypatch.setattr(shared, "make_rotary_embeddings", lambda *args, **kwargs: None) + monkeypatch.setattr(shared, "_postprocess_video", fake_postprocess) + monkeypatch.setattr(wan22, "mlx_wan22_dit_from_mlx_checkpoint", fake_load) + monkeypatch.setattr(wan22_sample, "sample_wan22_dmd", lambda *args, **kwargs: args[2]) + monkeypatch.setattr(wan_vae, "decode_latents_wan_vae_torch", fake_decode) + monkeypatch.setattr(diffusers_utils, "export_to_video", lambda *args, **kwargs: None) + monkeypatch.setattr(gc, "collect", tracked_collect) + monkeypatch.setattr(mx, "clear_cache", lambda: events.append("mlx-clear")) + monkeypatch.setattr(torch.backends.mps, "is_available", lambda: True) + monkeypatch.setattr(torch.mps, "empty_cache", lambda: events.append("mps-clear")) + monkeypatch.setattr( + sys, + "argv", + [ + "mlx_wan22_generate.py", + "--text-encoder-root", + str(tmp_path / "text"), + "--no-prompt-cache", + "--mlx-checkpoint", + str(checkpoint), + "--height", + "32", + "--width", + "32", + "--num-frames", + "1", + "--dmd-denoising-steps", + "1", + "--compile", + "--fast", + "--decode-backend", + "wan-vae", + "--vae-root", + str(tmp_path / "vae"), + "--output-path", + str(tmp_path / "out.mp4"), + ], + ) + + module.main() + + assert events == ["gc", "mlx-clear", "decode", "gc", "mps-clear", "rife"] diff --git a/fastvideo/tests/mlx/test_mlx_wan22_parity.py b/fastvideo/tests/mlx/test_mlx_wan22_parity.py new file mode 100644 index 000000000..7ffad9f98 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_wan22_parity.py @@ -0,0 +1,142 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Track D Rung 2: per-token-timestep parity for the Wan2.2-TI2V-5B MLX port. + +The 5B (FullAttn) differs from the ported 2.1 only in scale (config) and +*per-token* timestep conditioning (``expand_timesteps``: timestep is ``[B, L]``, +one noise level per patch token — how TI2V keeps the image frame at t=0 while +video frames are noised). Passing a 2-D timestep drives the torch model's +per-token path (``ts_seq_len``), so this compares ``MLXWan22DiT`` against the +torch reference on a tiny random-weight config. This is the run-6 prereq gate; +runs on CPU in CI. Real 5B-weight parity is a separate Metal-gated test. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +mx = pytest.importorskip("mlx.core", reason="MLX is required for the Wan2.2 parity test") + +from fastvideo.forward_context import set_forward_context # noqa: E402 +from fastvideo.mlx_runtime.fastwan import mlx_block_weights_from_torch # noqa: E402 +from fastvideo.mlx_runtime.wan22 import ( # noqa: E402 + MLXWan22DiT, + MLXWan22TransformerBlock, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch # noqa: E402 +from fastvideo.tests.mlx.tiny_wan import ( # noqa: E402 + TINY_ARCH, + TOP_LEVEL_KEY_MAP, + build_hf_config, + build_tiny_wan_config, + build_torch_model, + mlx_rotary_embeddings, +) + + +def _mlx_wan22_from_torch(model, hf_config, *, compile: bool = False) -> MLXWan22DiT: + state = {name: value.detach().float() for name, value in model.state_dict().items()} + inner_dim = int(hf_config["num_attention_heads"]) * int(hf_config["attention_head_dim"]) + weights = {} + for mlx_name, torch_name in TOP_LEVEL_KEY_MAP.items(): + tensor = state[torch_name] + if mlx_name == "patch_embedding.weight": + tensor = tensor.reshape(inner_dim, -1) + weights[mlx_name] = mx.array(tensor.numpy()) + blocks = [ + MLXWan22TransformerBlock( + mlx_block_weights_from_torch(tb), dim=inner_dim, ffn_dim=int(hf_config["ffn_dim"]), + num_heads=int(hf_config["num_attention_heads"]), eps=float(hf_config["eps"])) for tb in model.blocks + ] + return MLXWan22DiT(weights, blocks, dict(hf_config), compile=compile) + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan22_per_token_timestep_matches_torch() -> None: + torch_model = build_torch_model() + hf_config = build_hf_config(build_tiny_wan_config()) + mlx_model = _mlx_wan22_from_torch(torch_model, hf_config) + + # Tiny latent: 4 frames, 8x8 latent, patch (1,2,2) -> 4*4*4 = 64 tokens (L). + frames, height, width = 4, 8, 8 + p_t, p_h, p_w = TINY_ARCH["patch_size"] + num_tokens = (frames // p_t) * (height // p_h) * (width // p_w) + tokens_per_frame = num_tokens // (frames // p_t) + + gen = torch.Generator().manual_seed(11) + hidden = torch.randn(1, TINY_ARCH["in_channels"], frames, height, width, generator=gen, dtype=torch.float32) + text = torch.randn(1, 8, TINY_ARCH["text_dim"], generator=gen, dtype=torch.float32) + + # Per-token timestep [B, L]: frame 0 clean (t=0, I2V-style), rest noised (t=500). + per_frame = [0] + [500] * (frames // p_t - 1) + timestep = torch.tensor([[per_frame[i // tokens_per_frame] for i in range(num_tokens)]], dtype=torch.long) + + with torch.no_grad(), set_forward_context( + current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")): + ref = torch_model(hidden_states=hidden, encoder_hidden_states=text, timestep=timestep).detach().float().numpy() + + freqs_cis = mlx_rotary_embeddings(hidden) + out = mlx_model(mx.array(hidden.numpy()), mx.array(text.numpy()), mx.array(timestep.float().numpy()), freqs_cis) + mx.eval(out) + mlx_out = np.array(out.astype(mx.float32)) + + np.testing.assert_allclose(mlx_out, ref, atol=2e-3, rtol=2e-3) + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan22_compile_path_matches_eager() -> None: + torch_model = build_torch_model() + hf_config = build_hf_config(build_tiny_wan_config()) + eager_model = _mlx_wan22_from_torch(torch_model, hf_config) + compiled_model = _mlx_wan22_from_torch(torch_model, hf_config, compile=True) + + frames, height, width = 4, 8, 8 + p_t, p_h, p_w = TINY_ARCH["patch_size"] + num_tokens = (frames // p_t) * (height // p_h) * (width // p_w) + tokens_per_frame = num_tokens // (frames // p_t) + + gen = torch.Generator().manual_seed(12) + hidden = torch.randn( + 1, + TINY_ARCH["in_channels"], + frames, + height, + width, + generator=gen, + dtype=torch.float32, + ) + text = torch.randn( + 1, + 8, + TINY_ARCH["text_dim"], + generator=gen, + dtype=torch.float32, + ) + per_frame = [0] + [500] * (frames // p_t - 1) + timestep = torch.tensor( + [[per_frame[i // tokens_per_frame] for i in range(num_tokens)]], + dtype=torch.long, + ) + freqs_cis = mlx_rotary_embeddings(hidden) + + eager = eager_model( + mx.array(hidden.numpy()), + mx.array(text.numpy()), + mx.array(timestep.float().numpy()), + freqs_cis, + ) + compiled = compiled_model( + mx.array(hidden.numpy()), + mx.array(text.numpy()), + mx.array(timestep.float().numpy()), + freqs_cis, + ) + mx.eval(eager, compiled) + + np.testing.assert_allclose( + np.array(compiled.astype(mx.float32)), + np.array(eager.astype(mx.float32)), + atol=1e-4, + rtol=1e-4, + ) diff --git a/fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py b/fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py new file mode 100644 index 000000000..1a2e5495f --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py @@ -0,0 +1,318 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan2.2 prompt-embedding cache fingerprint regression tests.""" + +from __future__ import annotations + +import importlib.util +import json +import os +import sys +from pathlib import Path + +import numpy as np +import pytest + +from fastvideo.mlx_runtime.prompt_cache import load_prompt_cache, save_prompt_cache + + +def _load_wan22_module(): + root = Path(__file__).resolve().parents[3] + script = root / "examples/inference/basic/mlx_wan22_generate.py" + spec = importlib.util.spec_from_file_location( + "mlx_wan22_generate_for_test", script + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def _load_wan21_module(): + root = Path(__file__).resolve().parents[3] + script = root / "examples/inference/basic/mlx_wan_prompt_to_video.py" + spec = importlib.util.spec_from_file_location( + "mlx_wan_prompt_to_video_for_cache_test", script + ) + assert spec is not None and spec.loader is not None + module = importlib.util.module_from_spec(spec) + spec.loader.exec_module(module) + return module + + +def test_wan22_prompt_cache_rejects_mismatched_fingerprint(tmp_path: Path) -> None: + module = _load_wan22_module() + cache_path = tmp_path / "prompt.npy" + text_encoder_root = tmp_path / "encoder" + text_encoder_root.mkdir() + + fingerprint = module._prompt_cache_fingerprint( + prompt="a fox", + prompt_used="a fox, cinematic", + enhance_prompt=True, + enhance_prompt_backend="template", + text_encoder_root=text_encoder_root, + max_sequence_length=512, + dtype="fp16", + ) + embeds = np.ones((1, 512, 4), dtype=np.float32) + + save_prompt_cache(cache_path, embeds, fingerprint) + + assert load_prompt_cache(cache_path, fingerprint) is not None + + changed = dict(fingerprint) + changed["prompt"] = "a cat" + assert load_prompt_cache(cache_path, changed) is None + + +def test_wan22_prompt_cache_rejects_missing_metadata(tmp_path: Path) -> None: + module = _load_wan22_module() + cache_path = tmp_path / "prompt.npy" + np.save(cache_path, np.zeros((1, 512, 4), dtype=np.float32)) + + fingerprint = module._prompt_cache_fingerprint( + prompt="a fox", + prompt_used="a fox", + enhance_prompt=False, + enhance_prompt_backend="template", + text_encoder_root=tmp_path, + max_sequence_length=512, + dtype="fp16", + ) + + assert load_prompt_cache(cache_path, fingerprint) is None + + +def test_prompt_cache_is_best_effort_and_rejects_torn_data(tmp_path: Path) -> None: + cache_path = tmp_path / "prompt.npy" + fingerprint = {"prompt": "a fox"} + save_prompt_cache(cache_path, np.ones((1, 1, 1)), fingerprint) + + np.save(cache_path, np.full((1, 1, 1), 9.0)) + assert load_prompt_cache(cache_path, fingerprint) is None + + invalid_parent = tmp_path / "not-a-directory" + invalid_parent.write_text("occupied") + assert not save_prompt_cache( + invalid_parent / "prompt.npy", + np.ones((1, 1, 1)), + fingerprint, + ) + + incomplete_path = tmp_path / "incomplete.npy" + incomplete = {"text_encoder": {"complete": False}} + assert not save_prompt_cache(incomplete_path, np.ones((1, 1, 1)), incomplete) + assert not incomplete_path.exists() + + +def test_wan21_explicit_cache_reencodes_for_a_changed_prompt(monkeypatch, tmp_path: Path) -> None: + import torch + + module = _load_wan21_module() + calls = [] + + def fake_encode_prompt(**kwargs): + calls.append(kwargs["prompt"]) + return torch.tensor([[float(len(calls))]]) + + monkeypatch.setattr(module, "encode_prompt", fake_encode_prompt) + cache_path = tmp_path / "prompt.npy" + kwargs = { + "model_root": tmp_path / "model", + "max_sequence_length": 512, + "device_arg": "cpu", + "dtype_arg": "fp16", + "encode_mode": "inline", + "cache_path": cache_path, + } + first = module.get_prompt_embeds(prompt="first", **kwargs) + second = module.get_prompt_embeds(prompt="second", **kwargs) + second_cached = module.get_prompt_embeds(prompt="second", **kwargs) + + assert calls == ["first", "second"] + assert not torch.equal(first, second) + assert torch.equal(second, second_cached) + + +def test_wan22_missing_paths_download_only_selected_assets(monkeypatch, tmp_path: Path) -> None: + module = _load_wan22_module() + import huggingface_hub + + calls = [] + + def fake_snapshot_download(model_id, allow_patterns): + calls.append((model_id, allow_patterns)) + return str(tmp_path / ("wan21" if "2.1" in model_id else "wan22")) + + monkeypatch.setattr(huggingface_hub, "snapshot_download", fake_snapshot_download) + text_root, dit_checkpoint, dit_config, vae_root = module._resolve_model_paths( + text_encoder_root=None, + dit_checkpoint=None, + dit_config=None, + vae_root=None, + mlx_checkpoint=None, + decode_backend="taehv", + ) + + assert text_root == tmp_path / "wan21" + assert dit_checkpoint == tmp_path / "wan22/transformer/diffusion_pytorch_model.safetensors" + assert dit_config == tmp_path / "wan22/transformer/config.json" + assert vae_root is None + assert calls == [ + (module.FASTWAN21_MODEL_ID, ["tokenizer/*", "text_encoder/*"]), + ( + module.FASTWAN22_MODEL_ID, + ["transformer/diffusion_pytorch_model.safetensors", "transformer/config.json"], + ), + ] + + +def test_wan22_explicit_mlx_paths_do_not_download(monkeypatch, tmp_path: Path) -> None: + module = _load_wan22_module() + import huggingface_hub + + def unexpected_download(*args, **kwargs): + pytest.fail("explicit MLX and text paths must not download model assets") + + monkeypatch.setattr(huggingface_hub, "snapshot_download", unexpected_download) + text_root, dit_checkpoint, dit_config, vae_root = module._resolve_model_paths( + text_encoder_root=tmp_path / "text", + dit_checkpoint=None, + dit_config=None, + vae_root=None, + mlx_checkpoint=tmp_path / "mlx", + decode_backend="taehv", + ) + + assert text_root == tmp_path / "text" + assert dit_checkpoint is None + assert dit_config is None + assert vae_root is None + + +def test_wan22_prompt_cache_has_a_default_path(tmp_path: Path) -> None: + """The cache must work without an explicit --prompt-embeds-cache. + + It previously only engaged when handed a path, so every 5B run paid a full + UMT5 encode (~45s on an M4 Max) even for a repeat prompt, while the Wan2.1 + entrypoint cached by default. + """ + module = _load_wan22_module() + fingerprint = module._prompt_cache_fingerprint( + prompt="a fox", + prompt_used="a fox", + enhance_prompt=False, + enhance_prompt_backend="none", + text_encoder_root=tmp_path, + max_sequence_length=512, + dtype="fp16", + ) + path = module._default_prompt_cache_path(fingerprint) + assert path.suffix == ".npy" + assert path.parent.name == "prompt_embeds" + assert path.name.startswith("wan22_") + + +def test_wan22_default_cache_path_tracks_the_fingerprint(tmp_path: Path) -> None: + """Two prompts must not collide, and the same prompt must be stable.""" + module = _load_wan22_module() + + def fp(prompt: str): + return module._prompt_cache_fingerprint( + prompt=prompt, + prompt_used=prompt, + enhance_prompt=False, + enhance_prompt_backend="none", + text_encoder_root=tmp_path, + max_sequence_length=512, + dtype="fp16", + ) + + assert module._default_prompt_cache_path(fp("a fox")) == module._default_prompt_cache_path(fp("a fox")) + assert module._default_prompt_cache_path(fp("a fox")) != module._default_prompt_cache_path(fp("a cat")) + + +def test_wan22_default_cache_path_tracks_encoder_files(tmp_path: Path) -> None: + module = _load_wan22_module() + weights = tmp_path / "text_encoder/weights.safetensors" + weights.parent.mkdir() + weights.write_bytes(b"one") + + def path(): + fingerprint = module._prompt_cache_fingerprint( + prompt="a fox", + prompt_used="a fox", + enhance_prompt=False, + enhance_prompt_backend="none", + text_encoder_root=tmp_path, + max_sequence_length=512, + dtype="fp16", + ) + return module._default_prompt_cache_path(fingerprint) + + before = path() + weights.write_bytes(b"two") + stat = weights.stat() + os.utime(weights, ns=(stat.st_atime_ns, stat.st_mtime_ns + 1)) + assert path() != before + + +@pytest.mark.parametrize( + ("cache_args", "expected"), + [ + ([], "default"), + (["--no-prompt-cache"], None), + (["--prompt-embeds-cache", "explicit.npy"], "explicit"), + (["--no-prompt-cache", "--prompt-embeds-cache", "explicit.npy"], "explicit"), + ], +) +def test_wan22_main_resolves_prompt_cache( + monkeypatch, + tmp_path: Path, + cache_args: list[str], + expected: str | None, +) -> None: + module = _load_wan22_module() + checkpoint = tmp_path / "mlx" + checkpoint.mkdir() + (checkpoint / "mlx_dit.json").write_text( + json.dumps({"config": {"patch_size": [1, 2, 2], "in_channels": 1}}) + ) + explicit = tmp_path / "explicit.npy" + + class CacheProbe(Exception): + pass + + def probe(cache_path, fingerprint): + if expected == "default": + assert cache_path == module._default_prompt_cache_path(fingerprint) + elif expected == "explicit": + assert cache_path == explicit + else: + assert cache_path is None + raise CacheProbe + + resolved_args = [str(explicit) if arg == "explicit.npy" else arg for arg in cache_args] + monkeypatch.setenv("HOME", str(tmp_path)) + monkeypatch.setattr(module, "load_prompt_cache", probe) + monkeypatch.setattr( + sys, + "argv", + [ + "mlx_wan22_generate.py", + "--text-encoder-root", + str(tmp_path / "text"), + "--mlx-checkpoint", + str(checkpoint), + "--height", + "32", + "--width", + "32", + "--num-frames", + "1", + *resolved_args, + ], + ) + + with pytest.raises(CacheProbe): + module.main() diff --git a/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py b/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py new file mode 100644 index 000000000..e53715900 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_wan22_real_weights_parity.py @@ -0,0 +1,222 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Track D Rung 3: real-weight T2V parity for Wan2.2-TI2V-5B MLX vs torch. + +Loads the released ``FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers`` transformer +into both ``MLXWan22DiT`` and torch ``WanTransformer3DModel``, runs one forward +with a 2-D per-token timestep (frame 0 at t=0 I2V-style, remaining tokens at +t≈900), and asserts allclose. Gated on Metal + local weights so Linux +``mlx[cpu]`` CI stays green. +""" + +from __future__ import annotations + +import json +import os +from pathlib import Path + +import numpy as np +import pytest + +# The torch reference side brings up FastVideo's single-process distributed +# environment, which rendezvouses over env://. Match the other torch-backed +# tests in this repo and provide a local rendezvous before that happens. +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29513") + +mx = pytest.importorskip("mlx.core", reason="MLX is required for the Wan2.2 real-weight parity test") + +_ROOT = Path(os.environ.get("FASTVIDEO_WAN22_5B_ROOT", str(Path.home() / "models" / "fastwan22_5b"))) +_CHECKPOINT = _ROOT / "transformer" / "diffusion_pytorch_model.safetensors" +_CONFIG = _ROOT / "transformer" / "config.json" + +_HAS_METAL = bool(getattr(mx, "metal", None) and mx.metal.is_available()) +_HAS_WEIGHTS = _CHECKPOINT.exists() and _CONFIG.exists() and _CHECKPOINT.stat().st_size > 1_000_000_000 + +# The real-weight test loads ~10 GB fp16 weights + activations. Gate on the +# documented 36 GB requirement unless explicitly opted in. +def _check_memory_requirement() -> bool: + """Determine whether the system meets the memory requirement for the test. + + Returns: + bool: `true` if the system has at least 36 GiB of memory or low-memory + execution is explicitly enabled, `false` otherwise. + """ + if os.environ.get("FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY") == "1": + return True + try: + # On macOS, check unified memory via sysctl. + import subprocess + result = subprocess.run( + ["sysctl", "-n", "hw.memsize"], + capture_output=True, + text=True, + check=False, + timeout=5, + ) + if result.returncode == 0: + mem_bytes = int(result.stdout.strip()) + mem_gib = mem_bytes / (1024**3) + return mem_gib >= 36.0 + except Exception: + pass + # Conservative fallback: require opt-in if we can't detect memory. + return False + + +_HAS_SUFFICIENT_MEMORY = _check_memory_requirement() + +pytestmark = [ + pytest.mark.skipif(not _HAS_METAL, reason="Metal required for real 5B fp16 parity"), + pytest.mark.skipif( + not _HAS_WEIGHTS, + reason=f"Wan2.2-5B checkpoint not found/incomplete under {_ROOT} " + "(set FASTVIDEO_WAN22_5B_ROOT; expected ~10 GB safetensors)", + ), + pytest.mark.skipif( + not _HAS_SUFFICIENT_MEMORY, + reason="Real-weight 5B test requires >= 36 GB unified memory " + "(set FASTVIDEO_WAN22_5B_ALLOW_LOW_MEMORY=1 to override)", + ), +] + + +def _load_torch_wan22_from_diffusers(checkpoint: Path, config_path: Path, *, dtype): + """ + Instantiate a Wan2.2 transformer and load its Diffusers-format checkpoint weights. + + Parameters: + checkpoint (Path): Path to the safetensors checkpoint. + config_path (Path): Path to the Diffusers configuration file. + dtype: Data type to use for the returned model. + + Returns: + WanTransformer3DModel: An evaluation-mode model with the loaded weights. + """ + import torch + from safetensors.torch import load_file + + from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig + from fastvideo.models.dits.wanvideo import WanTransformer3DModel + from fastvideo.models.loader.utils import get_param_names_mapping, hf_to_custom_state_dict + + hf_config = json.loads(config_path.read_text()) + # Diffusers-only keys that are not arch fields. + arch_kwargs = { + k: v + for k, v in hf_config.items() + if k not in {"_class_name", "_name_or_path", "_diffusers_version"} + } + cfg = WanVideoConfig(arch_config=WanVideoArchConfig(**{ + "num_attention_heads": int(arch_kwargs["num_attention_heads"]), + "attention_head_dim": int(arch_kwargs["attention_head_dim"]), + "in_channels": int(arch_kwargs["in_channels"]), + "out_channels": int(arch_kwargs["out_channels"]), + "text_dim": int(arch_kwargs["text_dim"]), + "freq_dim": int(arch_kwargs["freq_dim"]), + "ffn_dim": int(arch_kwargs["ffn_dim"]), + "num_layers": int(arch_kwargs["num_layers"]), + "patch_size": tuple(arch_kwargs["patch_size"]), + "cross_attn_norm": bool(arch_kwargs.get("cross_attn_norm", True)), + "qk_norm": arch_kwargs.get("qk_norm", "rms_norm_across_heads"), + "eps": float(arch_kwargs.get("eps", 1e-6)), + "rope_max_seq_len": int(arch_kwargs.get("rope_max_seq_len", 1024)), + "added_kv_proj_dim": arch_kwargs.get("added_kv_proj_dim"), + "image_dim": arch_kwargs.get("image_dim"), + "pos_embed_seq_len": arch_kwargs.get("pos_embed_seq_len"), + })) + model = WanTransformer3DModel(config=cfg, hf_config=hf_config).eval() + raw = load_file(str(checkpoint), device="cpu") + mapping = get_param_names_mapping(model.param_names_mapping) + custom_sd, _ = hf_to_custom_state_dict(raw.items(), mapping) + missing, unexpected = model.load_state_dict(custom_sd, strict=False) + # Diffusers layout should map cleanly; unexpected keys are always an error. + assert not unexpected, f"unexpected keys: {unexpected[:8]}" + # Missing keys must be either empty buffers or explicitly approved. Reject + # missing parameters to prevent parity checks from running with uninitialized weights. + if missing: + # Allowlist for known empty buffers (if any exist for Wan2.2-5B). + # Update this list if legitimate empty buffers are identified. + allowed_missing = set() + unapproved = [k for k in missing if k not in allowed_missing] + assert not unapproved, ( + f"Missing keys in Wan2.2-5B state dict load: {unapproved[:8]}. " + "Cannot proceed with parity check; uninitialized parameters would " + "invalidate results." + ) + return model.to(dtype=dtype) + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan22_real_weights_mlx_matches_torch_per_token_timestep() -> None: + """ + Verify that the MLX and PyTorch Wan2.2 implementations produce equivalent outputs for per-token timesteps using real checkpoint weights. + """ + import torch + + from fastvideo.forward_context import set_forward_context + from fastvideo.mlx_runtime.wan22 import mlx_wan22_dit_from_diffusers_safetensors + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + from examples.inference.basic.mlx_wan_prompt_to_video import make_rotary_embeddings + + config = json.loads(_CONFIG.read_text()) + in_ch = int(config["in_channels"]) + text_dim = int(config["text_dim"]) + # Small latent so activations fit alongside the ~10 GB fp16 weights on 36 GB. + frames, height, width = 2, 16, 16 + p_t, p_h, p_w = tuple(config["patch_size"]) + tokens_per_frame = (height // p_h) * (width // p_w) + num_tokens = (frames // p_t) * tokens_per_frame + + # Frame 0 clean (t=0), remaining tokens noised (t=900) — I2V-style expand_timesteps. + per_frame_levels = [0] + [900] * (frames // p_t - 1) + timestep_1d = [per_frame_levels[i // tokens_per_frame] for i in range(num_tokens)] + timestep = torch.tensor([timestep_1d], dtype=torch.long) + + rng = np.random.default_rng(2026) + hidden_np = (rng.standard_normal((1, in_ch, frames, height, width)) * 0.5).astype(np.float32) + text_np = (rng.standard_normal((1, 32, text_dim)) * 0.1).astype(np.float32) + + # --- torch reference (CPU fp16, same weight dtype as MLX deploy path) --- + torch_model = _load_torch_wan22_from_diffusers(_CHECKPOINT, _CONFIG, dtype=torch.float16) + hidden_t = torch.from_numpy(hidden_np).to(torch.float16) + text_t = torch.from_numpy(text_np).to(torch.float16) + with torch.no_grad(), set_forward_context( + current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")): + ref = torch_model(hidden_states=hidden_t, encoder_hidden_states=text_t, timestep=timestep) + ref_np = ref.detach().float().cpu().numpy() + + # Free torch weights before loading MLX so ~10 GB fp16 + activations fit in 36 GB. + del torch_model + import gc + gc.collect() + + # --- MLX (fp16 weights, fp16 compute) --- + mlx_model = mlx_wan22_dit_from_diffusers_safetensors(_CHECKPOINT, _CONFIG, dtype="fp16") + freqs_cis = make_rotary_embeddings( + config, latent_frames=frames, latent_height=height, latent_width=width) + out = mlx_model( + mx.array(hidden_np).astype(mx.float16), + mx.array(text_np).astype(mx.float16), + mx.array(np.array(timestep_1d, dtype=np.float32)[None, :]), + freqs_cis, + ) + mx.eval(out) + mlx_np = np.array(out.astype(mx.float32)) + + assert mlx_np.shape == ref_np.shape + assert np.isfinite(mlx_np).all() + max_abs = float(np.abs(mlx_np - ref_np).max()) + mean_abs = float(np.abs(mlx_np - ref_np).mean()) + flat_a, flat_b = mlx_np.ravel(), ref_np.ravel() + cosine = float(np.dot(flat_a, flat_b) / (np.linalg.norm(flat_a) * np.linalg.norm(flat_b) + 1e-12)) + print(f"Wan2.2-5B real-weight MLX-fp16 vs torch-fp16: " + f"max|Δ|={max_abs:.3e} mean|Δ|={mean_abs:.3e} cosine={cosine:.6f}") + # Full 30-layer fp16 Metal vs torch SDPA drifts more than the tiny-config + # fp32 gate (2e-3) / CUDA-tiny dump (5e-3). Measured on M4 Max: mean 6.1e-3, + # cosine 0.999948, both stable across MLX versions. The worst single element + # is not: it was 9e-2 when this budget was written and reaches 1.65e-1 under + # MLX 0.32's kernel fusion, so the aggregate statistics carry the assertion + # and max|Δ| stays as a loose guard against a genuine divergence. + assert mean_abs <= 1e-2, f"mean|Δ|={mean_abs} over budget" + assert cosine >= 0.9995, f"cosine={cosine} too low" + assert max_abs <= 0.25, f"max|Δ|={max_abs} over budget" diff --git a/fastvideo/tests/mlx/test_mlx_wan_vae_denormalization.py b/fastvideo/tests/mlx/test_mlx_wan_vae_denormalization.py new file mode 100644 index 000000000..42bd3738d --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_wan_vae_denormalization.py @@ -0,0 +1,60 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan VAE latent denormalization regression tests.""" + +from __future__ import annotations + +from pathlib import Path + +import numpy as np +import pytest + +from fastvideo.mlx_runtime.wan_vae import ( + WanVAEConfigView, + denormalize_latents_np, +) + + +def test_denormalize_matches_prompt_to_video_formula() -> None: + """``latents / (1/std) + mean`` used by the hybrid script.""" + cfg = WanVAEConfigView( + z_dim=2, + latents_mean=(0.1, -0.2), + latents_std=(2.0, 0.5), + ) + lat = np.ones((1, 2, 1, 1, 1), dtype=np.float32) + out = denormalize_latents_np(lat, cfg) + # z * std + mean + np.testing.assert_allclose(out[0, 0, 0, 0, 0], 1.0 * 2.0 + 0.1) + np.testing.assert_allclose(out[0, 1, 0, 0, 0], 1.0 * 0.5 + (-0.2)) + + +def test_full_vae_config_and_denormalize_formula() -> None: + """Config load + denormalize formula vs the hybrid script's tensor math.""" + import torch + + snaps = list( + Path.home() + .joinpath( + ".cache/huggingface/hub/" + "models--FastVideo--FastWan2.2-TI2V-5B-FullAttn-Diffusers/" + "snapshots" + ) + .glob("*/vae") + ) + if not snaps: + pytest.skip("Wan2.2 VAE not in HF cache") + vae_dir = snaps[0] + cfg = WanVAEConfigView.from_vae_dir(vae_dir) + assert cfg.z_dim == 48 + rng = np.random.default_rng(1) + lat = (rng.standard_normal((1, 48, 3, 8, 8)) * 0.3).astype(np.float32) + lat_dn = denormalize_latents_np(lat, cfg) + mean = torch.tensor(cfg.latents_mean, dtype=torch.float32).view( + 1, -1, 1, 1, 1 + ) + inv_std = (1.0 / torch.tensor(cfg.latents_std, dtype=torch.float32)).view( + 1, -1, 1, 1, 1 + ) + lat_t = torch.from_numpy(lat) + ref = (lat_t / inv_std + mean).numpy() + np.testing.assert_allclose(lat_dn, ref, atol=1e-5, rtol=1e-5) diff --git a/fastvideo/tests/mlx/test_quant_backends.py b/fastvideo/tests/mlx/test_quant_backends.py new file mode 100644 index 000000000..814d5f634 --- /dev/null +++ b/fastvideo/tests/mlx/test_quant_backends.py @@ -0,0 +1,126 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Accuracy + capability probe for MLX block-scaled quant backends. + +For each backend in :mod:`fastvideo.mlx_runtime.quant_backends`, if the installed +MLX build supports the mode, assert that quantized matmul stays within a +per-format relative Frobenius error budget vs. the fp16 reference ``x @ w.T``. +Unsupported modes are skipped with an explicit reason. + +Run with ``-s`` to print the summary table (backend | supported | bytes/weight | +rel-error). +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("mlx.core", reason="MLX is required for quant backend tests") + +import mlx.core as mx # noqa: E402 + +from fastvideo.mlx_runtime.quant_backends import ( # noqa: E402 + BACKENDS, + bytes_per_weight, + is_supported, + quantize_weight, + quantized_matmul, + support_error, +) + +# Fixed shapes: weight cols must be divisible by all group sizes (64/32/16). +_WEIGHT_ROWS = 1024 +_WEIGHT_COLS = 1024 +_BATCH = 8 +_SEED = 0 + +# Relative Frobenius error ||y - ref||_F / ||ref||_F tolerances. +# Measured once on mlx 0.31.2 / Apple Silicon for N(0,1) fp16 weights: +# affine_int8_g64 ~ 5e-3, mxfp8 ~ 8e-2, mxfp4 ~ 1.2e-1, nvfp4 ~ 1.0e-1. +# Budgets keep ~2x headroom for RNG / platform variance. +_REL_ERROR_TOL: dict[str, float] = { + "affine_int8_g64": 2e-2, + "mxfp8": 1.5e-1, + "mxfp4": 2.5e-1, + "nvfp4": 2.5e-1, +} + + +def _frobenius_rel_error(y: mx.array, ref: mx.array) -> float: + """||y - ref||_F / ||ref||_F in float32 (avoids fp16 sum overflow).""" + y32 = y.astype(mx.float32) + ref32 = ref.astype(mx.float32) + diff = y32 - ref32 + num = mx.sqrt(mx.sum(diff * diff)) + den = mx.sqrt(mx.sum(ref32 * ref32)) + mx.array(1e-12, dtype=mx.float32) + mx.eval(num, den) + return float(num / den) + + +def _make_inputs() -> tuple[mx.array, mx.array, mx.array]: + """ + Generate deterministic random weights and inputs with their matrix multiplication reference. + + Returns: + tuple[mx.array, mx.array, mx.array]: The weights, inputs, and fp16 reference output. + """ + mx.random.seed(_SEED) + w = mx.random.normal((_WEIGHT_ROWS, _WEIGHT_COLS)).astype(mx.float16) + x = mx.random.normal((_BATCH, _WEIGHT_COLS)).astype(mx.float16) + ref = x @ w.T + mx.eval(w, x, ref) + return w, x, ref + + +@pytest.mark.parametrize("backend", list(BACKENDS)) +def test_quant_backend_matmul_accuracy(backend: str) -> None: + if not is_supported(backend): + err = support_error(backend) or "unknown error" + pytest.skip(reason=f"{backend} not supported by installed MLX: {err}") + + w, x, ref = _make_inputs() + qw = quantize_weight(w, backend) + y = quantized_matmul(x, qw) + mx.eval(y) + + rel = _frobenius_rel_error(y, ref) + tol = _REL_ERROR_TOL[backend] + assert rel < tol, ( + f"{backend}: relative Frobenius error {rel:.6e} exceeds tolerance {tol:.6e}" + ) + + +def test_quant_backends_summary_table(capsys: pytest.CaptureFixture[str]) -> None: + """ + Print a support, storage, and relative-error summary for each quantization backend. + + The table reports unsupported backends as unavailable and asserts that the affine + int8 baseline is supported. + """ + w, x, ref = _make_inputs() + rows: list[tuple[str, str, str, str]] = [] + + for backend in BACKENDS: + supported = is_supported(backend) + if not supported: + rows.append((backend, "no", "n/a", "n/a")) + continue + + bpw = bytes_per_weight(backend) + qw = quantize_weight(w, backend) + y = quantized_matmul(x, qw) + mx.eval(y) + rel = _frobenius_rel_error(y, ref) + rows.append((backend, "yes", f"{bpw:.6f}", f"{rel:.6e}")) + + # Header + rows for human inspection under pytest -s. + col_w = max(len(r[0]) for r in rows) + header = f"{'backend':<{col_w}} supported bytes/weight rel-error" + print("\n" + header) + print("-" * len(header)) + for backend, supported, bpw, rel in rows: + print(f"{backend:<{col_w}} {supported:<9} {bpw:>12} {rel}") + + # At least the affine baseline must work on any supported MLX install. + assert is_supported("affine_int8_g64"), ( + "affine_int8_g64 must be supported; cannot evaluate other backends" + ) diff --git a/fastvideo/tests/mlx/test_quant_capability.py b/fastvideo/tests/mlx/test_quant_capability.py new file mode 100644 index 000000000..ab8e014e1 --- /dev/null +++ b/fastvideo/tests/mlx/test_quant_capability.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MLX quantization-mode capability detection. + +The affine int8/int4 modes work on every MLX build the repo supports; the +mxfp8/mxfp4/nvfp4 mode strings only exist on newer MLX builds. These tests pin +the contract: supported modes probe clean, and unsupported ones surface as +``UnsupportedMLXQuantizationError`` with an actionable message *before* any +model weights are loaded. + +Runs on any MLX backend (Metal or ``mlx[cpu]``). +""" + +from __future__ import annotations + +import pytest + +pytest.importorskip("mlx.core", reason="MLX is required for quantization capability tests") + +from fastvideo.mlx_runtime.fastwan import ( # noqa: E402 + MLXQuantizationSpec, + UnsupportedMLXQuantizationError, + _QUANT_SUPPORT_CACHE, + ensure_quantization_supported, + quantization_support_error, +) + + +def test_affine_modes_are_supported_on_any_mlx_build() -> None: + for name in ("int8", "int4"): + spec = MLXQuantizationSpec.from_name(name) + assert quantization_support_error(spec) is None, f"{name} must be supported on every MLX build" + ensure_quantization_supported(spec) # must not raise + + +def test_none_spec_is_always_supported() -> None: + ensure_quantization_supported(None) # must not raise + + +def test_probe_result_is_cached_per_spec() -> None: + spec = MLXQuantizationSpec.from_name("int8") + quantization_support_error(spec) + key = (spec.mode, spec.bits, spec.group_size) + assert key in _QUANT_SUPPORT_CACHE + + +def test_unsupported_mode_raises_with_actionable_message(monkeypatch) -> None: + spec = MLXQuantizationSpec.from_name("nvfp4") + key = (spec.mode, spec.bits, spec.group_size) + monkeypatch.setitem(_QUANT_SUPPORT_CACHE, key, "ValueError: [quantize] Unknown mode 'nvfp4'") + + with pytest.raises(UnsupportedMLXQuantizationError) as excinfo: + ensure_quantization_supported(spec) + message = str(excinfo.value) + assert "nvfp4" in message + assert "int8" in message # points the user at the reliable fallback + + +def test_loader_rejects_unsupported_mode_before_reading_weights(monkeypatch, tmp_path) -> None: + """Verify that unsupported quantization modes are rejected before checkpoint weights are read.""" + from fastvideo.mlx_runtime.fastwan import mlx_dit_from_diffusers_safetensors + + spec_key = ("mxfp4", None, None) + monkeypatch.setitem(_QUANT_SUPPORT_CACHE, spec_key, "ValueError: [quantize] Unknown mode 'mxfp4'") + + config_path = tmp_path / "config.json" + config_path.write_text('{"num_layers": 1, "num_attention_heads": 4, "attention_head_dim": 16, ' + '"ffn_dim": 128, "in_channels": 16, "out_channels": 16, ' + '"patch_size": [1, 2, 2], "freq_dim": 64, "eps": 1e-6}') + # The checkpoint path deliberately does not exist: the capability check + # must fire before the loader ever opens the safetensors file. + with pytest.raises(UnsupportedMLXQuantizationError): + mlx_dit_from_diffusers_safetensors(tmp_path / "missing.safetensors", config_path, quantization="mxfp4") diff --git a/fastvideo/tests/mlx/test_taehv_decode.py b/fastvideo/tests/mlx/test_taehv_decode.py new file mode 100644 index 000000000..3f28da28e --- /dev/null +++ b/fastvideo/tests/mlx/test_taehv_decode.py @@ -0,0 +1,34 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Offline contracts for the vendored TAEHV decoder helper.""" + +import re + +import pytest + +from fastvideo.mlx_runtime import wan_vae +from fastvideo.mlx_runtime.taehv_decode import TAEW2_1_CHECKPOINT_SHA256, ensure_taew2_1_checkpoint + + +def test_taehv_checkpoint_pin_is_a_sha256() -> None: + assert len(TAEW2_1_CHECKPOINT_SHA256) == 64 + assert int(TAEW2_1_CHECKPOINT_SHA256, 16) >= 0 + + +def test_explicit_taehv_checkpoint_path_is_not_downloaded(tmp_path) -> None: + checkpoint = tmp_path / "custom-taehv.pth" + checkpoint.write_bytes(b"local test checkpoint") + assert ensure_taew2_1_checkpoint(checkpoint) == checkpoint + + +def test_taew2_2_checkpoint_pin_is_immutable_and_sha256() -> None: + assert "/main/" not in wan_vae.TAEW2_2_URL + assert re.search(r"/[0-9a-f]{40}/taew2_2\.pth$", wan_vae.TAEW2_2_URL) + assert re.fullmatch(r"[0-9a-f]{64}", wan_vae.TAEW2_2_SHA256) + + +def test_cached_taew2_2_checkpoint_is_verified(monkeypatch, tmp_path) -> None: + monkeypatch.setattr(wan_vae, "_cache_dir", lambda: tmp_path) + (tmp_path / "taew2_2.pth").write_bytes(b"tampered") + + with pytest.raises(RuntimeError, match="failed sha256 verification"): + wan_vae.ensure_taehv_checkpoint(z_dim=48) diff --git a/fastvideo/tests/mlx/test_wan22_sample.py b/fastvideo/tests/mlx/test_wan22_sample.py new file mode 100644 index 000000000..109e4a9f9 --- /dev/null +++ b/fastvideo/tests/mlx/test_wan22_sample.py @@ -0,0 +1,32 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for Wan2.2 warped DMD schedule construction.""" + +from __future__ import annotations + +import torch + +from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler +from fastvideo.mlx_runtime.wan22_sample import build_wan22_dmd_schedule + + +def test_warped_dmd_timesteps_differ_from_raw_indices() -> None: + schedule, warped = build_wan22_dmd_schedule([1000, 757, 522], flow_shift=5.0, warp_denoising_step=True) + _, raw = build_wan22_dmd_schedule([1000, 757, 522], flow_shift=5.0, warp_denoising_step=False) + assert raw == [1000.0, 757.0, 522.0] + # Warping maps step indices into the continuous flow-match schedule. + assert warped[0] == 1000.0 + assert warped[1] != 757.0 + assert warped[2] != 522.0 + # Sigmas are monotone-ish decreasing along the schedule. + sigmas = [schedule.sigma_for(t) for t in warped] + assert sigmas[0] >= sigmas[1] >= sigmas[2] + + +def test_warped_dmd_timesteps_match_training_scheduler() -> None: + steps = torch.tensor([1000, 757, 522], dtype=torch.long) + scheduler = FlowMatchEulerDiscreteScheduler(shift=5.0) + expected = torch.cat((scheduler.timesteps, torch.tensor([0.0])))[1000 - steps] + + _, warped = build_wan22_dmd_schedule(steps.tolist(), flow_shift=5.0, warp_denoising_step=True) + + torch.testing.assert_close(torch.tensor(warped), expected) diff --git a/fastvideo/tests/mlx/test_windowed_attention.py b/fastvideo/tests/mlx/test_windowed_attention.py new file mode 100644 index 000000000..ada241d9d --- /dev/null +++ b/fastvideo/tests/mlx/test_windowed_attention.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Correctness checks for chunked non-causal windowed attention (MLX).""" + +from __future__ import annotations + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core", reason="MLX is required for windowed-attention tests") + +from fastvideo.mlx_runtime.windowed_attention import ( # noqa: E402 + _key_ranges_for_block, + full_attention, + windowed_attention, +) + + +def _rand_qkv(b: int, h: int, s: int, d: int, seed: int = 0): + """ + Generate deterministic random query, key, and value tensors. + + Parameters: + b (int): Batch size. + h (int): Number of attention heads. + s (int): Sequence length. + d (int): Feature dimension. + seed (int): Random number generator seed. + + Returns: + tuple: Query, key, and value tensors with shape `(b, h, s, d)`. + """ + rng = np.random.default_rng(seed) + q = mx.array(rng.standard_normal((b, h, s, d)).astype(np.float32)) + k = mx.array(rng.standard_normal((b, h, s, d)).astype(np.float32)) + v = mx.array(rng.standard_normal((b, h, s, d)).astype(np.float32)) + return q, k, v + + +def _mean_cosine(a: "mx.array", b: "mx.array") -> float: + """ + Compute the mean cosine similarity between corresponding feature vectors. + + Parameters: + a: The first tensor. + b: The second tensor. + + Returns: + The mean cosine similarity, with protection against zero-norm vectors. + """ + a32 = np.array(a.astype(mx.float32)).reshape(-1, a.shape[-1]) + b32 = np.array(b.astype(mx.float32)).reshape(-1, b.shape[-1]) + num = np.sum(a32 * b32, axis=-1) + den = np.linalg.norm(a32, axis=-1) * np.linalg.norm(b32, axis=-1) + return float(np.mean(num / np.maximum(den, 1e-12))) + + +def test_windowed_matches_full_when_window_ge_seq() -> None: + """``window >= S`` must recover full attention (cosine > 0.999).""" + b, h, s, d = 1, 2, 64, 16 + q, k, v = _rand_qkv(b, h, s, d, seed=7) + scale = d**-0.5 + + ref = full_attention(q, k, v, scale=scale) + out = windowed_attention(q, k, v, window=s, sink=0, scale=scale) + mx.eval(ref, out) + + assert out.shape == q.shape + cos = _mean_cosine(ref, out) + assert cos > 0.999, f"expected cosine > 0.999 for window>=S, got {cos}" + + +def test_windowed_output_shape_equals_input() -> None: + b, h, s, d = 2, 3, 48, 8 + q, k, v = _rand_qkv(b, h, s, d, seed=11) + out = windowed_attention(q, k, v, window=16, sink=4) + mx.eval(out) + assert out.shape == q.shape + assert out.shape == (b, h, s, d) + + +def test_windowed_with_sinks_shape_and_finite() -> None: + b, h, s, d = 1, 2, 40, 8 + q, k, v = _rand_qkv(b, h, s, d, seed=19) + out = windowed_attention(q, k, v, window=8, sink=4, chunk_size=5) + mx.eval(out) + assert out.shape == (b, h, s, d) + arr = np.array(out) + assert np.isfinite(arr).all() + + +def test_windowed_result_does_not_depend_on_chunk_size() -> None: + q, k, v = _rand_qkv(1, 2, 40, 8, seed=23) + per_query = windowed_attention(q, k, v, window=8, sink=4, chunk_size=1) + chunked = windowed_attention(q, k, v, window=8, sink=4, chunk_size=5) + mx.eval(per_query, chunked) + + np.testing.assert_allclose(np.array(chunked), np.array(per_query), rtol=1e-5, atol=1e-5) + + +def test_sink_range_extends_beyond_local_window() -> None: + assert _key_ranges_for_block(qs=0, qe=1, seq_len=12, half=1, sink=5) == [(0, 5)] + + +def test_windowed_attention_enforces_per_query_local_and_sink_mask() -> None: + q = mx.zeros((1, 1, 6, 1)) + k = mx.zeros_like(q) + v = mx.arange(6).reshape(1, 1, 6, 1) + + out = windowed_attention(q, k, v, window=2, sink=3, chunk_size=4) + mx.eval(out) + + np.testing.assert_allclose(np.array(out).reshape(-1), [1.0, 1.0, 1.5, 2.0, 2.5, 2.4], atol=1e-6) + + +def test_invalid_window_raises() -> None: + q, k, v = _rand_qkv(1, 1, 8, 4, seed=1) + with pytest.raises(ValueError, match="window"): + windowed_attention(q, k, v, window=0) diff --git a/fastvideo/tests/mlx/tiny_wan.py b/fastvideo/tests/mlx/tiny_wan.py new file mode 100644 index 000000000..797d7fe15 --- /dev/null +++ b/fastvideo/tests/mlx/tiny_wan.py @@ -0,0 +1,290 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Tiny random-weight Wan DiT fixtures shared by the MLX runtime tests. + +Builds a miniature ``WanTransformer3DModel`` (all matmul dims are multiples of +the int8 group size 64, so the same model exercises the quantized paths) plus +the conversions into ``fastvideo.mlx_runtime.fastwan.MLXWanDiT``. +""" + +from __future__ import annotations + +import os + +import numpy as np +import torch + +os.environ.setdefault("FASTVIDEO_ATTENTION_BACKEND", "TORCH_SDPA") +os.environ.setdefault("MASTER_ADDR", "localhost") +os.environ.setdefault("MASTER_PORT", "29513") + +from fastvideo.configs.models.dits.wanvideo import ( # noqa: E402 + WanVideoArchConfig, + WanVideoConfig, +) +from fastvideo.forward_context import set_forward_context # noqa: E402 +from fastvideo.layers.rotary_embedding import get_rotary_pos_embed # noqa: E402 +from fastvideo.mlx_runtime.fastwan import ( # noqa: E402 + MLXQuantizationSpec, + MLXWanDiT, + MLXWanTransformerBlock, + mlx_block_weights_from_torch, + quantize_matrix, +) +from fastvideo.models.dits.wanvideo import WanTransformer3DModel # noqa: E402 +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch # noqa: E402 + +SEED = 2026 + +# All matmul dims must be multiples of the int8 group size (64) so the +# quantized variants of the tests can reuse the same tiny model. +TINY_ARCH = dict( + num_attention_heads=4, + attention_head_dim=16, + in_channels=16, + out_channels=16, + text_dim=64, + freq_dim=64, + ffn_dim=128, + num_layers=2, + patch_size=(1, 2, 2), + rope_max_seq_len=64, +) + + +def build_tiny_wan_config() -> WanVideoConfig: + """Build the deterministic miniature Wan video configuration. + + Returns: + WanVideoConfig: Configuration created from the tiny test architecture. + """ + return WanVideoConfig(arch_config=WanVideoArchConfig(**TINY_ARCH)) + + +def build_hf_config(config: WanVideoConfig) -> dict[str, object]: + """ + Builds the configuration dictionary required by the Hugging Face-style Wan model representation. + + Parameters: + config (WanVideoConfig): Wan video configuration to convert. + + Returns: + dict[str, object]: Configuration values for the Hugging Face-style model. + """ + return { + "num_attention_heads": config.num_attention_heads, + "attention_head_dim": config.attention_head_dim, + "in_channels": config.in_channels, + "out_channels": config.out_channels, + "text_dim": config.text_dim, + "freq_dim": config.freq_dim, + "ffn_dim": config.ffn_dim, + "num_layers": config.num_layers, + "patch_size": config.patch_size, + "text_len": config.text_len, + "rope_max_seq_len": config.rope_max_seq_len, + "eps": 1e-6, + } + + +def initialize_model_parameters(model: torch.nn.Module) -> None: + # ReplicatedLinear parameters are allocated with torch.empty and need an + # explicit initialization in tests to avoid undefined values. + """Initialize model parameters deterministically for reproducible tests. + + Parameters: + model (torch.nn.Module): Model whose parameters will be initialized. + """ + torch.manual_seed(SEED + 3) + with torch.no_grad(): + for name, param in model.named_parameters(): + if param.ndim <= 1: + if name.endswith("weight") and "norm" in name: + param.fill_(1.0) + else: + param.normal_(mean=0.0, std=0.02) + continue + torch.nn.init.xavier_uniform_(param) + + +def build_torch_model() -> WanTransformer3DModel: + """ + Build a deterministically initialized CPU float32 Wan transformer model. + + Returns: + WanTransformer3DModel: The initialized model in evaluation mode. + """ + config = build_tiny_wan_config() + model = WanTransformer3DModel(config=config, hf_config=build_hf_config(config)) + model = model.to(device="cpu", dtype=torch.float32) + initialize_model_parameters(model) + model.eval() + return model + + +def build_inputs() -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Create deterministic latent, text-encoder, and timestep inputs for the tiny Wan model. + + Returns: + tuple[torch.Tensor, torch.Tensor, torch.Tensor]: The latent hidden states, text-encoder hidden states, and timestep.""" + generator = torch.Generator(device="cpu").manual_seed(SEED + 1) + hidden_states = torch.randn(1, TINY_ARCH["in_channels"], 4, 8, 8, generator=generator, dtype=torch.float32) + encoder_hidden_states = torch.randn(1, 8, TINY_ARCH["text_dim"], generator=generator, dtype=torch.float32) + timestep = torch.tensor([10], dtype=torch.long) + return hidden_states, encoder_hidden_states, timestep + + +# Top-level weights: the MLX runtime uses the Diffusers key layout; the torch +# model uses FastVideo's module names (see WanVideoConfig.param_names_mapping). +TOP_LEVEL_KEY_MAP = { + "patch_embedding.weight": "patch_embedding.proj.weight", + "patch_embedding.bias": "patch_embedding.proj.bias", + "condition_embedder.time_embedder.linear_1.weight": "condition_embedder.time_embedder.mlp.fc_in.weight", + "condition_embedder.time_embedder.linear_1.bias": "condition_embedder.time_embedder.mlp.fc_in.bias", + "condition_embedder.time_embedder.linear_2.weight": "condition_embedder.time_embedder.mlp.fc_out.weight", + "condition_embedder.time_embedder.linear_2.bias": "condition_embedder.time_embedder.mlp.fc_out.bias", + "condition_embedder.time_proj.weight": "condition_embedder.time_modulation.linear.weight", + "condition_embedder.time_proj.bias": "condition_embedder.time_modulation.linear.bias", + "condition_embedder.text_embedder.linear_1.weight": "condition_embedder.text_embedder.fc_in.weight", + "condition_embedder.text_embedder.linear_1.bias": "condition_embedder.text_embedder.fc_in.bias", + "condition_embedder.text_embedder.linear_2.weight": "condition_embedder.text_embedder.fc_out.weight", + "condition_embedder.text_embedder.linear_2.bias": "condition_embedder.text_embedder.fc_out.bias", + "scale_shift_table": "scale_shift_table", + "proj_out.weight": "proj_out.weight", + "proj_out.bias": "proj_out.bias", +} + + +def mlx_dit_from_torch_model( + model: WanTransformer3DModel, + hf_config: dict[str, object], + *, + quantization: MLXQuantizationSpec | None = None, +) -> MLXWanDiT: + """ + Convert a PyTorch Wan DiT model to an MLX Wan DiT model. + + Parameters: + model (WanTransformer3DModel): PyTorch model whose parameters are converted. + hf_config (dict[str, object]): Model configuration used to construct the MLX representation. + quantization (MLXQuantizationSpec | None): Optional quantization specification for eligible weight matrices. + + Returns: + MLXWanDiT: The converted MLX model. + """ + import mlx.core as mx + + state = {name: value.detach().float() for name, value in model.state_dict().items()} + inner_dim = int(hf_config["num_attention_heads"]) * int(hf_config["attention_head_dim"]) # type: ignore[arg-type] + + weights = {} + for mlx_name, torch_name in TOP_LEVEL_KEY_MAP.items(): + tensor = state[torch_name] + if mlx_name == "patch_embedding.weight": + tensor = tensor.reshape(inner_dim, -1) + array = mx.array(tensor.numpy()) + if quantization is not None and mlx_name.endswith(".weight") and mlx_name != "scale_shift_table": + weights[mlx_name] = quantize_matrix(array, quantization) + else: + weights[mlx_name] = array + + blocks = [] + for torch_block in model.blocks: + block_weights = mlx_block_weights_from_torch(torch_block) + if quantization is not None: + block_weights = { + name: (quantize_matrix(value, quantization) + if name.endswith(".weight") and "norm" not in name and len(value.shape) >= 2 else value) + for name, value in block_weights.items() + } + blocks.append( + MLXWanTransformerBlock( + block_weights, + dim=inner_dim, + ffn_dim=int(hf_config["ffn_dim"]), # type: ignore[arg-type] + num_heads=int(hf_config["num_attention_heads"]), # type: ignore[arg-type] + eps=float(hf_config["eps"]), # type: ignore[arg-type] + )) + return MLXWanDiT(weights, blocks, dict(hf_config)) + + +def mlx_rotary_embeddings(hidden_states: torch.Tensor): + """ + Generate MLX rotary-position embedding tables matching the tiny Wan model. + + Parameters: + hidden_states (torch.Tensor): Input latent tensor whose spatial and temporal dimensions determine the embedding grid. + + Returns: + Tuple of MLX arrays containing the cosine and sine rotary-position embeddings. + """ + import mlx.core as mx + + _, _, frames, height, width = hidden_states.shape + p_t, p_h, p_w = TINY_ARCH["patch_size"] + head_dim = TINY_ARCH["attention_head_dim"] + hidden_size = TINY_ARCH["num_attention_heads"] * head_dim + rope_dim_list = [head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)] + freqs_cos, freqs_sin = get_rotary_pos_embed( + (frames // p_t, height // p_h, width // p_w), + hidden_size, + TINY_ARCH["num_attention_heads"], + rope_dim_list, + dtype=torch.float64, + rope_theta=10000, + ) + return ( + mx.array(freqs_cos.float().numpy()).astype(mx.float32), + mx.array(freqs_sin.float().numpy()).astype(mx.float32), + ) + + +def torch_reference_output(model, hidden_states, encoder_hidden_states, timestep) -> np.ndarray: + """ + Run the PyTorch model on the provided inputs and return its output as a NumPy array. + + Parameters: + model: The PyTorch Wan transformer model. + hidden_states: The latent hidden states provided to the model. + encoder_hidden_states: The encoder hidden states provided to the model. + timestep: The diffusion timestep provided to the model. + + Returns: + The model output as a CPU float32 NumPy array. + """ + with torch.no_grad(), set_forward_context( + current_timestep=0, + attn_metadata=None, + forward_batch=ForwardBatch(data_type="dummy"), + ): + output = model( + hidden_states=hidden_states, + encoder_hidden_states=encoder_hidden_states, + timestep=timestep, + ) + return output.detach().float().cpu().numpy() + + +def mlx_output(dit, hidden_states, encoder_hidden_states, timestep, freqs_cis) -> np.ndarray: + """ + Run the MLX DiT model on the provided inputs. + + Parameters: + dit: The MLX DiT model to evaluate. + hidden_states: Input latent states. + encoder_hidden_states: Text encoder states. + timestep: Diffusion timestep input. + freqs_cis: Rotary-position embedding frequencies. + + Returns: + np.ndarray: The model output as a float32 NumPy array. + """ + import mlx.core as mx + + out = dit( + mx.array(hidden_states.numpy()), + mx.array(encoder_hidden_states.numpy()), + mx.array(timestep.float().numpy()), + freqs_cis, + ) + mx.eval(out) + return np.array(out.astype(mx.float32)) diff --git a/fastvideo/tests/platforms/test_mps_vsa_error.py b/fastvideo/tests/platforms/test_mps_vsa_error.py new file mode 100644 index 000000000..7bb1444a6 --- /dev/null +++ b/fastvideo/tests/platforms/test_mps_vsa_error.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MPS must reject VSA before a model reaches an incompatible SDPA call.""" + +import pytest +import torch + +from fastvideo.platforms import AttentionBackendEnum +from fastvideo.platforms.mps import MpsPlatform + + +def test_mps_rejects_video_sparse_attention_with_actionable_error() -> None: + with pytest.raises(NotImplementedError, match="TORCH_SDPA"): + MpsPlatform.get_attn_backend_cls(AttentionBackendEnum.VIDEO_SPARSE_ATTN, 64, torch.float16) + + +def test_mps_resolves_sdpa_for_supported_backend() -> None: + assert MpsPlatform.get_attn_backend_cls(AttentionBackendEnum.TORCH_SDPA, 64, torch.float16).endswith( + "SDPABackend" + ) diff --git a/fastvideo/third_party/rife_mlx/LICENSE b/fastvideo/third_party/rife_mlx/LICENSE new file mode 100644 index 000000000..d040a4a3a --- /dev/null +++ b/fastvideo/third_party/rife_mlx/LICENSE @@ -0,0 +1,26 @@ +MIT License + +This directory vendors the inference path of rife-mlx +(https://github.com/xocialize/rife-mlx), an Apple MLX port of Practical-RIFE +(https://github.com/hzwer/Practical-RIFE). + +Copyright (c) 2025 xocialize (rife-mlx, MLX port) +Copyright (c) hzwer (Practical-RIFE / RIFE, original work) + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/fastvideo/third_party/rife_mlx/__init__.py b/fastvideo/third_party/rife_mlx/__init__.py new file mode 100644 index 000000000..04712887a --- /dev/null +++ b/fastvideo/third_party/rife_mlx/__init__.py @@ -0,0 +1,17 @@ +# SPDX-License-Identifier: MIT +"""Vendored rife-mlx (RIFE frame interpolation on Apple MLX). + +Upstream: https://github.com/xocialize/rife-mlx, itself an MLX port of +Practical-RIFE (https://github.com/hzwer/Practical-RIFE, (c) hzwer). Both are +MIT; see ``LICENSE`` for the text. + +Only the inference path is vendored, which is what ``--fast`` calls: +``utils.weights.build_model`` and ``pipeline_mlx.interpolate_pair``. Upstream's +video I/O and torch-to-MLX conversion helpers are omitted, so this package +needs no ``av``, ``PIL``, or torch import to load. Weights are still fetched at +runtime from the ``mlx-community`` Hugging Face org, exactly as upstream does. +""" + +from fastvideo.third_party.rife_mlx.config import DEFAULT_VERSION, VERSIONS, VersionConfig + +__all__ = ["DEFAULT_VERSION", "VERSIONS", "VersionConfig"] diff --git a/fastvideo/third_party/rife_mlx/config.py b/fastvideo/third_party/rife_mlx/config.py new file mode 100644 index 000000000..b42d23cf7 --- /dev/null +++ b/fastvideo/third_party/rife_mlx/config.py @@ -0,0 +1,45 @@ +"""Version registry — pinned config for the RIFE checkpoint(s) we port. + +Source of truth: the bundled `IFNet_HDv3.py` + `flownet.pkl` that ship together +per version (Google Drive), NOT the GitHub repo (which omits the version arch). +CONFIRM gate #3: pin the exact arch from the downloaded package, never guess. + +This round pins **RIFE 4.25** (upstream's recommended default; "anime scenes +significantly improved"). Arch PINNED in P1 from the bundled IFNet_HDv3.py + +flownet.pkl shapes: 5 IFBlocks (the 4.25 "more flow blocks" note), channels +[192,128,96,64,32], in_planes [15,28,28,28,28] (blocks 1-4 add the running +4-ch flow inside IFBlock.forward), scale_list [16,8,4,2,1] (÷ the --scale knob). +Activation is LeakyReLU(0.2). All F.interpolate use align_corners=False; warp +grid_sample uses align_corners=True. 40 teacher/caltime keys are train-only. +""" + +from __future__ import annotations + +from dataclasses import dataclass, field + + +@dataclass(frozen=True) +class VersionConfig: + name: str # internal key + hf_name: str # published HF repo name (mlx-community/) + gdrive_id: str # per-version package (model .py + flownet.pkl) + # arch params — PINNED from bundled v4.25 IFNet_HDv3.py + flownet.pkl (P1) + scale_list: tuple[float, ...] = (16, 8, 4, 2, 1) + block_channels: tuple[int, ...] = (192, 128, 96, 64, 32) + block_in_planes: tuple[int, ...] = (15, 28, 28, 28, 28) + pad_to: int = 64 # 5 blocks downsample; pad H,W to a multiple of 64 + # warp / interpolate semantics (verified from source) + grid_align_corners: bool = True # warp grid_sample + grid_padding_mode: str = "border" + interp_align_corners: bool = False # F.interpolate in IFBlock + + +VERSIONS: dict[str, VersionConfig] = { + "4.25": VersionConfig( + name="4.25", + hf_name="RIFE-4.25", + gdrive_id="1ZKjcbmt1hypiFprJPIKW0Tt0lr_2i7bg", + ), +} + +DEFAULT_VERSION = "4.25" diff --git a/fastvideo/third_party/rife_mlx/model/IFNet_HDv3.py b/fastvideo/third_party/rife_mlx/model/IFNet_HDv3.py new file mode 100644 index 000000000..c0fe82d1c --- /dev/null +++ b/fastvideo/third_party/rife_mlx/model/IFNet_HDv3.py @@ -0,0 +1,142 @@ +"""IFNet (RIFE 4.25) — MLX port, isomorphic to the bundled train_log/IFNet_HDv3.py. + +NHWC throughout. Translated line-by-line from the pinned v4.25 source; only +PyTorch ops -> MLX ops and NCHW -> NHWC change. Parity-locked in P3 (Gate B). + +Arch (pinned): 5 IFBlocks c=[192,128,96,64,32], in_planes=[15,28,28,28,28]; +conv = Conv2d + LeakyReLU(0.2); ResConv = conv*beta + x -> LeakyReLU; lastconv = +ConvTranspose2d(c,52,4,2,1) + PixelShuffle(2) -> 13ch (flow4+mask1+feat8); Head +encoder -> 4ch; F.interpolate align_corners=False; warp align_corners=True. + +MLX param keys (post-conversion) match torch except: conv0.{j}.0.* -> conv0.{j}.conv.* +and lastconv.0.* -> lastconv.* (see utils/convert.py). +""" + +from __future__ import annotations + +import mlx.core as mx +import mlx.nn as nn + +from ..ops.interpolate import interpolate_bilinear +from ..ops.spatial import pixel_shuffle_nhwc +from .warplayer import warp + +_SLOPE = 0.2 + + +def _lrelu(x): + return nn.leaky_relu(x, _SLOPE) + + +class Conv(nn.Module): + """conv = Conv2d + LeakyReLU(0.2). key:

.conv.weight""" + + def __init__(self, i, o, k=3, s=1, p=1): + super().__init__() + self.conv = nn.Conv2d(i, o, kernel_size=k, stride=s, padding=p) + + def __call__(self, x): + return _lrelu(self.conv(x)) + + +class ResConv(nn.Module): + """conv*beta + x -> LeakyReLU. keys: .conv.weight/bias, .beta""" + + def __init__(self, c): + super().__init__() + self.conv = nn.Conv2d(c, c, kernel_size=3, stride=1, padding=1) + self.beta = mx.ones((1, 1, 1, c)) # NHWC per-channel + + def __call__(self, x): + return _lrelu(self.conv(x) * self.beta + x) + + +class Head(nn.Module): + """Feature encoder. keys: cnn0,cnn1,cnn2 (Conv2d), cnn3 (ConvTranspose2d).""" + + def __init__(self): + super().__init__() + self.cnn0 = nn.Conv2d(3, 16, 3, stride=2, padding=1) + self.cnn1 = nn.Conv2d(16, 16, 3, stride=1, padding=1) + self.cnn2 = nn.Conv2d(16, 16, 3, stride=1, padding=1) + self.cnn3 = nn.ConvTranspose2d(16, 4, 4, stride=2, padding=1) + + def __call__(self, x): + x = _lrelu(self.cnn0(x)) + x = _lrelu(self.cnn1(x)) + x = _lrelu(self.cnn2(x)) + return self.cnn3(x) + + +class IFBlock(nn.Module): + def __init__(self, in_planes, c=64): + super().__init__() + self.conv0 = [Conv(in_planes, c // 2, 3, 2, 1), Conv(c // 2, c, 3, 2, 1)] + self.convblock = [ResConv(c) for _ in range(8)] + self.lastconv = nn.ConvTranspose2d(c, 4 * 13, 4, stride=2, padding=1) + + def __call__(self, x, flow=None, scale=1): + x = interpolate_bilinear(x, scale_factor=1.0 / scale, align_corners=False) + if flow is not None: + flow = interpolate_bilinear(flow, scale_factor=1.0 / scale, + align_corners=False) * (1.0 / scale) + x = mx.concatenate([x, flow], axis=-1) + feat = self.conv0[0](x) + feat = self.conv0[1](feat) + for blk in self.convblock: + feat = blk(feat) + tmp = pixel_shuffle_nhwc(self.lastconv(feat), 2) + tmp = interpolate_bilinear(tmp, scale_factor=scale, align_corners=False) + flow = tmp[..., :4] * scale + mask = tmp[..., 4:5] + feat = tmp[..., 5:] + return flow, mask, feat + + +class IFNet(nn.Module): + def __init__(self): + super().__init__() + self.block0 = IFBlock(7 + 8, c=192) + self.block1 = IFBlock(8 + 4 + 8 + 8, c=128) + self.block2 = IFBlock(8 + 4 + 8 + 8, c=96) + self.block3 = IFBlock(8 + 4 + 8 + 8, c=64) + self.block4 = IFBlock(8 + 4 + 8 + 8, c=32) + self.encode = Head() + + def __call__(self, x, timestep=0.5, scale_list=(16, 8, 4, 2, 1)): + """x: [N,H,W,6] (img0|img1 on channel axis). Returns (flow_list, mask, merged).""" + channel = x.shape[-1] // 2 + img0 = x[..., :channel] + img1 = x[..., channel:] + N, H, W, _ = img0.shape + timestep = mx.full((N, H, W, 1), float(timestep)) + + f0 = self.encode(img0[..., :3]) + f1 = self.encode(img1[..., :3]) + flow_list, merged, mask_list = [], [], [] + warped_img0, warped_img1 = img0, img1 + flow = None + mask = None + blocks = [self.block0, self.block1, self.block2, self.block3, self.block4] + for i in range(5): + if flow is None: + flow, mask, feat = blocks[i]( + mx.concatenate([img0[..., :3], img1[..., :3], f0, f1, timestep], axis=-1), + None, scale=scale_list[i]) + else: + wf0 = warp(f0, flow[..., :2]) + wf1 = warp(f1, flow[..., 2:4]) + fd, m0, feat = blocks[i]( + mx.concatenate([warped_img0[..., :3], warped_img1[..., :3], + wf0, wf1, timestep, mask, feat], axis=-1), + flow, scale=scale_list[i]) + mask = m0 + flow = flow + fd + mask_list.append(mask) + flow_list.append(flow) + warped_img0 = warp(img0, flow[..., :2]) + warped_img1 = warp(img1, flow[..., 2:4]) + merged.append((warped_img0, warped_img1)) + mask = mx.sigmoid(mask) + merged_final = warped_img0 * mask + warped_img1 * (1 - mask) + return flow_list, mask, merged_final diff --git a/fastvideo/third_party/rife_mlx/model/RIFE_HDv3.py b/fastvideo/third_party/rife_mlx/model/RIFE_HDv3.py new file mode 100644 index 000000000..635a3b80a --- /dev/null +++ b/fastvideo/third_party/rife_mlx/model/RIFE_HDv3.py @@ -0,0 +1,37 @@ +"""Model wrapper — MLX port of train_log/RIFE_HDv3.py. + +Wraps IFNet and exposes inference(img0, img1, timestep, scale) -> middle frame, +matching upstream: scale_list = [16,8,4,2,1] / scale; H,W padded to a multiple of +pad_to (64; the 5-block downsample factor) and cropped back after. +""" + +from __future__ import annotations + +import mlx.core as mx + +from .IFNet_HDv3 import IFNet + + +class Model: + def __init__(self, scale_list=(16, 8, 4, 2, 1), pad_to: int = 64) -> None: + self.flownet = IFNet() + self.scale_list = tuple(scale_list) + self.pad_to = pad_to + + def inference(self, img0: mx.array, img1: mx.array, + timestep: float = 0.5, scale: float = 1.0) -> mx.array: + """img0,img1: [N,H,W,3] in [0,1]. Returns the frame at `timestep`.""" + N, H, W, _ = img0.shape + # pad must scale with 1/scale so the coarsest-scale interpolation + # round-trips exactly (else the recovered flow size != padded size). + pad = max(self.pad_to, int(round(self.pad_to / scale))) + ph = ((H - 1) // pad + 1) * pad + pw = ((W - 1) // pad + 1) * pad + pad = [(0, 0), (0, ph - H), (0, pw - W), (0, 0)] + i0 = mx.pad(img0, pad) + i1 = mx.pad(img1, pad) + + scale_list = tuple(s / scale for s in self.scale_list) + x = mx.concatenate([i0, i1], axis=-1) + _flow, _mask, merged = self.flownet(x, timestep, scale_list) + return merged[:, :H, :W, :] diff --git a/fastvideo/third_party/rife_mlx/model/__init__.py b/fastvideo/third_party/rife_mlx/model/__init__.py new file mode 100644 index 000000000..2e0556e0d --- /dev/null +++ b/fastvideo/third_party/rife_mlx/model/__init__.py @@ -0,0 +1,5 @@ +"""IFNet / RIFE model modules — isomorphic to upstream train_log/.""" + +from .warplayer import warp + +__all__ = ["warp"] diff --git a/fastvideo/third_party/rife_mlx/model/warplayer.py b/fastvideo/third_party/rife_mlx/model/warplayer.py new file mode 100644 index 000000000..651736193 --- /dev/null +++ b/fastvideo/third_party/rife_mlx/model/warplayer.py @@ -0,0 +1,37 @@ +"""warp — MLX port of model/warplayer.py (backward warp via grid_sample). + +Isomorphic to upstream: builds a normalized identity grid, adds the +per-pixel flow (normalized by (size-1)/2), and bilinear grid_samples the +input with border padding + align_corners=True. + +NHWC throughout. flow is [N,H,W,2] with channel 0 = x (horizontal), 1 = y. +""" + +from __future__ import annotations + +import mlx.core as mx + +from ..ops.grid_sample import grid_sample_bilinear + +_grid_cache: dict[tuple, mx.array] = {} + + +def _identity_grid(N: int, H: int, W: int) -> mx.array: + key = (N, H, W) + if key not in _grid_cache: + xs = mx.linspace(-1.0, 1.0, W).reshape(1, 1, W, 1) + ys = mx.linspace(-1.0, 1.0, H).reshape(1, 1, H, 1) # placeholder shape + xs = mx.broadcast_to(xs, (N, H, W, 1)) + ys = mx.broadcast_to(mx.linspace(-1.0, 1.0, H).reshape(1, H, 1, 1), (N, H, W, 1)) + _grid_cache[key] = mx.concatenate([xs, ys], axis=-1) # [N,H,W,2] + return _grid_cache[key] + + +def warp(tenInput: mx.array, tenFlow: mx.array) -> mx.array: + """tenInput: [N,H,W,C]; tenFlow: [N,H,W,2] (x,y pixel-space flow).""" + N, H, W, _ = tenFlow.shape + grid = _identity_grid(N, H, W) + fx = tenFlow[..., 0:1] / ((W - 1.0) / 2.0) + fy = tenFlow[..., 1:2] / ((H - 1.0) / 2.0) + g = grid + mx.concatenate([fx, fy], axis=-1) + return grid_sample_bilinear(tenInput, g, align_corners=True, padding_mode="border") diff --git a/fastvideo/third_party/rife_mlx/ops/__init__.py b/fastvideo/third_party/rife_mlx/ops/__init__.py new file mode 100644 index 000000000..e01077970 --- /dev/null +++ b/fastvideo/third_party/rife_mlx/ops/__init__.py @@ -0,0 +1,6 @@ +"""MLX ops hand-rolled for RIFE (no native equivalents).""" + +from .grid_sample import grid_sample_bilinear +from .interpolate import interpolate_bilinear + +__all__ = ["grid_sample_bilinear", "interpolate_bilinear"] diff --git a/fastvideo/third_party/rife_mlx/ops/grid_sample.py b/fastvideo/third_party/rife_mlx/ops/grid_sample.py new file mode 100644 index 000000000..52ce7196a --- /dev/null +++ b/fastvideo/third_party/rife_mlx/ops/grid_sample.py @@ -0,0 +1,60 @@ +"""Bilinear grid_sample (NHWC) — MLX equivalent of torch.nn.functional.grid_sample. + +RIFE's `warp` calls grid_sample(mode='bilinear', padding_mode='border', +align_corners=True). No native MLX op — hand-rolled here. This is the crux +surface of the port; parity-locked vs torch in P2 (Gate A). + +Conventions (match torch): +- grid[..., 0] = x (width, normalized [-1,1]); grid[..., 1] = y (height). +- align_corners=True: -1 -> pixel 0, +1 -> pixel (size-1). +- padding_mode='border': neighbor integer indices are clamped to the valid + range; bilinear weights come from the UNCLAMPED continuous coordinate. +""" + +from __future__ import annotations + +import mlx.core as mx + + +def grid_sample_bilinear(inp: mx.array, grid: mx.array, + align_corners: bool = True, + padding_mode: str = "border") -> mx.array: + """inp: [N,H,W,C]; grid: [N,gH,gW,2] in [-1,1]. Returns [N,gH,gW,C].""" + if padding_mode != "border": + raise NotImplementedError("only padding_mode='border' is needed for RIFE") + N, H, W, C = inp.shape + _, gH, gW, _ = grid.shape + + gx = grid[..., 0] # [N,gH,gW] + gy = grid[..., 1] + if align_corners: + ix = (gx + 1) * 0.5 * (W - 1) + iy = (gy + 1) * 0.5 * (H - 1) + else: + ix = ((gx + 1) * W - 1) * 0.5 + iy = ((gy + 1) * H - 1) * 0.5 + + x0 = mx.floor(ix); y0 = mx.floor(iy) + x1 = x0 + 1; y1 = y0 + 1 + wx1 = ix - x0; wx0 = 1.0 - wx1 + wy1 = iy - y0; wy0 = 1.0 - wy1 + + def clampx(a): return mx.clip(a, 0, W - 1).astype(mx.int32) + def clampy(a): return mx.clip(a, 0, H - 1).astype(mx.int32) + x0c, x1c = clampx(x0), clampx(x1) + y0c, y1c = clampy(y0), clampy(y1) + + inp_flat = inp.reshape(N, H * W, C) + + def gather(yc, xc): # yc,xc: [N,gH,gW] int -> [N,gH,gW,C] + idx = (yc * W + xc).reshape(N, gH * gW, 1) + idx = mx.broadcast_to(idx, (N, gH * gW, C)) + g = mx.take_along_axis(inp_flat, idx, axis=1) + return g.reshape(N, gH, gW, C) + + v00 = gather(y0c, x0c); v01 = gather(y0c, x1c) + v10 = gather(y1c, x0c); v11 = gather(y1c, x1c) + + w00 = (wy0 * wx0)[..., None]; w01 = (wy0 * wx1)[..., None] + w10 = (wy1 * wx0)[..., None]; w11 = (wy1 * wx1)[..., None] + return v00 * w00 + v01 * w01 + v10 * w10 + v11 * w11 diff --git a/fastvideo/third_party/rife_mlx/ops/interpolate.py b/fastvideo/third_party/rife_mlx/ops/interpolate.py new file mode 100644 index 000000000..9774fd4a4 --- /dev/null +++ b/fastvideo/third_party/rife_mlx/ops/interpolate.py @@ -0,0 +1,51 @@ +"""Bilinear resize (NHWC) — MLX equivalent of F.interpolate(mode='bilinear'). + +RIFE uses bilinear resize for the coarse-to-fine scale pyramid inside IFNet and +for the `--scale` knob (0.25..4.0). align_corners is pinned from the bundled +IFNet_HDv3.py in P3 (RIFE 4.x typically uses align_corners=False). Parity-locked +in P2 alongside grid_sample. +""" + +from __future__ import annotations + +import mlx.core as mx + + +def _sample_coords(out: int, in_: int, align_corners: bool) -> mx.array: + dst = mx.arange(out, dtype=mx.float32) + if align_corners: + scale = (in_ - 1) / (out - 1) if out > 1 else 0.0 + return dst * scale + scale = in_ / out + return (dst + 0.5) * scale - 0.5 + + +def _bilinear_1d(x: mx.array, axis: int, out: int, align_corners: bool) -> mx.array: + """Resample one spatial axis with bilinear weights + border clamp.""" + in_ = x.shape[axis] + if in_ == out: + return x + src = _sample_coords(out, in_, align_corners) # [out] + i0 = mx.floor(src); i1 = i0 + 1 + w1 = src - i0; w0 = 1.0 - w1 + i0 = mx.clip(i0, 0, in_ - 1).astype(mx.int32) + i1 = mx.clip(i1, 0, in_ - 1).astype(mx.int32) + g0 = mx.take(x, i0, axis=axis) + g1 = mx.take(x, i1, axis=axis) + shape = [1] * x.ndim; shape[axis] = out + w0 = w0.reshape(shape); w1 = w1.reshape(shape) + return g0 * w0 + g1 * w1 + + +def interpolate_bilinear(x: mx.array, size: tuple[int, int] | None = None, + scale_factor: float | None = None, + align_corners: bool = False) -> mx.array: + """x: [N,H,W,C]. Resize H,W to `size` or by `scale_factor` (bilinear).""" + N, H, W, C = x.shape + if size is not None: + oH, oW = size + else: + oH, oW = int(round(H * scale_factor)), int(round(W * scale_factor)) + x = _bilinear_1d(x, axis=1, out=oH, align_corners=align_corners) + x = _bilinear_1d(x, axis=2, out=oW, align_corners=align_corners) + return x diff --git a/fastvideo/third_party/rife_mlx/ops/spatial.py b/fastvideo/third_party/rife_mlx/ops/spatial.py new file mode 100644 index 000000000..e42ea55ea --- /dev/null +++ b/fastvideo/third_party/rife_mlx/ops/spatial.py @@ -0,0 +1,18 @@ +"""NHWC pixel-shuffle — verified (C,r,r) ordering matching torch.nn.PixelShuffle. + +Reused from the realesrgan-mlx port (mlx-porting pitfall #7). RIFE's IFBlock +lastconv ends in nn.PixelShuffle(2). +""" + +from __future__ import annotations + +import mlx.core as mx + + +def pixel_shuffle_nhwc(x: mx.array, r: int) -> mx.array: + """[N,H,W,C*r^2] -> [N,H*r,W*r,C], channel split (C,r,r) (torch-matching).""" + n, h, w, c_in = x.shape + c = c_in // (r * r) + x = x.reshape((n, h, w, c, r, r)) + x = mx.transpose(x, (0, 1, 4, 2, 5, 3)) # (N,H,r_i,W,r_j,C) + return x.reshape((n, h * r, w * r, c)) diff --git a/fastvideo/third_party/rife_mlx/pipeline_mlx.py b/fastvideo/third_party/rife_mlx/pipeline_mlx.py new file mode 100644 index 000000000..fc2fbed60 --- /dev/null +++ b/fastvideo/third_party/rife_mlx/pipeline_mlx.py @@ -0,0 +1,80 @@ +"""High-level interpolation + CLI (image-pair + full video). P5. + + rife-mlx -i in.mp4 -o out.mp4 --multi 2 --scale 1.0 + rife-mlx --img0 a.png --img1 b.png -t 0.5 -o mid.png + +--multi N inserts N-1 frames between each pair (at t = i/N); output fps = N*input. +--scale {0.25,0.5,1,2,4} sets the coarse-to-fine pyramid scale (0.5 for 4K). +""" + +from __future__ import annotations + +import argparse + +import mlx.core as mx +import numpy as np + +from .config import DEFAULT_VERSION +from .utils.weights import build_model + + +def _to_nhwc(img_u8: np.ndarray) -> mx.array: + return mx.array(img_u8.astype(np.float32)[None] / 255.0) + + +def _to_u8(arr: mx.array) -> np.ndarray: + a = np.clip(np.array(arr)[0], 0.0, 1.0) + return (a * 255.0).round().astype(np.uint8) + + +def interpolate_pair(model, img0_u8, img1_u8, timestep=0.5, scale=1.0) -> np.ndarray: + out = model.inference(_to_nhwc(img0_u8), _to_nhwc(img1_u8), timestep, scale) + mx.eval(out) + return _to_u8(out) + + +def interpolate_sequence(model, frames, multi=2, scale=1.0): + """frames: list[HWC uint8] -> list with multi-1 frames inserted per pair.""" + out = [] + for i in range(len(frames) - 1): + out.append(frames[i]) + for j in range(1, multi): + out.append(interpolate_pair(model, frames[i], frames[i + 1], j / multi, scale)) + out.append(frames[-1]) + return out + + +def cli_main() -> None: + p = argparse.ArgumentParser(description="Practical-RIFE 4.25 MLX frame interpolation") + p.add_argument("-i", "--input", help="input video") + p.add_argument("--img0"); p.add_argument("--img1") + p.add_argument("-o", "--output", required=True) + p.add_argument("-t", "--timestep", type=float, default=0.5, help="image-pair mode") + p.add_argument("-m", "--multi", type=int, default=2, help="video fps multiplier") + p.add_argument("-s", "--scale", type=float, default=1.0, + choices=[0.25, 0.5, 1.0, 2.0, 4.0]) + p.add_argument("-n", "--version", default=DEFAULT_VERSION) + p.add_argument("--weights_dir", default=None) + args = p.parse_args() + + model = build_model(args.version, weights_dir=args.weights_dir) + + if args.img0 and args.img1: + from PIL import Image + a = np.asarray(Image.open(args.img0).convert("RGB")) + b = np.asarray(Image.open(args.img1).convert("RGB")) + mid = interpolate_pair(model, a, b, args.timestep, args.scale) + Image.fromarray(mid).save(args.output) + print(f"{args.img0} + {args.img1} @t={args.timestep} -> {args.output} {mid.shape}") + return + + from .video import read_frames, write_video + frames, fps = read_frames(args.input) + out_frames = interpolate_sequence(model, frames, args.multi, args.scale) + write_video(args.output, out_frames, fps * args.multi, audio_from=args.input) + print(f"{args.input} ({len(frames)}f @{fps:.2f}) -> {args.output} " + f"({len(out_frames)}f @{fps * args.multi:.2f}, audio passthrough)") + + +if __name__ == "__main__": + cli_main() diff --git a/fastvideo/third_party/rife_mlx/utils/__init__.py b/fastvideo/third_party/rife_mlx/utils/__init__.py new file mode 100644 index 000000000..84bdd8863 --- /dev/null +++ b/fastvideo/third_party/rife_mlx/utils/__init__.py @@ -0,0 +1 @@ +"""Weight loading + conversion utilities.""" diff --git a/fastvideo/third_party/rife_mlx/utils/weights.py b/fastvideo/third_party/rife_mlx/utils/weights.py new file mode 100644 index 000000000..bdf9b1770 --- /dev/null +++ b/fastvideo/third_party/rife_mlx/utils/weights.py @@ -0,0 +1,44 @@ +"""Model build + weight load (HF auto-download / local override). + +Resolution: weights_dir arg | $RIFE_MLX_WEIGHTS_DIR | dist/ | HF +mlx-community/. Each dir: model.safetensors + config.json. +""" + +from __future__ import annotations + +import os + +import mlx.core as mx + +from ..config import VERSIONS, VersionConfig +from ..model.RIFE_HDv3 import Model + +WEIGHTS_DIR_ENV = "RIFE_MLX_WEIGHTS_DIR" +HF_ORG = "mlx-community" + + +def resolve_repo_id(version: str) -> str: + return f"{HF_ORG}/{VERSIONS[version].hf_name}" + + +def _resolve_dir(cfg: VersionConfig, weights_dir: str | None) -> str: + if weights_dir: + return weights_dir + env = os.environ.get(WEIGHTS_DIR_ENV) + if env and os.path.isdir(env): + return env + local = os.path.join("dist", cfg.hf_name) + if os.path.isdir(local): + return local + from huggingface_hub import snapshot_download + return snapshot_download(repo_id=f"{HF_ORG}/{cfg.hf_name}") + + +def build_model(version: str = "4.25", weights_dir: str | None = None) -> Model: + cfg = VERSIONS[version] + wdir = _resolve_dir(cfg, weights_dir) + weights = dict(mx.load(os.path.join(wdir, "model.safetensors")).items()) + model = Model(scale_list=cfg.scale_list, pad_to=cfg.pad_to) + model.flownet.load_weights(list(weights.items()), strict=True) + mx.eval(model.flownet.parameters()) + return model diff --git a/fastvideo/third_party/taehv/LICENSE b/fastvideo/third_party/taehv/LICENSE new file mode 100644 index 000000000..6dc1edfbb --- /dev/null +++ b/fastvideo/third_party/taehv/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2025 Ollin Boer Bohan + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/fastvideo/third_party/taehv/__init__.py b/fastvideo/third_party/taehv/__init__.py new file mode 100644 index 000000000..0c9143c85 --- /dev/null +++ b/fastvideo/third_party/taehv/__init__.py @@ -0,0 +1,10 @@ +# SPDX-License-Identifier: MIT +"""Vendored TAEHV (Tiny AutoEncoder for Hunyuan/Wan video latents). + +Upstream: https://github.com/madebyollin/taehv -- see ``taehv.py``'s header for +provenance and ``LICENSE`` for the MIT license text. +""" + +from fastvideo.third_party.taehv.taehv import TAEHV + +__all__ = ["TAEHV"] diff --git a/fastvideo/third_party/taehv/taehv.py b/fastvideo/third_party/taehv/taehv.py new file mode 100644 index 000000000..9feb522c3 --- /dev/null +++ b/fastvideo/third_party/taehv/taehv.py @@ -0,0 +1,475 @@ +# Vendored from https://github.com/madebyollin/taehv (MIT License, (c) 2025 +# Ollin Boer Bohan; see LICENSE in this directory). +# +# Source file: https://raw.githubusercontent.com/madebyollin/taehv/main/taehv.py +# Accessed: 2026-07-02 +# Upstream file sha256: 1e228d34d47e2f95b17f1aa190f551d89ce5ef4af0f82be459a46bea81ba3fd2 +# +# Vendored unmodified below this header so the Apple Silicon TAEHV decode path +# (fastvideo/mlx_runtime/taehv_decode.py) does not download and execute remote +# code at runtime. To update, re-fetch the file, refresh the sha256 above, and +# re-run fastvideo/tests/mlx/. +#!/usr/bin/env python3 +""" +Tiny AutoEncoder for Hunyuan Video +(DNN for encoding / decoding videos to Hunyuan Video's latent space) +""" +import torch +import torch.nn as nn +import torch.nn.functional as F +from tqdm.auto import tqdm +from collections import namedtuple + +TWorkItem = namedtuple("TWorkItem", ("input_tensor", "block_index")) + +def conv(n_in, n_out, **kwargs): + return nn.Conv2d(n_in, n_out, 3, padding=1, **kwargs) + +class Clamp(nn.Module): + def forward(self, x): + return torch.tanh(x / 3) * 3 + +class MemBlock(nn.Module): + def __init__(self, n_in, n_out): + super().__init__() + self.conv = nn.Sequential(conv(n_in * 2, n_out), nn.ReLU(inplace=True), conv(n_out, n_out), nn.ReLU(inplace=True), conv(n_out, n_out)) + self.skip = nn.Conv2d(n_in, n_out, 1, bias=False) if n_in != n_out else nn.Identity() + self.act = nn.ReLU(inplace=True) + def forward(self, x, past): + return self.act(self.conv(torch.cat([x, past], 1)) + self.skip(x)) + +class TPool(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f*stride,n_f, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + return self.conv(x.reshape(-1, self.stride * C, H, W)) + +class TGrow(nn.Module): + def __init__(self, n_f, stride): + super().__init__() + self.stride = stride + self.conv = nn.Conv2d(n_f, n_f*stride, 1, bias=False) + def forward(self, x): + _NT, C, H, W = x.shape + x = self.conv(x) + return x.reshape(-1, C, H, W) + +def apply_model_with_memblocks_parallel(model, x, show_progress_bar): + """ + Apply a sequential model with memblocks to the given input, + with parallelization over the time axis and iteration over blocks. + + Args: + - model: nn.Sequential of blocks to apply + - x: input data, of dimensions NTCHW + - show_progress_bar: if True, enables tqdm progressbar display + + Returns NTCHW tensor of output data. + """ + assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor" + N, T, C, H, W = x.shape + x = x.reshape(N*T, C, H, W) + + # parallel over input timesteps, iterate over blocks + for b in tqdm(model, disable=not show_progress_bar): + if isinstance(b, MemBlock): + NT, C, H, W = x.shape + T = NT // N + _x = x.reshape(N, T, C, H, W) + # pad with zeros along time axis (i.e. empty memory), slice + block_memory = F.pad(_x, (0,0,0,0,0,0,1,0), value=0)[:,:T].reshape(x.shape) + x = b(x, block_memory) + else: + x = b(x) + NT, C, H, W = x.shape + T = NT // N + return x.view(N, T, C, H, W) + +def apply_model_with_memblocks_sequential_single_step(model, memory, work_queue, progress_bar=None): + """ + Process the work queue (a graph traversal over blocks and timesteps) + until an output frame is produced or the queue is empty. + Mutates memory and work_queue in place. + + Returns N1CHW output tensor, or None if the queue needs more input. + """ + while work_queue: + xt, i = work_queue.pop(0) + if progress_bar is not None and i == 0: + progress_bar.update(1) + if i == len(model): + return xt.unsqueeze(1) + b = model[i] + if isinstance(b, MemBlock): + # mem blocks are simple since we're visiting the graph in causal order + if memory[i] is None: + xt_new = b(xt, xt * 0) + else: + xt_new = b(xt, memory[i]) + memory[i] = xt + work_queue.insert(0, TWorkItem(xt_new, i+1)) + elif isinstance(b, TPool): + # pool blocks accumulate inputs until they have enough to pool + if memory[i] is None: + memory[i] = [] + memory[i].append(xt) + if len(memory[i]) > b.stride: + raise ValueError(f"TPool memory overflow: {len(memory[i])} items for stride {b.stride}") + elif len(memory[i]) == b.stride: + N, C, H, W = xt.shape + xt = b(torch.cat(memory[i], 1).view(N*b.stride, C, H, W)) + memory[i] = [] + work_queue.insert(0, TWorkItem(xt, i+1)) + elif isinstance(b, TGrow): + xt = b(xt) + NT, C, H, W = xt.shape + for xt_next in reversed(xt.view(NT//b.stride, b.stride*C, H, W).chunk(b.stride, 1)): + work_queue.insert(0, TWorkItem(xt_next, i+1)) + else: + xt = b(xt) + work_queue.insert(0, TWorkItem(xt, i+1)) + return None + +def apply_model_with_memblocks_sequential(model, x, show_progress_bar): + """ + Apply a sequential model with memblocks to the given input, + with iteration over timesteps as well as blocks. + + Args: + - model: nn.Sequential of blocks to apply + - x: input data, of dimensions NTCHW + - show_progress_bar: if True, enables tqdm progressbar display + + Returns NTCHW tensor of output data. + """ + assert x.ndim == 5, f"TAEHV operates on NTCHW tensors, but got {x.ndim}-dim tensor" + work_queue = [TWorkItem(xt, 0) for xt in x.unbind(1)] + memory = [None] * len(model) + progress_bar = tqdm(range(len(work_queue)), disable=not show_progress_bar) + out = [] + while work_queue: + xt = apply_model_with_memblocks_sequential_single_step(model, memory, work_queue, progress_bar) + if xt is not None: + out.append(xt) + progress_bar.close() + return torch.cat(out, 1) + +def apply_model_with_memblocks(model, x, parallel, show_progress_bar): + """ + Apply a sequential model with memblocks to the given input. + Args: + - model: nn.Sequential of blocks to apply + - x: input data, of dimensions NTCHW + - parallel: if True, parallelize over timesteps (fast but uses O(T) memory) + if False, each timestep will be processed sequentially (slow but uses O(1) memory) + - show_progress_bar: if True, enables tqdm progressbar display + + Returns NTCHW tensor of output data. + """ + if parallel: + return apply_model_with_memblocks_parallel(model, x, show_progress_bar) + else: + return apply_model_with_memblocks_sequential(model, x, show_progress_bar) + +class TAEHV(nn.Module): + def __init__(self, checkpoint_path="taehv.pth", encoder_time_downscale=(True, True, False), decoder_time_upscale=(False, True, True), decoder_space_upscale=(True, True, True), patch_size=1, latent_channels=16): + """Initialize pretrained TAEHV from the given checkpoint. + + Arg: + checkpoint_path: path to weight file to load. taehv.pth for Hunyuan, taew2_1.pth for Wan 2.1. + encoder_time_downscale: whether temporal downsampling is enabled for each block. + decoder_time_upscale: whether temporal upsampling is enabled for each block. upsampling can be disabled for a cheaper preview. + decoder_space_upscale: whether spatial upsampling is enabled for each block. upsampling can be disabled for a cheaper preview. + patch_size: input/output pixelshuffle patch-size for this model. + latent_channels: number of latent channels (z dim) for this model. + """ + super().__init__() + self.patch_size = patch_size + self.latent_channels = latent_channels + self.image_channels = 3 + if len(decoder_time_upscale) == 2: + decoder_time_upscale = (False, *decoder_time_upscale) + self.is_cogvideox = checkpoint_path is not None and "taecvx" in checkpoint_path + if checkpoint_path is not None and "taew2_2" in checkpoint_path: + self.patch_size, self.latent_channels = 2, 48 + if checkpoint_path is not None and "taehv1_5" in checkpoint_path: + self.patch_size, self.latent_channels = 2, 32 + if checkpoint_path is not None and "taeltx" in checkpoint_path: # same for both 2 and 2.3 + self.patch_size, self.latent_channels, encoder_time_downscale, decoder_time_upscale = 4, 128, (True, True, True), (True, True, True) + self.encoder = nn.Sequential( + conv(self.image_channels*self.patch_size**2, 64), nn.ReLU(inplace=True), + TPool(64, 2 if encoder_time_downscale[0] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + TPool(64, 2 if encoder_time_downscale[1] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + TPool(64, 2 if encoder_time_downscale[2] else 1), conv(64, 64, stride=2, bias=False), MemBlock(64, 64), MemBlock(64, 64), MemBlock(64, 64), + conv(64, self.latent_channels), + ) + n_f = [256, 128, 64, 64] + self.decoder = nn.Sequential( + Clamp(), conv(self.latent_channels, n_f[0]), nn.ReLU(inplace=True), + MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), MemBlock(n_f[0], n_f[0]), nn.Upsample(scale_factor=2 if decoder_space_upscale[0] else 1), TGrow(n_f[0], 2 if decoder_time_upscale[0] else 1), conv(n_f[0], n_f[1], bias=False), + MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), MemBlock(n_f[1], n_f[1]), nn.Upsample(scale_factor=2 if decoder_space_upscale[1] else 1), TGrow(n_f[1], 2 if decoder_time_upscale[1] else 1), conv(n_f[1], n_f[2], bias=False), + MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), MemBlock(n_f[2], n_f[2]), nn.Upsample(scale_factor=2 if decoder_space_upscale[2] else 1), TGrow(n_f[2], 2 if decoder_time_upscale[2] else 1), conv(n_f[2], n_f[3], bias=False), + nn.ReLU(inplace=True), conv(n_f[3], self.image_channels*self.patch_size**2), + ) + # computed properties + self.t_downscale = 2**sum(t.stride == 2 for t in self.encoder if isinstance(t, TPool)) + self.t_upscale = 2**sum(t.stride == 2 for t in self.decoder if isinstance(t, TGrow)) + self.frames_to_trim = self.t_upscale - 1 + + if checkpoint_path is not None: + self.load_state_dict(self.patch_tgrow_layers(torch.load(checkpoint_path, map_location="cpu", weights_only=True))) + + def patch_tgrow_layers(self, sd): + """Patch TGrow layers to use a smaller kernel if needed. + + Args: + sd: state dict to patch + """ + new_sd = self.state_dict() + for i, layer in enumerate(self.decoder): + if isinstance(layer, TGrow): + key = f"decoder.{i}.conv.weight" + if sd[key].shape[0] > new_sd[key].shape[0]: + # take the last-timestep output channels + sd[key] = sd[key][-new_sd[key].shape[0]:] + return sd + + def preprocess_input_frames(self, x): + """Preprocess RGB input frames prior to the main encoder sequence.""" + if self.patch_size > 1: x = F.pixel_unshuffle(x, self.patch_size) + return x + + def encode_video(self, x, parallel=True, show_progress_bar=True): + """Encode a sequence of frames. + + Args: + x: input NTCHW RGB (C=3) tensor with values in [0, 1]. + parallel: if True, all frames will be processed at once. + (this is faster but may require more memory). + if False, frames will be processed sequentially. + Returns NTCHW latent tensor with ~Gaussian values. + """ + x = self.preprocess_input_frames(x) + if x.shape[1] % self.t_downscale != 0: + # pad at end to multiple of self.t_downscale + n_pad = self.t_downscale - x.shape[1] % self.t_downscale + padding = x[:, -1:].repeat_interleave(n_pad, dim=1) + x = torch.cat([x, padding], 1) + return apply_model_with_memblocks(self.encoder, x, parallel, show_progress_bar) + + def postprocess_output_frames(self, x): + """Postprocess RGB frames after the main decoder sequence.""" + if self.patch_size > 1: x = F.pixel_shuffle(x, self.patch_size) + return x.clamp_(0, 1) + + def decode_video(self, x, parallel=True, show_progress_bar=True): + """Decode a sequence of frames. + + Args: + x: input NTCHW latent (C=self.latent_channels) tensor with ~Gaussian values. + parallel: if True, all frames will be processed at once. + (this is faster but may require more memory). + if False, frames will be processed sequentially. + Returns NTCHW RGB tensor with ~[0, 1] values. + """ + skip_trim = self.is_cogvideox and x.shape[1] % 2 == 0 + x = apply_model_with_memblocks(self.decoder, x, parallel, show_progress_bar) + x = self.postprocess_output_frames(x) + if skip_trim: + # skip trimming for cogvideox to make frame counts match. + # this still doesn't have correct temporal alignment for certain frame counts + # (cogvideox seems to pad at the start?), but for multiple-of-4 it's fine. + return x + return x[:, self.frames_to_trim:] + +class StreamingTAEHV(nn.Module): + def __init__(self, taehv): + """Streaming wrapper around TAEHV for real-time use-cases (where not all inputs are available immediately). + + Encode-decode (video-to-video) usage: + streaming = StreamingTAEHV(taehv) + for frame in video_frames: + latent = streaming.encode(frame_tensor) + decoded = streaming.decode(latent) # feeds latent if not None, then returns next frame + if decoded is not None: + display(decoded) + for frame in streaming.flush(): + display(frame) + + Decode-only (world model) usage: + streaming = StreamingTAEHV(taehv) + while running: + latent = world_model.step() # latent represents t_upscale frames + frame = streaming.decode(latent) # returns first frame immediately + while frame is not None: # retrieve remaining frames from this latent + display(frame) + frame = streaming.decode() + """ + super().__init__() + self.taehv = taehv + self.reset() + + def reset(self): + """Reset all internal state. Call this to start encoding/decoding a new stream.""" + self.encoder_work_queue, self.encoder_memory = [], [None] * len(self.taehv.encoder) + self.decoder_work_queue, self.decoder_memory = [], [None] * len(self.taehv.decoder) + self.n_frames_encoded, self.n_frames_decoded = 0, 0 + self._last_encoder_input_frame = None + + def encode(self, x=None): + """Feed an input frame (optional) and try to produce an encoder output. + + The encoder accumulates t_downscale input frames before producing one latent, + so most calls will return None. Use flush_encoder() at end-of-stream to pad and + drain any remaining latents. + + Args: + x: NTCHW RGB frame tensor with values in [0, 1], or None to just process pending work. + Returns: N1CHW latent tensor, or None if not enough input has been accumulated. + """ + if x is not None: + assert x.ndim == 5 and x.shape[2] == self.taehv.image_channels, f"Expected NTCHW frames but got {x.shape=}" + self._last_encoder_input_frame = x[:, -1:] + x = self.taehv.preprocess_input_frames(x) + self.encoder_work_queue.extend(TWorkItem(xt, 0) for xt in x.unbind(1)) + self.n_frames_encoded += x.shape[1] + xt = apply_model_with_memblocks_sequential_single_step( + self.taehv.encoder, self.encoder_memory, self.encoder_work_queue) + return xt + + def decode(self, x=None): + """Feed a latent (optional) and try to produce a decoded frame. + + Each latent produces t_upscale output frames due to temporal upscaling. The first + decode(latent) call returns the first of these frames; call decode() with no argument + to retrieve the rest, one at a time. Each call does the minimum decoder work needed to + produce one frame. + + Startup frames (the first frames_to_trim raw decoder outputs, used for causal alignment + with the reference VAE) are consumed internally and never returned. + + Args: + x: NTCHW latent tensor, or None to retrieve the next pending frame. + Returns: N1CHW decoded RGB frame tensor, or None if the queue needs more input. + """ + if x is not None: + assert x.ndim == 5 and x.shape[2] == self.taehv.latent_channels, f"Expected NTCHW latents but got {x.shape=}" + self.decoder_work_queue.extend(TWorkItem(xt, 0) for xt in x.unbind(1)) + while True: + xt = apply_model_with_memblocks_sequential_single_step( + self.taehv.decoder, self.decoder_memory, self.decoder_work_queue) + if xt is None: + return None + self.n_frames_decoded += 1 + # skip startup frames (to match decode_video trim behavior) + if not self.taehv.is_cogvideox and self.n_frames_decoded <= self.taehv.frames_to_trim: + continue + return self.taehv.postprocess_output_frames(xt) + + def flush_encoder(self): + """Pad (if needed) and drain all remaining latents from the encoder. + + Returns list of N1CHW latent tensors. + """ + latents = [] + if self._last_encoder_input_frame is not None and self.n_frames_encoded % self.taehv.t_downscale != 0: + n_pad = self.taehv.t_downscale - self.n_frames_encoded % self.taehv.t_downscale + for _ in range(n_pad): + lat = self.encode(self._last_encoder_input_frame) + if lat is not None: + latents.append(lat) + while (lat := self.encode()) is not None: + latents.append(lat) + return latents + + def flush_decoder(self): + """Drain all remaining decoded frames from the decoder. + + Returns list of N1CHW decoded RGB frame tensors. + """ + frames = [] + while (frame := self.decode()) is not None: + frames.append(frame) + return frames + + def flush(self): + """Flush encoder (with padding) and decoder, returning all remaining decoded frames. + + Returns list of N1CHW decoded RGB frame tensors. + """ + frames = [] + for latent in self.flush_encoder(): + frame = self.decode(latent) + if frame is not None: + frames.append(frame) + frames.extend(self.flush_decoder()) + return frames + +@torch.no_grad() +def main(): + """Run TAEHV roundtrip reconstruction on the given video paths.""" + import os + import sys + import cv2 # no highly esteemed deed is commemorated here + + class VideoTensorReader: + def __init__(self, video_file_path): + self.cap = cv2.VideoCapture(video_file_path) + assert self.cap.isOpened(), f"Could not load {video_file_path}" + self.fps = self.cap.get(cv2.CAP_PROP_FPS) + def __iter__(self): + return self + def __next__(self): + ret, frame = self.cap.read() + if not ret: + self.cap.release() + raise StopIteration # End of video or error + return torch.from_numpy(cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)).permute(2, 0, 1) # BGR HWC -> RGB CHW + + class VideoTensorWriter: + def __init__(self, video_file_path, width_height, fps=30): + self.writer = cv2.VideoWriter(video_file_path, cv2.VideoWriter_fourcc(*'mp4v'), fps, width_height) + assert self.writer.isOpened(), f"Could not create writer for {video_file_path}" + def write(self, frame_tensor): + assert frame_tensor.ndim == 3 and frame_tensor.shape[0] == 3, f"{frame_tensor.shape}??" + self.writer.write(cv2.cvtColor(frame_tensor.permute(1, 2, 0).numpy(), cv2.COLOR_RGB2BGR)) # RGB CHW -> BGR HWC + def __del__(self): + if hasattr(self, 'writer'): self.writer.release() + + dev = torch.device("cuda" if torch.cuda.is_available() else "mps" if torch.backends.mps.is_available() else "cpu") + dtype = torch.float16 + checkpoint_path = os.getenv("TAEHV_CHECKPOINT_PATH", "taehv.pth") + checkpoint_name = os.path.splitext(os.path.basename(checkpoint_path))[0] + print(f"Using device \033[31m{dev}\033[0m, dtype \033[32m{dtype}\033[0m, checkpoint \033[34m{checkpoint_name}\033[0m ({checkpoint_path})") + taehv = TAEHV(checkpoint_path=checkpoint_path).to(dev, dtype) + for video_path in sys.argv[1:]: + print(f"Processing {video_path}...") + video_in = VideoTensorReader(video_path) + video = torch.stack(list(video_in), 0)[None] + vid_dev = video.to(dev, dtype).div_(255.0) + # convert to device tensor + if video.numel() < 100_000_000: + print(f" {video_path} seems small enough, will process all frames in parallel") + # convert to device tensor + vid_enc = taehv.encode_video(vid_dev) + print(f" Encoded {video_path} -> {vid_enc.shape}. Decoding...") + vid_dec = taehv.decode_video(vid_enc) + print(f" Decoded {video_path} -> {vid_dec.shape}") + else: + print(f" {video_path} seems large, will process each frame sequentially") + # convert to device tensor + vid_enc = taehv.encode_video(vid_dev, parallel=False) + print(f" Encoded {video_path} -> {vid_enc.shape}. Decoding...") + vid_dec = taehv.decode_video(vid_enc, parallel=False) + print(f" Decoded {video_path} -> {vid_dec.shape}") + video_out_path = video_path + f".reconstructed_by_{checkpoint_name}.mp4" + video_out = VideoTensorWriter(video_out_path, (vid_dec.shape[-1], vid_dec.shape[-2]), fps=int(round(video_in.fps))) + for frame in vid_dec.clamp_(0, 1).mul_(255).round_().byte().cpu()[0]: + video_out.write(frame) + print(f" Saved to {video_out_path}") + +if __name__ == "__main__": + main() diff --git a/mkdocs.yml b/mkdocs.yml index e47f2fe4b..caf72e898 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -151,6 +151,7 @@ nav: - Home: index.md - Getting Started: - Installation: getting_started/installation.md + - Apple Silicon FastWan: getting_started/installation/mps.md - Quick Start: getting_started/quick_start.md - V1 API: getting_started/v1_api.md - Inference: diff --git a/pyproject.toml b/pyproject.toml index d7630ba34..abf075ce0 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -201,6 +201,12 @@ rocm = [ "amdsmi", ] +# Apple Silicon FastWan release runtime. The range is intentionally pinned to +# the MLX version used for the checked-in parity and checkpoint tests. +mlx = [ + "mlx>=0.31.2,<0.33; platform_system == 'Darwin' and platform_machine == 'arm64'", +] + [project.scripts] fastvideo = "fastvideo.entrypoints.cli.main:main" dreamverse-server = "dreamverse.server_entry:cli" @@ -212,6 +218,8 @@ exclude = ["assets*", "docker*", "docs", "scripts*", "apps*", "dreamverse.tests* [tool.setuptools.package-data] dreamverse = ["prompts/*.md"] +"fastvideo.third_party.rife_mlx" = ["LICENSE"] +"fastvideo.third_party.taehv" = ["LICENSE"] [tool.wheel] exclude = ["assets*", "docker*", "docs", "scripts*", "apps*"]