Skip to content

Commit 8537dcd

Browse files
authored
[feat]: Apple Silicon MLX runtime — INT8 Wan2.1 and Wan2.2 inference (#1638)
Co-authored-by: Aryan Kumar <aryan5v@users.noreply.github.com>
1 parent 8208536 commit 8537dcd

78 files changed

Lines changed: 13180 additions & 5 deletions

File tree

Some content is hidden

Large Commits have some content hidden by default. Use the searchbox below for content that may be hidden.

.github/workflows/ci-macos-mlx.yml

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
name: macOS MLX Smoke
2+
3+
on:
4+
pull_request:
5+
branches: [main]
6+
paths:
7+
- ".github/workflows/ci-macos-mlx.yml"
8+
- "fastvideo/mlx_runtime/**"
9+
- "fastvideo/tests/mlx/**"
10+
- "fastvideo/tests/platforms/test_mps_vsa_error.py"
11+
- "fastvideo/platforms/mps.py"
12+
- "fastvideo/platforms/__init__.py"
13+
- "fastvideo/__init__.py"
14+
- "examples/inference/basic/mlx_*.py"
15+
- "fastvideo/benchmarks/mlx_*.py"
16+
- "pyproject.toml"
17+
workflow_dispatch:
18+
19+
permissions:
20+
contents: read
21+
22+
concurrency:
23+
group: macos-mlx-${{ github.ref }}
24+
cancel-in-progress: true
25+
26+
jobs:
27+
mlx-smoke:
28+
if: github.event_name == 'workflow_dispatch' || github.event.pull_request.draft != true
29+
runs-on: macos-15
30+
timeout-minutes: 25
31+
env:
32+
FASTVIDEO_ATTENTION_BACKEND: TORCH_SDPA
33+
TOKENIZERS_PARALLELISM: "false"
34+
MASTER_ADDR: localhost
35+
MASTER_PORT: "29513"
36+
steps:
37+
- uses: actions/checkout@v4
38+
39+
- uses: actions/setup-python@v5
40+
with:
41+
python-version: "3.12"
42+
cache: pip
43+
44+
- uses: astral-sh/setup-uv@v3
45+
46+
- name: Install lightweight MLX smoke dependencies
47+
run: |
48+
uv pip install --system \
49+
--index-url https://download.pytorch.org/whl/cpu \
50+
torch==2.11.0 torchvision torchaudio
51+
uv pip install --system \
52+
pytest numpy scipy pillow imageio einops cloudpickle filelock \
53+
PyYAML diffusers huggingface_hub remote-pdb safetensors loguru mlx \
54+
"ftfy>=6.3.1" "opencv-python>=4.10.0.84" psutil "transformers>=5.0.0"
55+
56+
- name: Show Apple runtime
57+
run: |
58+
python - <<'PY'
59+
import platform
60+
import mlx.core as mx
61+
import torch
62+
63+
print("machine:", platform.machine())
64+
print("processor:", platform.processor())
65+
print("mlx default device:", mx.default_device())
66+
memory_size = mx.metal.device_info().get("memory_size") if mx.metal.is_available() else "metal unavailable"
67+
print("mlx memory_size:", memory_size)
68+
print("torch:", torch.__version__)
69+
print("torch mps available:", torch.backends.mps.is_available())
70+
PY
71+
72+
- name: Run MLX smoke tests
73+
run: |
74+
python -m pytest \
75+
fastvideo/tests/mlx/test_dmd_sampling.py \
76+
fastvideo/tests/mlx/test_memory_limits.py \
77+
fastvideo/tests/mlx/test_quant_capability.py \
78+
fastvideo/tests/mlx/test_mlx_dit_parity.py \
79+
fastvideo/tests/mlx/test_mlx_compile_parity.py \
80+
fastvideo/tests/mlx/test_mlx_checkpoint.py \
81+
fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \
82+
fastvideo/tests/mlx/test_taehv_decode.py \
83+
fastvideo/tests/mlx/test_frame_upsample.py \
84+
fastvideo/tests/mlx/test_mlx_fast_spatial.py \
85+
fastvideo/tests/mlx/test_mlx_refine.py \
86+
fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \
87+
fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \
88+
fastvideo/tests/mlx/test_wan22_sample.py \
89+
fastvideo/tests/mlx/test_windowed_attention.py \
90+
fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_download_unavailable_has_specific_error \
91+
fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_backend_regression_is_not_skip_eligible \
92+
fastvideo/tests/platforms/test_mps_vsa_error.py \
93+
-q
94+
95+
# Same tests on MLX's CPU backend. Hosted macOS runners are scarce and
96+
# slower to schedule; this Linux job gives fast PR signal on the identical
97+
# graph (the parity tests were designed to be backend-agnostic), while the
98+
# macOS job above stays the source of truth for Metal behavior.
99+
mlx-smoke-linux-cpu:
100+
if: github.event_name == 'workflow_dispatch' || github.event.pull_request.draft != true
101+
runs-on: ubuntu-latest
102+
timeout-minutes: 20
103+
env:
104+
FASTVIDEO_ATTENTION_BACKEND: TORCH_SDPA
105+
TOKENIZERS_PARALLELISM: "false"
106+
MASTER_ADDR: localhost
107+
MASTER_PORT: "29513"
108+
steps:
109+
- uses: actions/checkout@v4
110+
111+
- uses: actions/setup-python@v5
112+
with:
113+
python-version: "3.12"
114+
cache: pip
115+
116+
- uses: astral-sh/setup-uv@v3
117+
118+
- name: Install lightweight MLX smoke dependencies (CPU backend)
119+
run: |
120+
uv pip install --system \
121+
--index-url https://download.pytorch.org/whl/cpu \
122+
torch==2.11.0 torchvision torchaudio
123+
uv pip install --system \
124+
pytest numpy scipy pillow imageio einops cloudpickle filelock \
125+
PyYAML diffusers huggingface_hub remote-pdb safetensors loguru "mlx[cpu]" \
126+
"ftfy>=6.3.1" "opencv-python>=4.10.0.84" psutil "transformers>=5.0.0"
127+
128+
- name: Run MLX smoke tests (CPU backend)
129+
run: |
130+
python -m pytest \
131+
fastvideo/tests/mlx/test_dmd_sampling.py \
132+
fastvideo/tests/mlx/test_memory_limits.py \
133+
fastvideo/tests/mlx/test_quant_capability.py \
134+
fastvideo/tests/mlx/test_mlx_dit_parity.py \
135+
fastvideo/tests/mlx/test_mlx_compile_parity.py \
136+
fastvideo/tests/mlx/test_mlx_checkpoint.py \
137+
fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \
138+
fastvideo/tests/mlx/test_taehv_decode.py \
139+
fastvideo/tests/mlx/test_frame_upsample.py \
140+
fastvideo/tests/mlx/test_mlx_fast_spatial.py \
141+
fastvideo/tests/mlx/test_mlx_refine.py \
142+
fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \
143+
fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \
144+
fastvideo/tests/mlx/test_wan22_sample.py \
145+
fastvideo/tests/mlx/test_windowed_attention.py \
146+
fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_download_unavailable_has_specific_error \
147+
fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_backend_regression_is_not_skip_eligible \
148+
fastvideo/tests/platforms/test_mps_vsa_error.py \
149+
-q

README.md

Lines changed: 5 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -62,6 +62,11 @@ UV_TORCH_BACKEND=cu126 uv pip install fastvideo
6262
Use `UV_TORCH_BACKEND=cu130` on CUDA 13. Apple silicon users should follow the
6363
[MPS installation guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
6464

65+
> **On an Apple Silicon Mac?** FastVideo runs FastWan text-to-video natively
66+
> through an MLX runtime — a 5-second 480p clip generated locally, no cloud,
67+
> no discrete GPU. Install with `uv pip install -e '.[mlx]'` and follow the
68+
> [Apple Silicon guide](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/mps/).
69+
6570
Please see our [docs](https://hao-ai-lab.github.io/FastVideo/getting_started/installation/) for more detailed installation instructions.
6671

6772
> **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/).
Lines changed: 128 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,128 @@
1+
# Fast mode (RIFE) — Apple Silicon
2+
3+
`--fast` makes local generation ~2.7× faster by **generating fewer frames and
4+
interpolating the rest** with an Apple-Silicon-native RIFE model, instead of
5+
denoising every frame. Video-diffusion denoise is dominated by self-attention,
6+
which is O(tokens²); halving the frames cuts the token count ~2× and the denoise
7+
compute ~3.7×, so the wall-clock drops far more than 2×. RIFE (which estimates
8+
its own optical flow — no motion vectors needed) fills the dropped frames back
9+
in for ~1.4 s, and a light unsharp pass counters its softening.
10+
11+
Measured on the 1.3B INT8 QAD model (fox, 480×832×81, M4): generate 41 + RIFE→81
12+
runs in ~35 s of denoise vs ~90 s full, at reconstruction MS-SSIM **0.97**.
13+
Reproduce with `python -m fastvideo.benchmarks.eval_metalfx_rife --mode int8`.
14+
15+
> **Note:** Apple's *MetalFX* frame interpolation is **not** usable here — it
16+
> requires game-engine motion vectors + depth, which diffusion output lacks. We
17+
> use the video-native **`rife-mlx`** model instead (Metal-backed, torch-free).
18+
19+
## Install
20+
21+
```bash
22+
uv pip install -e ".[mlx]" # RIFE ships vendored; this only needs MLX
23+
```
24+
25+
## Use
26+
27+
```bash
28+
python examples/inference/basic/mlx_wan_prompt_to_video.py \
29+
--mlx-checkpoint <FastWan2.1-T2V-1.3B-INT8-QAD> \
30+
--prompt "A red fox trotting through a snowy pine forest at golden hour, cinematic" \
31+
--num-frames 81 --fast \
32+
--output-path video_samples/fox_fast.mp4
33+
```
34+
35+
`--num-frames` stays the *target* length; fast mode generates the smallest
36+
VAE-aligned keyframe count that RIFE can interpolate to that target.
37+
38+
| Flag | Default | Meaning |
39+
|---|---|---|
40+
| `--fast` / `--no-fast` | off | enable fast mode |
41+
| `--fast-factor` | 2 | generate 1/factor of the frames (2 = half) |
42+
| `--fast-sharpen` | 0.6 | light unsharp strength to counter RIFE softness (0 disables) |
43+
44+
Fast mode composes with everything else (`--mlx-quantization int8`,
45+
`--mlx-compile`, TAEHV vs `--decode-backend wan-vae`). Keep `--fast-factor` at 2
46+
for quality — larger temporal gaps are where RIFE starts inventing motion.
47+
48+
## Spatial fast mode (`--fast-spatial`)
49+
50+
The spatial twin of `--fast`: instead of dropping frames, drop pixels. Denoise
51+
*and decode* at `height/width // fast-spatial-scale`, then resample the decoded
52+
frames up to the requested size. Self-attention is O(tokens²), so halving each
53+
spatial axis cuts the token count 4× and the denoise time far more than that —
54+
measured on the 1.3B INT8 QAD model at 480×832×81, M4 Max: **86.1 s → 10.3 s**
55+
of denoise. It composes with `--fast`; both together run the same clip in
56+
**4.5 s** of denoise.
57+
58+
```bash
59+
python examples/inference/basic/mlx_wan_prompt_to_video.py \
60+
--prompt "A red fox trotting through a snowy pine forest at golden hour, cinematic" \
61+
--height 480 --width 832 --num-frames 81 --fast-spatial \
62+
--output-path video_samples/fox_fast_spatial.mp4
63+
```
64+
65+
| Flag | Default | Meaning |
66+
|---|---|---|
67+
| `--fast-spatial` / `--no-fast-spatial` | off | enable spatial fast mode |
68+
| `--fast-spatial-scale` | 2 | denoise at 1/scale of each spatial axis |
69+
| `--fast-spatial-upsample-mode` | `lanczos` | pixel interpolation kernel (`lanczos`, `cubic`, `bilinear`, `nearest`) |
70+
| `--fast-spatial-sharpen` | 0.4 | light unsharp strength to counter resampling softness (0 disables) |
71+
72+
### The upsample must happen in pixel space
73+
74+
This is the one thing to get right. The obvious implementation — bilinearly
75+
upsample the finished latents and decode at the target size — **does not work**,
76+
and produces a distinctive failure: correct composition and silhouette under a
77+
smeared, hazy veil, with ringing along strong edges.
78+
79+
A Wan latent cell is a *learned code* for an 8×8 (Wan2.1) or 16×16 (Wan2.2)
80+
pixel block, not a low-pass sample of the image. The average of two adjacent
81+
codes is not the code of the averaged blocks; it is a vector the decoder was
82+
never trained on. Measured on Wan2.1-1.3B at 480×832, a 2× bilinear latent
83+
upsample destroys **62%** of the latent's high-frequency energy while leaving
84+
its overall magnitude intact — exactly the signature of that veil. At Wan2.2-5B
85+
the same operation degrades to black or noise.
86+
87+
Decoded RGB frames have no such problem: an image *is* a sampled 2-D signal, so
88+
Lanczos interpolation is the operation it was defined for. The result is soft —
89+
it carries stage-1's real detail budget and no more — but clean and coherent.
90+
91+
`--refine` gets away with a latent-space upsample only because a second DMD pass
92+
re-denoises the hand-off; spatial fast mode passes the latent straight to the
93+
decoder, so it cannot.
94+
95+
## Refine (`--refine`) stage-2 timesteps
96+
97+
`--refine` hands stage 1 to stage 2 as `(1 - sigma) * upsampled + sigma * noise`,
98+
where `sigma` comes from the *first* stage-2 timestep. FastWan's DMD grid opens
99+
at `t=1000`, which is `sigma == 1` exactly — so a stage-2 grid that starts there
100+
weights the stage-1 result at zero and refine silently degrades into a plain
101+
full-resolution run at twice the cost.
102+
103+
Left unset, `--refine-dmd-denoising-steps` now derives the stage-2 grid from the
104+
stage-1 one with leading full-noise steps dropped (`1000,757,522``757,522`).
105+
That keeps the pass on timesteps the distilled student was trained on while
106+
letting stage-1 structure through: hand-off `sigma = 0.757`, stage-1 weight
107+
`0.243`. Passing a grid that starts at full noise is now an error rather than a
108+
silently wasted pass.
109+
110+
The run prints the resolved hand-off so it is visible:
111+
112+
```
113+
[refine] stage-2 hand-off sigma=0.7568 (stage-1 weight 0.2432)
114+
```
115+
116+
There is a trade-off in choosing that grid. Later start = more of the draft
117+
survives, but fewer stage-2 steps. On Wan2.1 the default `757,522` gives weight
118+
0.243 with two steps; `--refine-dmd-denoising-steps 522` gives weight 0.478 with
119+
one. `--refine-sigma` decouples the noise level from the timestep entirely — it
120+
logs a warning, because the DiT is then told a timestep that does not match the
121+
noise it receives.
122+
123+
**Wan2.2-5B has a lower ceiling.** Its warped schedule maps `1000,757,522` to
124+
sigmas `1.000, 0.940, 0.845`, so the best available stage-1 weight is **0.060**
125+
(vs 0.243 at 1.3B). Un-warped (`--no-warp`) the same grid gives `1.000, 0.757,
126+
0.522` and a weight of 0.243 — but warping is what matches the FastVideo
127+
sampling schedule, so turning it off changes the timesteps the distilled student
128+
sees. Which is better at 5B is unresolved and needs a run on real 5B weights.

docs/getting_started/installation/mps.md

Lines changed: 6 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -49,18 +49,20 @@ brew install ffmpeg
4949

5050
### Installation
5151

52+
FastWan's native Apple Silicon runtime requires the `mlx` extra.
53+
5254
#### With uv (recommended)
5355

5456
```bash
55-
uv pip install fastvideo
57+
uv pip install "fastvideo[mlx]"
5658
```
5759

5860
#### With Conda environment (alternative)
5961

6062
`uv` works inside an active conda env too, so prefer `uv pip` for the actual install:
6163

6264
```bash
63-
uv pip install fastvideo
65+
uv pip install "fastvideo[mlx]"
6466
```
6567

6668
### Installation from Source
@@ -76,13 +78,13 @@ git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo
7678
Basic installation:
7779

7880
```bash
79-
uv pip install -e .
81+
uv pip install -e ".[mlx]"
8082
```
8183

8284
Alternative with Conda environment:
8385

8486
```bash
85-
uv pip install -e .
87+
uv pip install -e ".[mlx]"
8688
```
8789

8890
## Development Environment Setup

docs/inference/support_matrix.md

Lines changed: 11 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -178,6 +178,17 @@ optimizations: absence means **untested**, not incompatible.
178178
| Matrix Game 3.0 Base Distilled | `FastVideo/Matrix-Game-3.0-Base-Distilled-Diffusers` | 720x1280 ||||||
179179
| GEN3C Cosmos 7B | `FastVideo/GEN3C-Cosmos-7B-Diffusers` | 704px1280p ||||||
180180

181+
## Apple Silicon native runtime
182+
183+
| Release path | Model | Mode | Validated hardware | Status |
184+
| --- | --- | --- | --- | --- |
185+
| 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 |
186+
187+
This is a text-to-video-only source-install release. It is validated on the
188+
hardware listed above; MLX allocator caps are not evidence of support for a
189+
physical 16 GB Mac. See [Apple Silicon FastWan](../getting_started/installation/mps.md)
190+
for the supported command and release gates.
191+
181192
**Note**: Wan2.2 TI2V 5B has some quality issues when performing I2V generation. We are working on fixing this issue.
182193

183194
***Lucy Edit Dev uses a non-commercial model license. FastVideo support is
Lines changed: 46 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,46 @@
1+
# SPDX-License-Identifier: Apache-2.0
2+
"""Tiny MLX RIFE frame-interpolation smoke test."""
3+
4+
from __future__ import annotations
5+
6+
import argparse
7+
import time
8+
9+
import numpy as np
10+
11+
from fastvideo.mlx_runtime.rife_interp import interpolate, load_model
12+
13+
14+
def main() -> None:
15+
parser = argparse.ArgumentParser(
16+
description="MLX RIFE 4.25 frame interpolation smoke test."
17+
)
18+
parser.add_argument(
19+
"--self-test",
20+
action="store_true",
21+
help="Run a tiny two-frame interpolation test.",
22+
)
23+
args = parser.parse_args()
24+
if not args.self_test:
25+
raise SystemExit("Nothing to do; pass --self-test")
26+
27+
frame0 = np.zeros((64, 96, 3), dtype=np.uint8)
28+
frame1 = np.zeros((64, 96, 3), dtype=np.uint8)
29+
frame1[:, :, 0] = 255
30+
start = time.perf_counter()
31+
model = load_model()
32+
load_s = time.perf_counter() - start
33+
start = time.perf_counter()
34+
frames = interpolate([frame0, frame1], factor=2, model=model)
35+
interp_s = time.perf_counter() - start
36+
assert len(frames) == 3
37+
assert frames[1].shape == frame0.shape
38+
assert frames[1].dtype == np.uint8
39+
print(
40+
"MLX RIFE self-test passed: "
41+
f"load_s={load_s:.3f} interp_s={interp_s:.3f} shape={frames[1].shape}"
42+
)
43+
44+
45+
if __name__ == "__main__":
46+
main()

0 commit comments

Comments
 (0)