diff --git a/.github/workflows/ci-macos-mlx.yml b/.github/workflows/ci-macos-mlx.yml new file mode 100644 index 0000000000..0be4cda846 --- /dev/null +++ b/.github/workflows/ci-macos-mlx.yml @@ -0,0 +1,124 @@ +name: macOS MLX Smoke + +on: + pull_request: + branches: [main] + paths: + - ".github/workflows/ci-macos-mlx.yml" + - "fastvideo/mlx_runtime/**" + - "fastvideo/tests/mlx/**" + - "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 + + - 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_checkpoint.py \ + fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \ + fastvideo/tests/mlx/test_taehv_decode.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]" + + - 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_checkpoint.py \ + fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py \ + fastvideo/tests/mlx/test_taehv_decode.py \ + -q diff --git a/docs/design/apple_silicon_qad_blog_draft.md b/docs/design/apple_silicon_qad_blog_draft.md new file mode 100644 index 0000000000..26f7eecb00 --- /dev/null +++ b/docs/design/apple_silicon_qad_blog_draft.md @@ -0,0 +1,50 @@ +# FastWan-QAD-INT8-1.3B: Local Text-to-Video on Apple Silicon + +> Draft — do not publish until the release owner selects raw or EMA from the +> final visual review and the model-card checksums are available. + +## TL;DR + +FastVideo's Apple-native FastWan runtime generates a five-second 480p-class +text-to-video clip locally on an Apple M4 Max. The release candidate uses an +MLX INT8 DiT, a three-step DMD schedule, and TAEHV decode. The recorded +480x832x81 result is 123.7 seconds end to end, including 117.6 seconds of +denoising, at 5.63 GiB MLX peak memory. + +This is a source-install, text-to-video release for the validated M4 Max, +36 GB-class configuration. It is not a general 16 GB Mac claim, and it does +not include image-to-video. + +## What ships + +- FastWan-QAD-INT8-1.3B in Diffusers form plus a pre-quantized MLX DiT + checkpoint that avoids requantizing weights at startup. +- One source-tree generation command using MLX for denoising and MPS for text + encoding and TAEHV decode. +- A reproducible benchmark harness and Metal/MLX-CPU runtime tests. + +TAEHV is vendored under its MIT license; FastVideo source is Apache-2.0. + +## Results and release decision + +The candidate raw and EMA checkpoints measure respectively 0.9360 and 0.9331 +mean MS-SSIM when each INT8 result is compared with the same model's FP16 +result. That metric measures quantization consistency only. It does not score +absolute visual quality, so it cannot select a release checkpoint. + +**Release owner decision: [TODO: raw or EMA after final motion7 visual review].** + +The final post must include the selected revision, checksums, model-card link, +fixed generation command, and the reviewed visual examples. It must not claim +bitwise fake-quant parity, use the invalid 0.9860 run-1 EMA number, or claim +that 16 GB Macs have been validated. + +## Run it + +```console +uv pip install -e '.[mlx]' +python examples/inference/basic/mlx_wan_prompt_to_video.py [release command from the model card] +``` + +See the Apple Silicon FastWan guide for the full command, hardware statement, +and troubleshooting. diff --git a/docs/design/apple_silicon_qad_runbook.md b/docs/design/apple_silicon_qad_runbook.md new file mode 100644 index 0000000000..a09312377d --- /dev/null +++ b/docs/design/apple_silicon_qad_runbook.md @@ -0,0 +1,162 @@ +# Runbook: Wan2.1-1.3B → 3-step INT8 QAD on the DGX B200 + +Operator instructions for launching the Mac-targeted quantization-aware DMD +distillation (roadmap M4 Phase B) on a DGX with B200 GPUs. Everything below +runs from a clone of `aryan5v/FastVideo` on branch +the approved future QAD integration branch. The training recipe is +`examples/train/configs/distribution_matching/wan/dmd2_t2v_mlx_int8.yaml`: +frozen Wan2.1-T2V-1.3B teacher + critic, trainable student whose linear +weights are fake-quantized every forward onto MLX's affine INT8 deploy +grid (the `mlx_qat` callback; serialized quantizer state is pinned exactly and +the Metal/PyTorch forward representation is bounded to one source-dtype epsilon by +`fastvideo/tests/mlx/test_mlx_affine_qat_parity.py`), distilled to the +3-step FastWan schedule `[1000, 757, 522]`. + +Use **4 GPUs** (the recipe default) — do not grab all 8 unless told to. +Expected wall time for the full run is roughly 4–8 hours on 4×B200, but run +the smoke test first and extrapolate from its measured seconds/step. + +**If only N (< 4) GPUs are free:** the run works on any count ≥ 1; the only +rule is that the HSDP shard dim must equal the GPU count. Add these overrides +to every launch command below (shown for N=3): + +```bash +NUM_GPUS=3 bash examples/train/run.sh \ + --training.distributed.num_gpus 3 \ + --training.distributed.hsdp_shard_dim 3 \ + ... +``` + +Global batch scales with GPU count (1 per GPU), which is fine for DMD at this +scale; wall time scales inversely (~6–11 h on 3 GPUs). Prefer waiting for a +4th GPU only if it frees up within the hour; otherwise just run on 3. The +same rule applies upward: if 8 are genuinely free and idle, `NUM_GPUS=8` with +`hsdp_shard_dim 8` roughly halves the wall time. + +## 0. Preflight + +```bash +nvidia-smi # expect B200s; confirm >= 4 idle +python3 --version # 3.10–3.12 +``` + +Verify credentials are present (both are expected to be preconfigured on the +box — do not write them into the repo): + +```bash +test -n "$WANDB_API_KEY" && echo "wandb ok" +hf auth whoami || huggingface-cli whoami # HF auth for model + dataset pulls +``` + +## 1. Clone and install + +```bash +git clone https://github.com/aryan5v/FastVideo.git && cd FastVideo +git checkout aryan/future/fastwan-qad-5b-i2v +uv venv --python 3.12 && source .venv/bin/activate +uv pip install -e ".[dev,mlx]" +``` + +B200 is sm_100: torch must be a recent CUDA build (the pinned deps are). +If any attention backend fails to import or pick a kernel at startup, force +the portable one — correctness is identical, it is only somewhat slower: + +```bash +export FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA +``` + +Sanity-check the QAT machinery on this box before spending GPU time: + +```bash +pytest fastvideo/tests/mlx/test_mlx_affine_qat_parity.py -q +pytest fastvideo/tests/training/test_mlx_qat_callback.py -q +``` + +## 2. Dataset + +```bash +python scripts/huggingface/download_hf.py \ + --repo_id "FastVideo/Wan-Syn_77x448x832_600k" \ + --local_dir "data/Wan-Syn_77x448x832_600k" \ + --repo_type "dataset" +``` + +This is large (order of a terabyte) — start it in `tmux` early and check free +disk first (`df -h .`). The recipe reads it from +`data/Wan-Syn_77x448x832_600k` relative to the repo root. + +## 3. Smoke run (mandatory, ~30–60 min) + +```bash +tmux new -s qad-smoke +NUM_GPUS=4 bash examples/train/run.sh \ + examples/train/configs/distribution_matching/wan/dmd2_t2v_mlx_int8.yaml \ + --training.loop.max_train_steps 100 \ + --training.checkpoint.output_dir outputs/smoke_mlx_int8 \ + --training.tracker.run_name wan2.1_qad_int8_smoke \ + --callbacks.validation.every_steps 50 +``` + +What to verify before proceeding: + +1. The log contains a line like `mlx_qat: fake-quantizing N weights (int8, + group_size=64, ...)` with N in the hundreds — that is the QAT callback + arming. If instead it raises `mlx_qat matched no weights`, stop and + report. +2. No crash referencing `parametrizations` together with FSDP/HSDP/DTensor. + The QAT callback wraps weights with torch parametrizations after model + setup; a sharding interaction here is the one known integration risk. If + it crashes, capture the full traceback and stop — do not work around it + by removing the callback (that silently turns the run into plain DMD). +3. Loss values are finite, W&B run `wan2.1_qad_int8_smoke` is logging under + project `distillation_wan`, and the step-50/100 validation clips are not + black/NaN garbage (blurry is fine at step 100). +4. Note the steady seconds/step from the log and extrapolate: full run = + 4000 × s/step. Report that number back. + +## 4. Full run + +```bash +tmux new -s qad-full +NUM_GPUS=4 bash examples/train/run.sh \ + examples/train/configs/distribution_matching/wan/dmd2_t2v_mlx_int8.yaml \ + --callbacks.validation.every_steps 200 +``` + +Output/checkpoints land in `outputs/wan2.1_dmd2_3steps_mlx_int8` +(checkpoint every 20 steps, last 3 kept), W&B run +`wan2.1_dmd2_3steps_mlx_int8`. The job is resumable: +`--training.checkpoint.resume_from_checkpoint /checkpoint-`. +Monitor W&B; the DMD generator loss is noisy by nature — judge by the +validation clips trending sharper/more coherent, not by the loss curve alone. + +## 5. Export (1 GPU) — both raw and EMA + +Export twice: the raw student, and the EMA shadow weights (usually visibly +smoother). Both are evaluated on the Mac side; the better one ships. + +```bash +python -m fastvideo.train.entrypoint.dcp_to_diffusers \ + --checkpoint outputs/wan2.1_dmd2_3steps_mlx_int8 \ + --output-dir outputs/wan2.1_qad_int8_diffusers \ + --role student + +python -m fastvideo.train.entrypoint.dcp_to_diffusers \ + --checkpoint outputs/wan2.1_dmd2_3steps_mlx_int8 \ + --output-dir outputs/wan2.1_qad_int8_ema_diffusers \ + --role student --ema +``` + +Each auto-picks the latest checkpoint and writes a Diffusers-style model dir. +If the `--ema` export errors, report the traceback and still deliver the raw +export — it unblocks Mac evaluation while the EMA path gets fixed. + +## 6. Deliverables + +Report back: (a) the W&B run URL, (b) measured seconds/step and total wall +time, (c) the path to `outputs/wan2.1_qad_int8_diffusers`, and (d) 2–3 +validation clips from late in training. Mac-side evaluation then happens per +`docs/design/apple_silicon_fastvideo.md` (M4 exit criteria): load the export +through the MLX runtime, quantize INT8 on load (the grid the student trained +on), `--save-mlx-checkpoint`, and run the benchmark suite against the +INT8-PTQ baseline. diff --git a/docs/design/fastwan_qad_int8_1_3b_release.md b/docs/design/fastwan_qad_int8_1_3b_release.md new file mode 100644 index 0000000000..a5219d29f0 --- /dev/null +++ b/docs/design/fastwan_qad_int8_1_3b_release.md @@ -0,0 +1,39 @@ +# FastWan-QAD-INT8-1.3B Apple Release Record + +This is the committed launch-evidence and publication checklist. Do not commit +generated videos or visual grids. + +## Reproducible evidence + +| Field | Recorded result | +| --- | --- | +| Hardware | Apple M4 Max, 36 GB-class unified memory (MLX: 38.65 GB) | +| Runtime | macOS 14+, Python 3.12, MLX 0.31.2 | +| Shape | 480x832, 81 frames, 3-step DMD (`1000,757,522`) | +| Path | MPS prompt encode, MLX INT8 DiT, TAEHV decode | +| Total / denoise / peak | 123.7 s / 117.6 s / 5.63 GiB | + +QAD v2 raw and EMA measure 0.9360 and 0.9331 mean MS-SSIM against each +model's own FP16 output. This is quantization consistency, not absolute quality +and not a checkpoint-selection metric. + +## Mandatory pre-publication gate + +1. Fresh source install with `uv pip install -e '.[mlx]'`; retain passing + Metal and MLX CPU suite output. +2. Run `motion7` at the recorded shape; validate all MP4s, timings, memory, + Diffusers-to-MLX conversion, and cold pre-quantized-checkpoint reload. +3. A named release owner must visually select raw or EMA. Neither is selected + by this branch; do not publish model or blog first. +4. Do not claim 16 GB support without a separate pass on a physical 16 GB Mac. + +## Model-card and blog requirements + +The model card must identify the chosen raw/EMA checkpoint, base provenance, +exact revision, model artifacts, SHA-256 checksums, fixed generation command, +hardware/software evidence, intended T2V-only use, and limitations. It must +include FastVideo's Apache-2.0 notice and TAEHV's MIT notice. + +The launch blog may quote the M4 Max measurement. It must not claim physical +16 GB support, use the invalid 0.9860 run-1 EMA score, call fake-quant parity +bitwise, claim everything is Apache-2.0, or preselect raw versus EMA. diff --git a/docs/getting_started/installation/mps.md b/docs/getting_started/installation/mps.md index df80c38a5d..04ccaeb66d 100644 --- a/docs/getting_started/installation/mps.md +++ b/docs/getting_started/installation/mps.md @@ -1,103 +1,60 @@ -# MPS (Apple Silicon) +# Apple Silicon FastWan -Instructions to install FastVideo for Apple Silicon. +FastWan-QAD-INT8-1.3B is the Apple-native, text-to-video release candidate. +Its DiT denoising loop runs in MLX; prompt encoding and TAEHV decode use +PyTorch MPS. It does not provide image-to-video support. -## Requirements +## Validated configuration -- **OS: macOS 14 or newer** -- **Python: 3.12.4** +The recorded release result is from an Apple M4 Max with 36 GB-class unified +memory (MLX reports 38.65 GB), macOS 14+, Python 3.12, and MLX 0.31.2: +480x832, 81 frames, three-step DMD, INT8 DiT + TAEHV decode in 123.7 seconds +end to end (117.6 seconds denoise; 5.63 GiB MLX peak). -## Set up using Python +This is the only launch-supported hardware configuration. Allocator-cap tests +are useful engineering evidence, not a claim that a physical 16 GB Mac works. -### Create a new Python environment - -#### uv -Recommended default: use [uv](https://docs.astral.sh/uv/) for faster and more stable environment setup. - -Please follow the [documentation](https://docs.astral.sh/uv/#getting-started) to install `uv`. After installing `uv`, create a new environment using: +## Install from source ```console -# (Recommended) Create a new uv environment. Use `--seed` to install `pip` and `setuptools`. +brew install ffmpeg +git clone https://github.com/hao-ai-lab/FastVideo.git +cd FastVideo uv venv --python 3.12 --seed source .venv/bin/activate +uv pip install -e '.[mlx]' ``` -#### Conda (alternative) - -You can also create a Python environment using [Conda](https://docs.conda.io/projects/conda/en/stable/user-guide/getting-started.html). - -##### 1. Install Miniconda (if not already installed) - -```bash -wget https://repo.anaconda.com/miniconda/Miniconda3-latest-MacOSX-arm64.sh -bash Miniconda3-latest-MacOSX-arm64.sh -source ~/.zshrc -``` - -##### 2. Create and activate a Conda environment for FastVideo - -```bash -conda create -n fastvideo python=3.12.4 -y -conda activate fastvideo -``` - -### Dependencies +The `mlx` extra is pinned to the MLX 0.31.2 compatibility range and only +resolves on Apple Silicon. -``` -brew install ffmpeg -``` - -### Installation - -#### With uv (recommended) - -```bash -uv pip install fastvideo -``` +## Generate a video -#### With Conda environment (alternative) +After publication, download the release model and run the one supported +source-tree entrypoint: -`uv` works inside an active conda env too, so prefer `uv pip` for the actual install: - -```bash -uv pip install fastvideo -``` - -### Installation from Source - -#### 1. Clone the FastVideo repository - -```bash -git clone https://github.com/hao-ai-lab/FastVideo.git && cd FastVideo -``` - -#### 2. Install FastVideo - -Basic installation: - -```bash -uv pip install -e . -``` - -Alternative with Conda environment: +```console +huggingface-cli download FastVideo/FastWan-QAD-INT8-1.3B-Diffusers \ + --local-dir ~/models/FastWan-QAD-INT8-1.3B -```bash -uv pip install -e . +python examples/inference/basic/mlx_wan_prompt_to_video.py \ + --model-root ~/models/FastWan-QAD-INT8-1.3B \ + --mlx-checkpoint ~/models/FastWan-QAD-INT8-1.3B/mlx_dit \ + --mlx-quantization int8 --height 480 --width 832 --num-frames 81 \ + --num-inference-steps 3 --dmd-denoising-steps 1000,757,522 \ + --decode-backend taehv \ + --prompt "A fox runs through a misty pine forest, leaves kicking up behind it." \ + --output-path video_samples/fox.mp4 ``` -## Development Environment Setup - -If you're planning to contribute to FastVideo please see the following page: -[Contributor Guide](../../contributing/overview.md) - -## Hardware Requirements - -### For Basic Inference - -- Mac M1, M2, M3, or M4 (at least 32 GB RAM is preferable for high quality video generation) +The release repository must include the hybrid Diffusers components and the +verified `mlx_dit/` directory. TAEHV's checkpoint is SHA-256 verified before +use; its vendored source is MIT-licensed, while FastVideo is Apache-2.0. ## Troubleshooting -If you encounter any issues during installation, please open an issue on our [GitHub repository](https://github.com/hao-ai-lab/FastVideo). - -You can also join our [Slack community](https://join.slack.com/t/fastvideo/shared_invite/zt-38u6p1jqe-yDI1QJOCEnbtkLoaI5bjZQ) for additional support. +- VSA is unsupported on MPS: unset `FASTVIDEO_ATTENTION_BACKEND` or set it + to `TORCH_SDPA`. +- Do not infer physical-16-GB support from allocator-cap experiments. +- Keep `mlx_dit/` beside the model's `transformer/`, `text_encoder/`, + `tokenizer/`, VAE, and scheduler files. diff --git a/docs/inference/support_matrix.md b/docs/inference/support_matrix.md index d0eec15644..bc3904264b 100644 --- a/docs/inference/support_matrix.md +++ b/docs/inference/support_matrix.md @@ -79,6 +79,17 @@ pipeline initialization and sampling. | 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_wan_prompt_to_video.py b/examples/inference/basic/mlx_wan_prompt_to_video.py new file mode 100644 index 0000000000..3cba39e86c --- /dev/null +++ b/examples/inference/basic/mlx_wan_prompt_to_video.py @@ -0,0 +1,523 @@ +"""Generate a small FastWan video with the experimental MLX FP16 DiT runtime. + +This is intentionally a hybrid proof of concept: + +- Hugging Face/torch encodes the prompt with UMT5. +- MLX runs the FastWan DiT denoising loop in FP16. +- Diffusers/torch decodes the final latents with AutoencoderKLWan. + +That makes it the first real prompt -> scheduler -> MLX DiT -> VAE -> MP4 path, +while keeping the initial MLX milestone focused on the denoiser runtime. +""" + +from __future__ import annotations + +import argparse +import gc +import json +import subprocess +import sys +import tempfile +import time +from pathlib import Path + +import numpy as np + +from fastvideo.mlx_runtime.memory import add_memory_limit_args, apply_memory_limits + + +DEFAULT_MODEL_ROOT = ( + Path.home() + / ".cache/huggingface/hub/models--FastVideo--FastWan2.1-T2V-1.3B-Diffusers/" + "snapshots/25e7ed7f41fd8ce2fdd108688c65e8caf0ce3aef" +) + + +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 _cleanup_torch() -> None: + import torch + + gc.collect() + if torch.backends.mps.is_available(): + torch.mps.empty_cache() + + +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.float16 if dtype_arg == "fp16" else torch.float32 + 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, + ) + prompt_embeds = prompt_embeds.cpu().contiguous() + del text_encoder, tokenizer, text_inputs, text_input_ids, mask + _cleanup_torch() + 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, + "--torch-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 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 + + if cache_path is not None and cache_path.exists(): + return torch.from_numpy(np.load(cache_path)).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: + cache_path.parent.mkdir(parents=True, exist_ok=True) + np.save(cache_path, prompt_embeds.cpu().numpy()) + 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: + 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.float16 if dtype_arg == "fp16" else torch.float32 + if backend == "taehv": + 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() + return + + if backend != "wan-vae": + raise ValueError(f"Unsupported decode backend: {backend}") + + 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, video + _cleanup_torch() + + +def main() -> None: + parser = argparse.ArgumentParser(description="Hybrid prompt-to-video FastWan generation using MLX for the DiT") + 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("--output-path", type=Path, default=Path("video_samples/mlx_fastwan_prompt_to_video.mp4")) + parser.add_argument("--height", type=int, default=256) + parser.add_argument("--width", type=int, default=448) + parser.add_argument("--num-frames", type=int, default=33) + 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("--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("--torch-device", default="auto", help="'auto', 'mps', or 'cpu' for text/VAE components.") + parser.add_argument("--torch-dtype", choices=("fp16", "fp32"), default="fp16") + parser.add_argument("--mlx-dtype", choices=("fp16", "bf16", "fp32"), default="fp16") + parser.add_argument( + "--mlx-quantization", + choices=("none", "int8", "int4", "mxfp8", "mxfp4", "nvfp4"), + default="none", + ) + 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="wan-vae") + 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) + 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() + + 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() + + if args.encode_prompt_only is not None: + 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, + ) + 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.fastwan import mlx_dit_from_diffusers_safetensors + from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step + + mx.random.seed(args.seed) + torch.manual_seed(args.seed) + + 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 + 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() + prompt_start = time.perf_counter() + prompt_embeds = get_prompt_embeds( + 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, + encode_mode=args.prompt_encode_mode, + cache_path=args.prompt_embeds_cache, + ) + 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) + config = dit.config + else: + 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() + + 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(args.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() + 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") + 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}") + + decode_start = time.perf_counter() + decode_latents_to_video( + model_root=args.model_root, + latents_np=latents_np, + output_path=args.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 + total_time = time.perf_counter() - total_start + + 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") + 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"Output written to: {args.output_path}") + + if args.metrics_json is not None: + metrics = { + "prompt": args.prompt, + "height": args.height, + "width": args.width, + "num_frames": 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()], + "mlx_dtype": args.mlx_dtype, + "mlx_quantization": args.mlx_quantization, + "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_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, + "total_s": total_time, + "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/train/configs/distribution_matching/wan/dmd2_t2v_mlx_int8.yaml b/examples/train/configs/distribution_matching/wan/dmd2_t2v_mlx_int8.yaml new file mode 100644 index 0000000000..d1affb88df --- /dev/null +++ b/examples/train/configs/distribution_matching/wan/dmd2_t2v_mlx_int8.yaml @@ -0,0 +1,106 @@ +# Quantization-aware DMD2 distillation targeting the Apple/MLX runtime: +# Wan 2.1 T2V 1.3B teacher -> 3-step INT8 student for Mac deployment. +# +# Identical to dmd2_t2v.yaml except: +# - the mlx_qat callback fake-quantizes the student's linear weights with +# MLX's exact affine INT8 quantizer (group size 64) every forward, so the +# student learns weights that survive `mx.quantize` at deploy time +# (numerics pinned bitwise by fastvideo/tests/mlx/test_mlx_affine_qat_parity.py); +# - 3 denoising steps [1000, 757, 522], matching the FastWan preset the MLX +# runtime ships (fastvideo/mlx_runtime + FAST_WAN_T2V_480P). +# +# Export after training: dcp_to_diffusers -> diffusers safetensors -> the MLX +# loader quantizes to the same grid the student trained on +# (`--save-mlx-checkpoint` then persists the pre-quantized weights). + +models: + student: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + teacher: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + disable_custom_init_weights: true + critic: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + disable_custom_init_weights: true + +method: + _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method + rollout_mode: simulate + generator_update_interval: 5 + real_score_guidance_scale: 4.5 + dmd_denoising_steps: [1000, 757, 522] + + # Critic optimizer (required — no fallback to training.optimizer) + fake_score_learning_rate: 8.0e-6 + fake_score_betas: [0.0, 0.999] + fake_score_lr_scheduler: constant + +training: + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + + data: + data_path: data/Wan-Syn_77x448x832_600k + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 20 + num_height: 448 + num_width: 832 + num_frames: 77 + + optimizer: + learning_rate: 2.0e-6 + betas: [0.0, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 4000 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/wan2.1_dmd2_3steps_mlx_int8 + training_state_checkpointing_steps: 20 + checkpoints_total_limit: 3 + + tracker: + project_name: distillation_wan + run_name: wan2.1_dmd2_3steps_mlx_int8 + + model: + enable_gradient_checkpointing_type: full + +callbacks: + mlx_qat: + group_size: 64 + bits: 8 + simulate_dtype: fp16 + grad_clip: + max_grad_norm: 1.0 + validation: + pipeline_target: fastvideo.pipelines.basic.wan.wan_dmd_pipeline.WanDMDPipeline + dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json + every_steps: 50 + sampling_steps: [3] + sampling_timesteps: [1000, 757, 522] + guidance_scale: 6.0 + ema: + _target_: fastvideo.train.callbacks.ema.EMACallback + decay: 0.98 + start_iter: 0 # delay EMA updates until this iteration + +pipeline: + flow_shift: 8 diff --git a/examples/train/configs/distribution_matching/wan/dmd2_t2v_mlx_int8_v2.yaml b/examples/train/configs/distribution_matching/wan/dmd2_t2v_mlx_int8_v2.yaml new file mode 100644 index 0000000000..4234c238d6 --- /dev/null +++ b/examples/train/configs/distribution_matching/wan/dmd2_t2v_mlx_int8_v2.yaml @@ -0,0 +1,104 @@ +# Run 2 of the Mac-targeted QAD distillation. Changes vs dmd2_t2v_mlx_int8.yaml: +# +# - Student and critic initialize from FastVideo/FastWan2.1-T2V-1.3B-Diffusers +# (already 3-step DMD-distilled): training becomes "adapt a good 3-step +# model to the INT8 grid" instead of learning distillation from scratch +# through a quantizer. The teacher stays the base 50-step Wan2.1. +# - gradient_accumulation_steps 4 (effective batch 16): run 1's global batch +# of 4 left the critic noisy and the student under-trained (weights moved +# only ~0.2% from init); motion coherence is the expected beneficiary. +# - Requires the world-size-portable EMA checkpoint state (ema.py) — run 1's +# EMA was unrecoverable from its checkpoint. Wall time ~4x run 1 per step +# count; same 4-GPU footprint. + +models: + student: + _target_: fastvideo.train.models.wan.WanModel + init_from: FastVideo/FastWan2.1-T2V-1.3B-Diffusers + trainable: true + teacher: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: false + disable_custom_init_weights: true + critic: + _target_: fastvideo.train.models.wan.WanModel + init_from: FastVideo/FastWan2.1-T2V-1.3B-Diffusers + trainable: true + disable_custom_init_weights: true + +method: + _target_: fastvideo.train.methods.distribution_matching.dmd2.DMD2Method + rollout_mode: simulate + generator_update_interval: 5 + real_score_guidance_scale: 4.5 + dmd_denoising_steps: [1000, 757, 522] + + # Critic optimizer (required — no fallback to training.optimizer) + fake_score_learning_rate: 8.0e-6 + fake_score_betas: [0.0, 0.999] + fake_score_lr_scheduler: constant + +training: + distributed: + num_gpus: 4 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 4 + + data: + data_path: data/Wan-Syn_77x448x832_600k + dataloader_num_workers: 4 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 20 + num_height: 448 + num_width: 832 + num_frames: 77 + + optimizer: + learning_rate: 2.0e-6 + betas: [0.0, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 4000 + gradient_accumulation_steps: 4 + + checkpoint: + output_dir: outputs/wan2.1_dmd2_3steps_mlx_int8_v2 + training_state_checkpointing_steps: 20 + checkpoints_total_limit: 3 + + tracker: + project_name: distillation_wan + run_name: wan2.1_dmd2_3steps_mlx_int8_v2 + + model: + enable_gradient_checkpointing_type: full + +callbacks: + mlx_qat: + group_size: 64 + bits: 8 + simulate_dtype: fp16 + grad_clip: + max_grad_norm: 1.0 + validation: + pipeline_target: fastvideo.pipelines.basic.wan.wan_dmd_pipeline.WanDMDPipeline + dataset_file: examples/training/finetune/Wan2.1-VSA/Wan-Syn-Data/validation_4.json + every_steps: 50 + sampling_steps: [3] + sampling_timesteps: [1000, 757, 522] + guidance_scale: 6.0 + ema: + _target_: fastvideo.train.callbacks.ema.EMACallback + decay: 0.98 + start_iter: 0 # delay EMA updates until this iteration + +pipeline: + flow_shift: 8 diff --git a/fastvideo/benchmarks/__init__.py b/fastvideo/benchmarks/__init__.py new file mode 100644 index 0000000000..9881313609 --- /dev/null +++ b/fastvideo/benchmarks/__init__.py @@ -0,0 +1 @@ +# SPDX-License-Identifier: Apache-2.0 diff --git a/fastvideo/benchmarks/mlx_fastwan_bench.py b/fastvideo/benchmarks/mlx_fastwan_bench.py new file mode 100644 index 0000000000..d429506c45 --- /dev/null +++ b/fastvideo/benchmarks/mlx_fastwan_bench.py @@ -0,0 +1,820 @@ +# 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 + +# 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 + + ssim_values = compute_video_ssim_torchvision(str(reference_video), str(candidate_video), use_ms_ssim=True) + 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) + + 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(): + """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/benchmarks/mlx_wan22_5b_bench.py b/fastvideo/benchmarks/mlx_wan22_5b_bench.py new file mode 100644 index 0000000000..9618f82e58 --- /dev/null +++ b/fastvideo/benchmarks/mlx_wan22_5b_bench.py @@ -0,0 +1,271 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Memory / latency prove-out for Wan2.2-TI2V-5B on Apple Silicon (Track D Rung 3). + +Loads the released 5B Diffusers transformer in fp16 and int8, runs a few 3-step +DMD denoise steps at a product-class latent shape, and records weight footprint, +peak unified memory, and steady per-step latency. Denoise-side only — random +text embeds (content-independent). Decode is noted separately: the 5B VAE is +z_dim=48; TAEHV taew2_1.pth targets Wan2.1 and is not a drop-in for 2.2. + + PYTHONPATH=$PWD python -m fastvideo.benchmarks.mlx_wan22_5b_bench \ + --model-root ~/models/fastwan22_5b --modes fp16,int8 +""" + +from __future__ import annotations + +import argparse +import json +import statistics +import time +from pathlib import Path + +import numpy as np + + +def _mx_dtype(name: str): + import mlx.core as mx + + return {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[name] + + +def _weight_bytes(model) -> int: + """Approximate resident weight bytes (plain + quantized matrices).""" + total = 0 + for value in model.weights.values(): + total += _array_nbytes(value) + for block in model.blocks: + for value in block.weights.values(): + total += _array_nbytes(value) + return total + + +def _array_nbytes(value) -> int: + # QuantizedMatrix from fastwan has .weight / .scales / .biases. + if hasattr(value, "weight") and hasattr(value, "scales"): + n = int(np.prod(value.weight.shape)) * _dtype_itemsize(value.weight) + n += int(np.prod(value.scales.shape)) * _dtype_itemsize(value.scales) + if getattr(value, "biases", None) is not None: + n += int(np.prod(value.biases.shape)) * _dtype_itemsize(value.biases) + return n + try: + return int(np.prod(value.shape)) * _dtype_itemsize(value) + except Exception: # noqa: BLE001 + return 0 + + +def _dtype_itemsize(arr) -> int: + name = str(getattr(arr, "dtype", "float16")) + if "64" in name: + return 8 + if "32" in name: + return 4 + if "16" in name or "bfloat" in name: + return 2 + if "8" in name or "uint" in name or "int8" in name: + return 1 + return 2 + + +def _build_per_token_timestep(mx, *, batch: int, frames: int, height: int, width: int, patch_size, video_t: float): + """2-D timestep: frame 0 at 0 (I2V image lock), remaining tokens at ``video_t``.""" + p_t, p_h, p_w = patch_size + tokens_per_frame = (height // p_h) * (width // p_w) + num_tokens = (frames // p_t) * tokens_per_frame + levels = [0.0] + [float(video_t)] * (frames // p_t - 1) + flat = [levels[i // tokens_per_frame] for i in range(num_tokens)] + return mx.array(np.array([flat] * batch, dtype=np.float32)) + + +def run_mode( + *, + model_root: Path, + mode: str, + height: int, + width: int, + num_frames: int, + dmd_steps: list[int], + flow_shift: float, + seed: int, +) -> dict: + import mlx.core as mx + + from examples.inference.basic.mlx_wan_prompt_to_video import make_rotary_embeddings + from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step, pred_noise_to_pred_video + from fastvideo.mlx_runtime.wan22 import mlx_wan22_dit_from_diffusers_safetensors + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + import torch + + config = json.loads((model_root / "transformer" / "config.json").read_text()) + # Wan2.2 VAE is 16× spatial (z_dim=48); Diffusers TI2V-5B uses latent H=height/16. + # For denoise-only we treat height/width as *pixel* size and derive latent size + # from the VAE spatial factor 16 (config does not store it; z_dim=48 ⇒ 2.2 VAE). + latent_h, latent_w = height // 16, width // 16 + # Latent temporal: TI2V keeps (num_frames-1)//4+1 style for 4× time compress. + latent_frames = (num_frames - 1) // 4 + 1 + if latent_frames < 1: + latent_frames = 1 + # Ensure divisible by patch temporal size. + p_t = int(config["patch_size"][0]) + if latent_frames % p_t != 0: + latent_frames += p_t - (latent_frames % p_t) + + quant = None if mode == "fp16" else mode + mx.clear_cache() + mx.reset_peak_memory() + load_start = time.perf_counter() + model = mlx_wan22_dit_from_diffusers_safetensors( + model_root / "transformer" / "diffusion_pytorch_model.safetensors", + model_root / "transformer" / "config.json", + dtype="fp16", + quantization=quant, + ) + # Force materialization of weights so peak memory reflects the load. + probe = next(iter(model.weights.values())) + if hasattr(probe, "weight"): + mx.eval(probe.weight) + else: + mx.eval(probe) + load_s = time.perf_counter() - load_start + load_peak = mx.get_peak_memory() / (1024**3) + weight_gib = _weight_bytes(model) / (1024**3) + + text_dim = int(config["text_dim"]) + text_len = 64 # short random embeds (latency is content-independent) + rng = np.random.default_rng(seed) + text = mx.array((rng.standard_normal((1, text_len, text_dim)) * 0.1).astype(np.float16)) + noise = mx.array( + rng.standard_normal((1, int(config["in_channels"]), latent_frames, latent_h, latent_w)).astype(np.float16)) + freqs_cis = make_rotary_embeddings(config, + latent_frames=latent_frames, + latent_height=latent_h, + latent_width=latent_w) + + scheduler = FlowMatchEulerDiscreteScheduler(shift=flow_shift) + scheduler.set_timesteps(1000, device="cpu") + schedule = MLXDMDSchedule.from_torch_scheduler(scheduler) + steps = torch.tensor(dmd_steps, dtype=torch.long) + warped = torch.cat((scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))) + timesteps = [float(t) for t in warped[1000 - steps]] + + # Warmup one step so compile/caches settle before timing. + ts0 = _build_per_token_timestep(mx, + batch=1, + frames=latent_frames, + height=latent_h, + width=latent_w, + patch_size=tuple(config["patch_size"]), + video_t=timesteps[0]) + _ = model(noise, text, ts0, freqs_cis) + mx.eval(_) + + mx.reset_peak_memory() + step_latencies: list[float] = [] + latents = noise + for i, ts_val in enumerate(timesteps): + ts = _build_per_token_timestep(mx, + batch=1, + frames=latent_frames, + height=latent_h, + width=latent_w, + patch_size=tuple(config["patch_size"]), + video_t=ts_val) + t0 = time.perf_counter() + pred = model(latents, text, ts, freqs_cis) + mx.eval(pred) + step_latencies.append(time.perf_counter() - t0) + if i < len(timesteps) - 1: + renoise = mx.random.normal(latents.shape).astype(mx.float32) + latents = dmd_step( + latents=latents.astype(mx.float32), + noise_input_latent=latents.astype(mx.float32), + pred_noise=pred.astype(mx.float32), + schedule=schedule, + timestep=ts_val, + next_timestep=timesteps[i + 1], + noise=renoise, + ).astype(mx.float16) + else: + latents = pred_noise_to_pred_video(pred.astype(mx.float32), latents.astype(mx.float32), + schedule.sigma_for(ts_val)).astype(mx.float16) + mx.eval(latents) + + denoise_peak = mx.get_peak_memory() / (1024**3) + steady = statistics.median(step_latencies) + return { + "mode": mode, + "quantization": quant or "none", + "pixel_hw": f"{height}x{width}", + "latent_shape": f"1x{config['in_channels']}x{latent_frames}x{latent_h}x{latent_w}", + "dmd_steps": dmd_steps, + "flow_shift": flow_shift, + "load_s": round(load_s, 2), + "weight_gib": round(weight_gib, 3), + "load_peak_gib": round(load_peak, 3), + "denoise_peak_gib": round(denoise_peak, 3), + "step_latencies_s": [round(x, 3) for x in step_latencies], + "steady_step_s": round(steady, 3), + "total_denoise_s": round(sum(step_latencies), 2), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description="Wan2.2-5B MLX memory/latency benchmark.") + parser.add_argument( + "--model-root", + type=Path, + default=Path.home() / "models" / "fastwan22_5b", + help="Root with transformer/config.json + diffusion_pytorch_model.safetensors", + ) + parser.add_argument("--modes", default="fp16,int8", help="Comma-separated: fp16,int8,...") + parser.add_argument("--height", type=int, default=480, help="Pixel height (latent = H/16).") + parser.add_argument("--width", type=int, default=832, help="Pixel width (latent = W/16).") + parser.add_argument("--num-frames", type=int, default=33, help="Pixel frames (latent T ≈ (F-1)/4+1).") + parser.add_argument("--dmd-denoising-steps", default="1000,757,522") + parser.add_argument("--flow-shift", type=float, default=5.0, help="Wan2.2-TI2V default shift.") + parser.add_argument("--seed", type=int, default=0) + parser.add_argument("--metrics-out", type=Path, default=None) + args = parser.parse_args() + + ckpt = args.model_root / "transformer" / "diffusion_pytorch_model.safetensors" + if not ckpt.exists(): + raise SystemExit(f"checkpoint not found: {ckpt}") + + dmd_steps = [int(s) for s in args.dmd_denoising_steps.split(",") if s.strip()] + modes = [m.strip() for m in args.modes.split(",") if m.strip()] + rows = [] + for mode in modes: + print(f"\n=== mode={mode} ===", flush=True) + row = run_mode( + model_root=args.model_root, + mode=mode, + height=args.height, + width=args.width, + num_frames=args.num_frames, + dmd_steps=dmd_steps, + flow_shift=args.flow_shift, + seed=args.seed, + ) + print(json.dumps(row, indent=2), flush=True) + rows.append(row) + + payload = { + "model": + "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers", + "model_root": + str(args.model_root), + "rows": + rows, + "decode_note": ("Wan2.2 VAE z_dim=48; TAEHV taew2_1.pth is Wan2.1-only. Full Wan2.2 VAE " + "decode on torch-MPS is the decode path until a 2.2-compatible TAE lands; " + "chunked/tiled decode may be needed for memory on 32 GB."), + } + print("\n=== summary ===") + print(json.dumps(payload, indent=2)) + if args.metrics_out is not None: + args.metrics_out.parent.mkdir(parents=True, exist_ok=True) + args.metrics_out.write_text(json.dumps(payload, indent=2)) + print(f"wrote {args.metrics_out}") + + +if __name__ == "__main__": + main() diff --git a/fastvideo/dataset/preprocessing_datasets.py b/fastvideo/dataset/preprocessing_datasets.py index 9d1ec44e77..902d7d2da7 100644 --- a/fastvideo/dataset/preprocessing_datasets.py +++ b/fastvideo/dataset/preprocessing_datasets.py @@ -11,6 +11,9 @@ import numpy as np import torch +# The Stateful base class below needs this submodule loaded; without the +# explicit import this module only works when something else imported it first. +import torch.distributed.checkpoint.stateful import torchvision from einops import rearrange from PIL import Image diff --git a/fastvideo/layers/quantization/mlx_affine_qat.py b/fastvideo/layers/quantization/mlx_affine_qat.py new file mode 100644 index 0000000000..0c76d0fd41 --- /dev/null +++ b/fastvideo/layers/quantization/mlx_affine_qat.py @@ -0,0 +1,119 @@ +# SPDX-License-Identifier: Apache-2.0 +"""PyTorch twin of MLX's affine quantizer, for Mac-targeted QAT. + +The Apple Silicon deployment path quantizes DiT linear weights with +``mx.quantize(w, group_size=64, bits=8, mode="affine")``. For +quantization-aware training to transfer to that runtime, the train-time +fake-quantization must reproduce MLX's quantizer exactly — a different +rounding rule or zero-point convention silently erases the QAT gains at +deploy time. + +This module transcribes the affine algorithm from MLX's CPU kernel +(``mlx/backend/cpu/quantized.cpp::quantize`` at v0.31.2), whose non-obvious +details are: + +- per-group min/max is computed in fp32 regardless of the input dtype, +- the scale is *negative* when ``|w_max| >= |w_min|`` (the quantizer anchors + at the endpoint with the larger magnitude), +- the anchor endpoint is re-expressed as an exact integer multiple of the + scale (``q0 = rint(edge / scale); scale = edge / q0; bias = edge``), and + MLX uses that adjusted scale to produce the integer codes, +- rounding is ``rint`` (round-half-to-even), matching ``torch.round``, +- codes are clamped to ``[0, 2^bits - 1]`` and scales/biases are cast to the + input dtype at the end. + +``fastvideo/tests/mlx/test_mlx_affine_qat_parity.py`` pins this against the +real ``mx.quantize``/``mx.dequantize`` — that test is the gate the roadmap +requires before any GPU is spent on a QAT run. +""" + +from __future__ import annotations + +import torch + +DEFAULT_GROUP_SIZE = 64 +DEFAULT_BITS = 8 +_EPS = 1e-7 + + +def _group(w: torch.Tensor, group_size: int) -> torch.Tensor: + if w.shape[-1] % group_size != 0: + raise ValueError(f"Last dim {w.shape[-1]} is not divisible by group_size {group_size}; " + "MLX affine quantization groups along the last axis.") + return w.reshape(*w.shape[:-1], w.shape[-1] // group_size, group_size) + + +def mlx_affine_quantize_reference( + w: torch.Tensor, + *, + group_size: int = DEFAULT_GROUP_SIZE, + bits: int = DEFAULT_BITS, +) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + """Quantize exactly like ``mx.quantize(..., mode="affine")``. + + Returns ``(codes, scales, biases)`` with unpacked integer ``codes`` (one + uint8-range value per element, not MLX's packed uint32 words); ``scales`` + and ``biases`` are cast to ``w.dtype`` like MLX casts to its input dtype. + """ + n_bins = float((1 << bits) - 1) + grouped = _group(w, group_size).float() + + w_min = grouped.min(dim=-1).values + w_max = grouped.max(dim=-1).values + mask = w_min.abs() > w_max.abs() + quant_scale = ((w_max - w_min) / n_bins).clamp_min(_EPS) + quant_scale = torch.where(mask, quant_scale, -quant_scale) + edge = torch.where(mask, w_min, w_max) + + q0 = torch.round(edge / quant_scale) + nonzero_q0 = q0 != 0 + scale = torch.where( + nonzero_q0, + edge / torch.where(nonzero_q0, q0, torch.ones_like(q0)), + quant_scale, + ) + bias = torch.where(nonzero_q0, edge, torch.zeros_like(edge)) + codes = torch.round((grouped - bias.unsqueeze(-1)) / scale.unsqueeze(-1)) + codes = codes.clamp(min=0.0, max=n_bins) + return codes.to(torch.int32), scale.to(w.dtype), bias.to(w.dtype) + + +def mlx_affine_dequantize_reference( + codes: torch.Tensor, + scales: torch.Tensor, + biases: torch.Tensor, + *, + out_shape: torch.Size | None = None, +) -> torch.Tensor: + """Dequantize exactly like MLX's kernels: ``code * scale + bias`` in the + scales' dtype (elementwise, no fp32 upcast of the fused expression).""" + dtype = scales.dtype + deq = codes.to(dtype) * scales.unsqueeze(-1) + biases.unsqueeze(-1) + if out_shape is not None: + deq = deq.reshape(out_shape) + return deq + + +def fake_quantize_mlx_affine( + w: torch.Tensor, + *, + group_size: int = DEFAULT_GROUP_SIZE, + bits: int = DEFAULT_BITS, + simulate_dtype: torch.dtype = torch.float16, +) -> torch.Tensor: + """Straight-through-estimator fake quantization for QAT forwards. + + Simulates the full deployment pipeline: cast the (bf16/fp32 master) + weight to ``simulate_dtype`` — the MLX loader casts checkpoints to fp16 + before quantizing — then quantize/dequantize with MLX's affine rules. + + Returns the dequantized weight in fp32, which represents every fp16 + deploy-time value exactly; casting back to a bf16 master dtype would + re-round and break the bit-for-bit correspondence with ``mx.dequantize``. + Gradients pass through to ``w`` unchanged (STE). + """ + w_sim = w.detach().to(simulate_dtype) + codes, scales, biases = mlx_affine_quantize_reference(w_sim, group_size=group_size, bits=bits) + deq = mlx_affine_dequantize_reference(codes, scales, biases, out_shape=w.shape).float() + w32 = w.float() + return w32 + (deq - w32).detach() diff --git a/fastvideo/mlx_runtime/__init__.py b/fastvideo/mlx_runtime/__init__.py new file mode 100644 index 0000000000..2775dd550c --- /dev/null +++ b/fastvideo/mlx_runtime/__init__.py @@ -0,0 +1,53 @@ +# 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, +) + +__all__ = [ + "AppliedMemoryLimits", + "FastWanShape", + "MLXQuantizationSpec", + "MLXWanDiT", + "MLXWanTransformerBlock", + "UnsupportedMLXQuantizationError", + "add_memory_limit_args", + "apply_memory_limits", + "ensure_quantization_supported", + "fastwan_shape", + "fastwan_shape_from_config", + "gib_to_bytes", + "load_mlx_dit_checkpoint", + "mlx_dit_from_diffusers_safetensors", + "mlx_block_weights_from_diffusers_safetensors", + "mlx_block_weights_from_torch", + "quantization_support_error", + "save_mlx_dit_checkpoint", +] diff --git a/fastvideo/mlx_runtime/causal.py b/fastvideo/mlx_runtime/causal.py new file mode 100644 index 0000000000..07bd483340 --- /dev/null +++ b/fastvideo/mlx_runtime/causal.py @@ -0,0 +1,156 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Causal (streaming) self-attention for the MLX FastWan runtime — Track C. + +Ports the KV-cached inference path of ``CausalWanSelfAttention`` from +``fastvideo/models/dits/causal_wanvideo.py`` to MLX. The porting insight this +module exists to exploit (see ``docs/design/mac_streaming_causal_guide.md``): + + mask-free cached decoding of one frame-block at a time is identical to a + single block-causal *masked* pass over the whole sequence. + +So the runtime needs no attention mask and no flex-attention — each chunk's +queries attend densely over the cached ``[0:local_end]`` window via +``mx.fast.scaled_dot_product_attention``. The rolling eviction with sink tokens +(which bounds memory for long/streaming rollouts) is reproduced index-for-index +from the torch reference. +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import TYPE_CHECKING + +from fastvideo.mlx_runtime.fastwan import apply_rotary_emb + +if TYPE_CHECKING: + import mlx.core as mx + +# When ``local_attn_size == -1`` the torch model keeps a fixed 21-latent-frame +# window for compatibility (``causal_wanvideo.py``); mirror that constant. +GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES = 21 + + +@dataclass +class MLXCausalKVCache: + """Preallocated rolling K/V cache for one attention layer. + + Mirrors the torch ``kv_cache`` dict (``k``, ``v``, ``global_end_index``, + ``local_end_index``) plus the ``sink_tokens`` count. ``k``/``v`` are + ``[batch, max_tokens, num_heads, head_dim]``; the *roped* keys are stored + (rotary is applied before the write, at global positions). + """ + + k: mx.array + v: mx.array + global_end_index: int + local_end_index: int + sink_tokens: int + + @classmethod + def allocate( + cls, + *, + batch: int, + max_tokens: int, + num_heads: int, + head_dim: int, + sink_tokens: int = 0, + dtype=None, + ) -> MLXCausalKVCache: + import mlx.core as mx + + dtype = dtype if dtype is not None else mx.float16 + shape = (batch, max_tokens, num_heads, head_dim) + return cls( + k=mx.zeros(shape, dtype=dtype), + v=mx.zeros(shape, dtype=dtype), + global_end_index=0, + local_end_index=0, + sink_tokens=sink_tokens, + ) + + +def max_attention_size(local_attn_size: int, frame_seqlen: int) -> int: + """Attention-window size in tokens, matching the torch reference.""" + if local_attn_size == -1: + return GLOBAL_ATTN_COMPAT_MAX_LATENT_FRAMES * frame_seqlen + return local_attn_size * frame_seqlen + + +def causal_self_attention_step( + q: mx.array, + k: mx.array, + v: mx.array, + cos: mx.array, + sin: mx.array, + cache: MLXCausalKVCache, + *, + current_start: int, + local_attn_size: int, + frame_seqlen: int, + scale: float | None = None, +) -> mx.array: + """One cached, mask-free causal attention step for a frame-block. + + ``q``/``k``/``v`` are ``[batch, num_new_tokens, num_heads, head_dim]`` for + the current chunk; ``cos``/``sin`` are the rotary tables for this chunk's + *global* positions (i.e. the caller has already offset by ``current_start``). + Writes the new roped K/V into ``cache`` (rolling out the oldest tokens past + the first ``sink_tokens`` on overflow) and returns the chunk's attention + output ``[batch, num_new_tokens, num_heads, head_dim]``. + """ + import mlx.core as mx + + num_new = q.shape[1] + head_dim = q.shape[-1] + scale = scale if scale is not None else head_dim**-0.5 + + roped_query = apply_rotary_emb(q, cos, sin, is_neox_style=False).astype(v.dtype) + roped_key = apply_rotary_emb(k, cos, sin, is_neox_style=False).astype(v.dtype) + + current_end = current_start + num_new + sink_tokens = cache.sink_tokens + window = max_attention_size(local_attn_size, frame_seqlen) + kv_cache_size = cache.k.shape[1] + global_end = cache.global_end_index + local_end_prev = cache.local_end_index + + overflow = (local_attn_size != -1 and current_end > global_end and num_new + local_end_prev > kv_cache_size) + if overflow: + # Discard the oldest tokens after the sinks by shifting content left. + num_evicted = num_new + local_end_prev - kv_cache_size + num_rolled = local_end_prev - num_evicted - sink_tokens + # Chunk larger than the non-sink capacity would make num_rolled negative and + # the subsequent local_start:local_end write would clobber the sink region. + if num_rolled < 0: + raise ValueError( + f"Chunk size ({num_new}) exceeds available cache capacity " + f"({kv_cache_size - sink_tokens} after sinks); cannot evict " + f"without overwriting sink tokens.") + # Copy the source slice first (mx slices are new arrays, so no aliasing). + rolled_k = cache.k[:, sink_tokens + num_evicted:sink_tokens + num_evicted + num_rolled] + rolled_v = cache.v[:, sink_tokens + num_evicted:sink_tokens + num_evicted + num_rolled] + cache.k[:, sink_tokens:sink_tokens + num_rolled] = rolled_k + cache.v[:, sink_tokens:sink_tokens + num_rolled] = rolled_v + local_end = local_end_prev + current_end - global_end - num_evicted + else: + local_end = local_end_prev + current_end - global_end + local_start = local_end - num_new + + cache.k[:, local_start:local_end] = roped_key + cache.v[:, local_start:local_end] = v + + win_start = max(0, local_end - window) + key_window = cache.k[:, win_start:local_end] + value_window = cache.v[:, win_start:local_end] + + attn = mx.fast.scaled_dot_product_attention( + roped_query.transpose(0, 2, 1, 3), + key_window.transpose(0, 2, 1, 3), + value_window.transpose(0, 2, 1, 3), + scale=scale, + ).transpose(0, 2, 1, 3) + + cache.global_end_index = current_end + cache.local_end_index = local_end + return attn diff --git a/fastvideo/mlx_runtime/causal_dit.py b/fastvideo/mlx_runtime/causal_dit.py new file mode 100644 index 0000000000..bbae1afae1 --- /dev/null +++ b/fastvideo/mlx_runtime/causal_dit.py @@ -0,0 +1,362 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Causal (streaming) Wan DiT forward for the MLX runtime — Track C, rung 4. + +Ports ``CausalWanTransformer3DModel._forward_inference`` (the KV-cached, +chunk-at-a-time path from ``fastvideo/models/dits/causal_wanvideo.py``) to MLX. +Reuses the dense port's patch-embed / condition / output / block weights +(``fastvideo.mlx_runtime.fastwan``) unchanged — the causal checkpoint has the +same weight layout — and swaps the self-attention for the cached, mask-free +``causal_self_attention_step`` plus a per-generation cross-attention cache and +per-frame timestep modulation. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +from fastvideo.mlx_runtime.causal import MLXCausalKVCache, causal_self_attention_step +from fastvideo.mlx_runtime.fastwan import ( + 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 + + +def _modulate_per_frame(x, scale, shift, *, temb_seq_len: int, tokens_per_temb: int): + """``(x reshaped to per-frame) * (1 + scale) + shift`` then flattened back. + + Mirrors the torch block's ``unflatten(1,(S,tpt)) * (1+scale) + shift`` so the + modulation is applied per timestep-frame (``scale``/``shift`` are + ``[B, S, 1, dim]``). + """ + batch, seq, dim = x.shape + x = x.reshape(batch, temb_seq_len, tokens_per_temb, dim) + x = x * (1.0 + scale) + shift + return x.reshape(batch, seq, dim) + + +def _gate_per_frame(x, gate, *, temb_seq_len: int, tokens_per_temb: int): + batch, seq, dim = x.shape + x = x.reshape(batch, temb_seq_len, tokens_per_temb, dim) + x = x * gate + return x.reshape(batch, seq, dim) + + +class MLXCausalWanTransformerBlock: + """Causal Wan block: cached self-attention + cross-attention cache.""" + + 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 + + def _cross_attention(self, x, encoder_hidden_states, crossattn_cache) -> mx.array: + 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 crossattn_cache is not None and crossattn_cache.get("is_init"): + key = crossattn_cache["k"] + value = crossattn_cache["v"] + else: + key = linear(encoder_hidden_states, self.weights["attn2.to_k.weight"], self.weights.get("attn2.to_k.bias")) + key = rms_norm(key, self.weights["attn2.norm_k.weight"], + eps=self.eps).reshape(batch, -1, self.num_heads, self.head_dim) + value = linear(encoder_hidden_states, self.weights["attn2.to_v.weight"], + self.weights.get("attn2.to_v.bias")).reshape(batch, -1, self.num_heads, self.head_dim) + if crossattn_cache is not None: + crossattn_cache["k"] = key + crossattn_cache["v"] = value + crossattn_cache["is_init"] = True + + attended = mx.fast.scaled_dot_product_attention( + q.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) + attended = attended.reshape(batch, -1, self.dim) + return linear(attended, self.weights["attn2.to_out.weight"], self.weights.get("attn2.to_out.bias")) + + def __call__( + self, + hidden_states, + encoder_hidden_states, + timestep_proj, + cos, + sin, + *, + kv_cache: MLXCausalKVCache, + crossattn_cache, + current_start: int, + local_attn_size: int, + frame_seqlen: int, + ) -> mx.array: + import mlx.core as mx + + orig_dtype = hidden_states.dtype + batch, seq_length, _ = hidden_states.shape + temb_seq_len = timestep_proj.shape[1] + tokens_per_temb = seq_length // temb_seq_len + + # e = scale_shift_table[1,6,dim] + timestep_proj[B,S,6,dim] -> [B,S,6,dim] + 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.reshape(batch, temb_seq_len, 1, self.dim) for part in mx.split(e, 6, axis=2) + ] + + # 1. Self-attention with per-frame modulation. + norm_hidden = layer_norm(hidden_states.astype(mx.float32), eps=self.eps) + norm_hidden = _modulate_per_frame(norm_hidden, + scale_msa, + shift_msa, + temb_seq_len=temb_seq_len, + tokens_per_temb=tokens_per_temb) + norm_hidden = norm_hidden.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) + + attn = causal_self_attention_step(query, + key, + value, + cos, + sin, + kv_cache, + current_start=current_start, + local_attn_size=local_attn_size, + frame_seqlen=frame_seqlen) + attn = attn.reshape(batch, -1, self.dim) + attn = linear(attn, self.weights["to_out.weight"], self.weights.get("to_out.bias")) + + # Residual (per-frame gate) + norm; self-attn residual-norm has null shift/scale. + residual = hidden_states + _gate_per_frame( + attn, gate_msa, temb_seq_len=temb_seq_len, tokens_per_temb=tokens_per_temb) + norm_hidden = layer_norm(residual.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) + hidden_states = residual.astype(orig_dtype) + + # 2. Cross-attention (cached), then residual-norm with per-frame shift/scale. + cross = self._cross_attention(norm_hidden, encoder_hidden_states, crossattn_cache) + residual = hidden_states + cross + norm_hidden = layer_norm(residual.astype(mx.float32), eps=self.eps) + norm_hidden = _modulate_per_frame(norm_hidden, + c_scale_msa, + c_shift_msa, + temb_seq_len=temb_seq_len, + tokens_per_temb=tokens_per_temb) + norm_hidden = norm_hidden.astype(orig_dtype) + hidden_states = residual.astype(orig_dtype) + + # 3. Feed-forward with per-frame 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 + _gate_per_frame( + ff, c_gate_msa, temb_seq_len=temb_seq_len, tokens_per_temb=tokens_per_temb) + return hidden_states.astype(orig_dtype) + + +class MLXCausalWanDiT: + """Causal Wan DiT with KV-cached, chunk-at-a-time inference (Track C).""" + + def __init__( + self, + weights: dict[str, mx.array], + blocks: list[MLXCausalWanTransformerBlock], + config: dict, + *, + local_attn_size: int = -1, + sink_size: int = 0, + num_frames_per_block: int = 1, + ) -> None: + 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.text_len = int(config.get("text_len", 512)) + self.local_attn_size = local_attn_size + self.sink_size = sink_size + self.num_frames_per_block = num_frames_per_block + + def allocate_caches(self, *, batch: int, frame_seqlen: int, dtype=None) -> tuple: + """One KV cache + cross-attn cache per block, sized to the window.""" + from fastvideo.mlx_runtime.causal import max_attention_size + + window = max_attention_size(self.local_attn_size, frame_seqlen) + kv_caches = [ + MLXCausalKVCache.allocate(batch=batch, + max_tokens=window, + num_heads=self.num_heads, + head_dim=self.head_dim, + sink_tokens=self.sink_size * frame_seqlen, + dtype=dtype) for _ in self.blocks + ] + crossattn_caches: list[dict] = [{"is_init": False} for _ in self.blocks] + return kv_caches, crossattn_caches + + 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-frame timestep conditioning. ``timestep`` is ``[B, frames]``.""" + import mlx.core as mx + + batch, frames = 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, frames, 6, self.hidden_size) + + # Pad the text sequence to text_len with zeros, exactly like the torch + # causal model does before the text embedder (causal_wanvideo.py). + if encoder_hidden_states.shape[1] > self.text_len: + # Match the reference model's fixed text context rather than + # letting a long prompt overflow the cached cross-attention shape. + encoder_hidden_states = encoder_hidden_states[:, :self.text_len] + pad = self.text_len - encoder_hidden_states.shape[1] + if pad > 0: + encoder_hidden_states = mx.concatenate([ + encoder_hidden_states, + mx.zeros((encoder_hidden_states.shape[0], pad, encoder_hidden_states.shape[2]), + dtype=encoder_hidden_states.dtype) + ], + axis=1) + + 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, frames, 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-frame output modulation: scale_shift_table[1,2,dim] + temb[B,F,1,dim]. + e = self.weights["scale_shift_table"][None] + temb_out[:, :, None, :].astype(mx.float32) + shift, scale = [part.reshape(batch, frames, 1, self.hidden_size) for part in mx.split(e, 2, axis=2)] + norm = layer_norm(hidden_states.astype(mx.float32), eps=self.eps) + tokens_per_frame = norm.shape[1] // frames + norm = _modulate_per_frame(norm, scale, shift, temb_seq_len=frames, tokens_per_temb=tokens_per_frame) + norm = norm.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_chunk( + self, + hidden_states, + encoder_hidden_states, + timestep, + cos, + sin, + kv_caches, + crossattn_caches, + *, + current_start: int, + ) -> mx.array: + """Denoise one frame-block; ``cos``/``sin`` are its global-position rotary.""" + batch, _, frames, height, width = hidden_states.shape + frame_seqlen = (height // self.patch_size[1]) * (width // self.patch_size[2]) + + hidden = self._patch_embed(hidden_states) + temb_out, timestep_proj, ehs = self._condition(timestep, encoder_hidden_states) + + if len(kv_caches) != len(self.blocks) or len(crossattn_caches) != len(self.blocks): + raise ValueError("KV and cross-attention cache counts must match the transformer block count") + for block, kv_cache, crossattn_cache in zip(self.blocks, kv_caches, crossattn_caches, strict=True): + hidden = block(hidden, + ehs, + timestep_proj, + cos, + sin, + kv_cache=kv_cache, + crossattn_cache=crossattn_cache, + current_start=current_start, + local_attn_size=self.local_attn_size, + frame_seqlen=frame_seqlen) + + return self._output(hidden, temb_out, batch=batch, frames=frames, height=height, width=width) + + +def mlx_causal_dit_from_diffusers_safetensors( + checkpoint_path: str | Path, + config_path: str | Path, + *, + dtype: str = "fp16", + num_blocks: int | None = None, + quantization=None, + local_attn_size: int = -1, + sink_size: int = 0, + num_frames_per_block: int = 1, +) -> MLXCausalWanDiT: + """Load a causal Wan DiT from a Diffusers checkpoint into ``MLXCausalWanDiT``. + + Reuses the dense Diffusers loader (the causal checkpoint has the same weight + layout) and re-wraps its blocks as causal blocks — only the forward differs. + """ + 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 = [ + MLXCausalWanTransformerBlock(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 MLXCausalWanDiT(dense.weights, + blocks, + dense.config, + local_attn_size=local_attn_size, + sink_size=sink_size, + num_frames_per_block=num_frames_per_block) diff --git a/fastvideo/mlx_runtime/causal_sampler.py b/fastvideo/mlx_runtime/causal_sampler.py new file mode 100644 index 0000000000..569309149d --- /dev/null +++ b/fastvideo/mlx_runtime/causal_sampler.py @@ -0,0 +1,137 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Streaming block-autoregressive DMD sampler for the MLX causal Wan runtime. + +Track C, rung 5. Drives ``MLXCausalWanDiT`` the way the torch causal DMD stage +(``fastvideo/pipelines/stages/causal_denoising.py``) does: one frame-block at a +time, each block denoised over the few-step DMD schedule while the KV cache holds +the *clean* latents of every earlier block. After a block's few steps, a +context-update forward at ``timestep=context_noise`` rewrites that block's K/V +with its clean values before the next block starts — so subsequent blocks attend +to clean history. Yields each block's latents as it finalizes, which is what +makes the preview stream. +""" + +from __future__ import annotations + +from typing import TYPE_CHECKING +from collections.abc import Iterator + +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.causal_dit import MLXCausalWanDiT + + +def build_dmd_schedule( + dmd_denoising_steps: list[int], + *, + flow_shift: float = 8.0, + warp_denoising_step: bool = True, +) -> tuple[MLXDMDSchedule, list[float]]: + """Return ``(schedule, timesteps)`` for causal DMD. + + Mirrors the torch stage: the raw denoising steps (e.g. ``[1000, 750, 500, + 250]``) index into a warped 1000-step flow-match schedule when + ``warp_denoising_step`` is set. + """ + import torch + + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + + scheduler = FlowMatchEulerDiscreteScheduler(shift=flow_shift) + scheduler.set_timesteps(1000, device="cpu") + steps = torch.tensor(dmd_denoising_steps, dtype=torch.long) + if warp_denoising_step: + warped = torch.cat((scheduler.timesteps.cpu(), torch.tensor([0], dtype=torch.float32))) + timesteps = warped[1000 - steps] + else: + timesteps = steps.to(torch.float32) + schedule = MLXDMDSchedule.from_torch_scheduler(scheduler) + return schedule, [float(t) for t in timesteps] + + +def stream_causal_latents( + model: MLXCausalWanDiT, + encoder_hidden_states: mx.array, + noise_latents: mx.array, + cos_full: mx.array, + sin_full: mx.array, + schedule: MLXDMDSchedule, + timesteps: list[float], + *, + frame_seqlen: int, + context_noise: float = 0.0, + seed: int = 0, + kv_caches: list | None = None, + crossattn_caches: list | None = None, +) -> Iterator[tuple[int, mx.array]]: + """Yield ``(block_index, clean_block_latents)`` as each block finalizes. + + ``noise_latents`` is ``[B, C, T, H, W]`` pure noise; ``cos_full``/``sin_full`` + are the rotary tables for the whole clip (frame-major), sliced per block. + """ + import mlx.core as mx + + mx.random.seed(seed) + dtype = noise_latents.dtype + batch, _, total_frames, _, _ = noise_latents.shape + nfb = model.num_frames_per_block + if total_frames % nfb != 0: + raise ValueError(f"total latent frames {total_frames} not divisible by num_frames_per_block {nfb}") + + if (kv_caches is None) ^ (crossattn_caches is None): + raise ValueError("kv_caches and crossattn_caches must both be provided or both omitted") + if kv_caches is None: + kv_caches, crossattn_caches = model.allocate_caches(batch=batch, frame_seqlen=frame_seqlen, dtype=dtype) + block_tokens = nfb * frame_seqlen + last = len(timesteps) - 1 + + for block_index in range(total_frames // nfb): + start = block_index * nfb + current_start = start * frame_seqlen + cos_blk = cos_full[current_start:current_start + block_tokens] + sin_blk = sin_full[current_start:current_start + block_tokens] + current = noise_latents[:, :, start:start + nfb] + + for i, timestep in enumerate(timesteps): + ts = mx.full((batch, nfb), timestep, dtype=mx.float32) + pred_noise = model.forward_chunk(current, + encoder_hidden_states, + ts, + cos_blk, + sin_blk, + kv_caches, + crossattn_caches, + current_start=current_start) + noise_input = current.astype(mx.float32) + pred = pred_noise.astype(mx.float32) + if i < last: + renoise = mx.random.normal(current.shape).astype(mx.float32) + current = dmd_step(latents=noise_input, + noise_input_latent=noise_input, + pred_noise=pred, + schedule=schedule, + timestep=timestep, + next_timestep=timesteps[i + 1], + noise=renoise).astype(dtype) + else: + current = pred_noise_to_pred_video(pred, noise_input, schedule.sigma_for(timestep)).astype(dtype) + + # Context update: rewrite this block's K/V from its clean latents so later + # blocks attend to clean history (output discarded). + ts_ctx = mx.full((batch, nfb), context_noise, dtype=mx.float32) + model.forward_chunk(current, + encoder_hidden_states, + ts_ctx, + cos_blk, + sin_blk, + kv_caches, + crossattn_caches, + current_start=current_start) + # The context-update output is discarded, so force its cache writes now. + # Otherwise MLX retains an O(T) lazy graph even though cache tensors are + # bounded, invalidating the long-rollout memory guarantee. + mx.eval(current, *[cache.k for cache in kv_caches], *[cache.v for cache in kv_caches]) + yield block_index, current diff --git a/fastvideo/mlx_runtime/checkpoint.py b/fastvideo/mlx_runtime/checkpoint.py new file mode 100644 index 0000000000..906e30f76f --- /dev/null +++ b/fastvideo/mlx_runtime/checkpoint.py @@ -0,0 +1,196 @@ +# 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 +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: + 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): + import mlx.core as mx + + return {"fp16": mx.float16, "bf16": mx.bfloat16, "fp32": mx.float32}[name] + + +def _flatten_weights(dit: MLXWanDiT) -> dict[str, Any]: + 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: + """Persist ``dit`` (plain or quantized) into ``checkpoint_dir``.""" + import mlx.core as mx + + checkpoint_dir = Path(checkpoint_dir) + checkpoint_dir.mkdir(parents=True, exist_ok=True) + + 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, + } + + weights_path = checkpoint_dir / WEIGHTS_FILENAME + mx.save_safetensors(str(weights_path), arrays) + (checkpoint_dir / MANIFEST_FILENAME).write_text(json.dumps(manifest, indent=2)) + 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: + """Rebuild an ``MLXWanDiT`` saved by :func:`save_mlx_dit_checkpoint`.""" + 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): + 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/fastwan.py b/fastvideo/mlx_runtime/fastwan.py new file mode 100644 index 0000000000..514d6e9495 --- /dev/null +++ b/fastvideo/mlx_runtime/fastwan.py @@ -0,0 +1,920 @@ +# SPDX-License-Identifier: Apache-2.0 +"""FastWan-oriented helpers for the experimental MLX runtime path.""" + +from __future__ import annotations + +import json +import math +import statistics +import time +from dataclasses import dataclass +from pathlib import Path + +import numpy as np + +from fastvideo.logger import init_logger + +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) + + +# This must be a Python float. A NumPy scalar multiplying a traced mx.array +# dispatches through NumPy, which evaluates the array and breaks mx.compile. +_GELU_TANH_COEF = math.sqrt(2.0 / math.pi) + + +def gelu_tanh(x): + import mlx.core as mx + + return 0.5 * x * (1.0 + mx.tanh(_GELU_TANH_COEF * (x + 0.044715 * mx.power(x, 3.0)))) + + +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) + + attn_output = 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_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 = 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 + + if self._compiled_forward is None: + self._compiled_forward = mx.compile(self._forward) + try: + return self._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 + 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) + + +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.""" + import mlx.core as mx + from safetensors import safe_open + + prefix = f"blocks.{block_index}." + 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", + } + + 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: + for source_name, target_name in key_map.items(): + array = _load_mx_array_from_safetensor(handle, prefix + source_name, 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, +) -> 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: + for name in top_level_names: + 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) + + +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 = { + "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", + } + + state = {} + with safe_open(str(checkpoint_path), framework="pt", device="cpu") as handle: + for source_name, target_name in key_map.items(): + state[target_name] = handle.get_tensor(prefix + source_name).float() + return state diff --git a/fastvideo/mlx_runtime/hardware_tier.py b/fastvideo/mlx_runtime/hardware_tier.py new file mode 100644 index 0000000000..4cb787a13b --- /dev/null +++ b/fastvideo/mlx_runtime/hardware_tier.py @@ -0,0 +1,286 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Hardware-adaptive model + quant tiering for Apple Silicon MLX. + +Picks a practical FastWan config from the Mac's unified-memory size so a 16 GB +machine defaults to a small INT8 model with a tight MLX allocator cap, while +32/64 GB machines can run higher fidelity (fp16, fuller decoders, larger caps). + +Pure recommendation helpers are unit-testable with injected memory sizes; the +detection path reads ``sysctl hw.memsize`` on macOS and falls back to MLX +``device_info`` / a safe default when Metal or sysctl is unavailable. +""" + +from __future__ import annotations + +import platform +import subprocess +from dataclasses import dataclass +from typing import Any, Literal + +GIB = 1024**3 + +# Thresholds are inclusive upper bounds for each tier band (easy to retune). +TIER_SMALL_MAX_GIB = 18.0 +TIER_MEDIUM_MAX_GIB = 40.0 + +# Recommended MLX allocator caps leave headroom for the OS, torch/MPS encode, +# and decode. Distinct from the benchmark *stress* presets that pin closer to +# the machine class (e.g. mac-16gb uses a 16 GiB stress cap). +TIER_SMALL_MLX_CAP_GIB = 12.0 +TIER_MEDIUM_MLX_CAP_GIB = 24.0 +TIER_LARGE_MLX_CAP_GIB = 48.0 + +MODEL_1_3B_REPO = "FastVideo/FastWan2.1-T2V-1.3B-Diffusers" + +# TODO(Track D): once the 5B MLX port is parity-green on Metal, set this to the +# Diffusers repo id (candidate: "FastVideo/FastWan2.2-TI2V-5B-Diffusers"). +# Until then ``recommend_tier(..., prefer_5b=True)`` falls back to 1.3B configs +# when this is ``None``. +FIVE_B_MODEL_REPO: str | None = None + +# Non-Mac / undetectable-memory fallback: choose the safest small tier. +DEFAULT_ASSUMED_MEMORY_GIB = 16.0 + +QuantName = Literal["int8", "none"] +DecoderName = Literal["taehv", "wan-vae"] + + +@dataclass(frozen=True) +class HardwareTier: + """Immutable recommended config for one unified-memory band.""" + + name: str + """Human-readable tier id: ``small`` / ``medium`` / ``large``.""" + + max_memory_gib: float | None + """Inclusive upper bound of this band; ``None`` means unbounded (large).""" + + model_repo: str + """Hugging Face Diffusers repo id for the recommended checkpoint.""" + + quantization: QuantName + """``"int8"`` for weight-only affine quant; ``"none"`` for fp16 weights.""" + + decoder: DecoderName + """Default decode path for this tier.""" + + mlx_memory_limit_gib: float + """Suggested MLX allocator cap (GiB).""" + + benchmark_preset: str + """Matching key in ``BENCHMARK_PRESETS`` (``mac-16gb`` / ``mac-32gb`` / ``mac-64gb``).""" + + modes: str + """Comma-separated benchmark modes applied by ``--auto-tier``.""" + + decoders: str + """Comma-separated benchmark decoders applied by ``--auto-tier``.""" + + uses_5b: bool = False + """True when the recommendation selected a 5B checkpoint (Track D).""" + + +def _tier_small() -> HardwareTier: + return HardwareTier( + name="small", + max_memory_gib=TIER_SMALL_MAX_GIB, + model_repo=MODEL_1_3B_REPO, + quantization="int8", + decoder="taehv", + mlx_memory_limit_gib=TIER_SMALL_MLX_CAP_GIB, + benchmark_preset="mac-16gb", + modes="int8", + decoders="taehv", + uses_5b=False, + ) + + +def _tier_medium(*, prefer_5b: bool, five_b_model_repo: str | None) -> HardwareTier: + use_5b = bool(prefer_5b and five_b_model_repo) + if use_5b: + # 5B INT8 + TAEHV once Track D lands. + return HardwareTier( + name="medium", + max_memory_gib=TIER_MEDIUM_MAX_GIB, + model_repo=five_b_model_repo, # type: ignore[arg-type] + quantization="int8", + decoder="taehv", + mlx_memory_limit_gib=TIER_MEDIUM_MLX_CAP_GIB, + benchmark_preset="mac-32gb", + modes="int8", + decoders="taehv", + uses_5b=True, + ) + # Fallback until 5B is available: 1.3B fp16 (higher quality than small's int8). + return HardwareTier( + name="medium", + max_memory_gib=TIER_MEDIUM_MAX_GIB, + model_repo=MODEL_1_3B_REPO, + quantization="none", + decoder="taehv", + mlx_memory_limit_gib=TIER_MEDIUM_MLX_CAP_GIB, + benchmark_preset="mac-32gb", + modes="fp16", + decoders="taehv", + uses_5b=False, + ) + + +def _tier_large(*, prefer_5b: bool, five_b_model_repo: str | None) -> HardwareTier: + use_5b = bool(prefer_5b and five_b_model_repo) + if use_5b: + return HardwareTier( + name="large", + max_memory_gib=None, + model_repo=five_b_model_repo, # type: ignore[arg-type] + quantization="none", + decoder="wan-vae", + mlx_memory_limit_gib=TIER_LARGE_MLX_CAP_GIB, + benchmark_preset="mac-64gb", + modes="fp16", + decoders="wan-vae", + uses_5b=True, + ) + return HardwareTier( + name="large", + max_memory_gib=None, + model_repo=MODEL_1_3B_REPO, + quantization="none", + decoder="wan-vae", + mlx_memory_limit_gib=TIER_LARGE_MLX_CAP_GIB, + benchmark_preset="mac-64gb", + modes="fp16", + decoders="wan-vae", + uses_5b=False, + ) + + +def _bytes_to_gib(num_bytes: int | float) -> float: + return float(num_bytes) / float(GIB) + + +def _sysctl_memsize_bytes() -> int | None: + """Return total physical/unified memory via macOS ``sysctl``, or None.""" + if platform.system() != "Darwin": + return None + try: + out = subprocess.check_output(["sysctl", "-n", "hw.memsize"], text=True, timeout=2.0) + value = int(out.strip()) + return value if value > 0 else None + except (OSError, subprocess.SubprocessError, ValueError): + return None + + +def _mlx_device_memory_bytes(mx_module: Any | None = None) -> int | None: + """Return ``memory_size`` from MLX device_info when Metal is available.""" + try: + if mx_module is None: + import mlx.core as mlx_core + mx_module = mlx_core + info = None + if hasattr(mx_module, "device_info"): + info = mx_module.device_info() + elif hasattr(mx_module, "metal") and hasattr(mx_module.metal, "device_info"): + # Older MLX: metal.device_info (deprecated in favour of mx.device_info). + if not mx_module.metal.is_available(): + return None + info = mx_module.metal.device_info() + if not isinstance(info, dict): + return None + memory_size = info.get("memory_size") + if memory_size is None: + return None + value = int(memory_size) + return value if value > 0 else None + except Exception: # noqa: BLE001 - optional path; any failure → caller falls back. + return None + + +def _proc_meminfo_bytes() -> int | None: + """Linux fallback: parse MemTotal from ``/proc/meminfo``.""" + try: + with open("/proc/meminfo", encoding="utf-8") as fh: + for line in fh: + if line.startswith("MemTotal:"): + # kB + parts = line.split() + kib = int(parts[1]) + return kib * 1024 if kib > 0 else None + except (OSError, ValueError, IndexError): + return None + return None + + +def detect_unified_memory_gib(*, mx_module: Any | None = None) -> float: + """Detect total unified/system memory in GiB. + + Preference order: + 1. macOS ``sysctl -n hw.memsize`` (true unified-memory size) + 2. MLX ``device_info()["memory_size"]`` when Metal is present + 3. Linux ``/proc/meminfo`` MemTotal + 4. :data:`DEFAULT_ASSUMED_MEMORY_GIB` (safe small-tier assumption) + """ + for probe in ( + _sysctl_memsize_bytes, + lambda: _mlx_device_memory_bytes(mx_module), + _proc_meminfo_bytes, + ): + num_bytes = probe() + if num_bytes is not None and num_bytes > 0: + return _bytes_to_gib(num_bytes) + return DEFAULT_ASSUMED_MEMORY_GIB + + +def recommend_tier( + memory_gib: float | None = None, + *, + prefer_5b: bool = True, + five_b_model_repo: str | None = None, + mx_module: Any | None = None, +) -> HardwareTier: + """Recommend a :class:`HardwareTier` for the given (or detected) memory. + + Args: + memory_gib: Injected unified memory in GiB. When ``None``, calls + :func:`detect_unified_memory_gib`. Pure/unit tests should pass this + explicitly so no Metal or sysctl is required. + prefer_5b: When True and a 5B repo id is known, medium/large tiers select + the 5B checkpoint. When the 5B id is unset (Track D not landed), + falls back to 1.3B configs. + five_b_model_repo: Explicit 5B Diffusers repo id for this call. When + ``None``, falls back to :data:`FIVE_B_MODEL_REPO` (also ``None`` + until Track D). Tests inject a fake id here to exercise the 5B path. + mx_module: Optional MLX module for detection (injection / tests). + """ + repo = five_b_model_repo if five_b_model_repo is not None else FIVE_B_MODEL_REPO + + if memory_gib is None: + memory_gib = detect_unified_memory_gib(mx_module=mx_module) + if memory_gib <= 0: + # Defensive: treat nonsense as the safe small tier. + memory_gib = DEFAULT_ASSUMED_MEMORY_GIB + + if memory_gib <= TIER_SMALL_MAX_GIB: + return _tier_small() + if memory_gib <= TIER_MEDIUM_MAX_GIB: + return _tier_medium(prefer_5b=prefer_5b, five_b_model_repo=repo) + return _tier_large(prefer_5b=prefer_5b, five_b_model_repo=repo) + + +def apply_tier_to_namespace(args: Any, tier: HardwareTier) -> HardwareTier: + """Mutate an argparse namespace with tier modes, decoders, and MLX caps. + + Leaves ``model_root`` alone (local path vs HF id); callers that resolve + checkpoints from a repo id should read ``tier.model_repo``. Sets + ``args.auto_tier_name`` / ``args.auto_tier_model_repo`` for metrics. + """ + args.modes = tier.modes + args.decoders = tier.decoders + args.mlx_memory_limit_gib = tier.mlx_memory_limit_gib + # Memory-tier runs disable the MLX cache so the allocator cap is meaningful. + args.mlx_disable_cache = True + args.auto_tier_name = tier.name + args.auto_tier_model_repo = tier.model_repo + args.auto_tier_quantization = tier.quantization + args.auto_tier_benchmark_preset = tier.benchmark_preset + return tier diff --git a/fastvideo/mlx_runtime/memory.py b/fastvideo/mlx_runtime/memory.py new file mode 100644 index 0000000000..958577ee3a --- /dev/null +++ b/fastvideo/mlx_runtime/memory.py @@ -0,0 +1,152 @@ +# 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 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]: + 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, value in self.errors.items(): + metrics[f"{name}_error"] = value + return metrics + + +def gib_to_bytes(value: float | None) -> int | None: + 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 _set_mps_env(name: str, value: float | None) -> float | None: + 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 + callers should invoke this before importing torch. If a high watermark is + requested without a low watermark, the low watermark is set to ``0.0`` to + avoid PyTorch's default low watermark exceeding the requested high cap. + """ + 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_module + + if memory_bytes is not None: + previous["mlx_memory_limit"] = int(mx_module.set_memory_limit(memory_bytes)) + applied["mlx_memory_limit"] = memory_bytes + if cache_bytes is not None: + previous["mlx_cache_limit"] = int(mx_module.set_cache_limit(cache_bytes)) + applied["mlx_cache_limit"] = cache_bytes + if wired_bytes is not None: + try: + previous["mlx_wired_limit"] = int(mx_module.set_wired_limit(wired_bytes)) + applied["mlx_wired_limit"] = wired_bytes + except Exception as exc: # noqa: BLE001 - macOS/system-limit dependent. + errors["mlx_wired_limit"] = 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 shared Apple Silicon memory-tier flags to an argparse parser.""" + 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/sampling.py b/fastvideo/mlx_runtime/sampling.py new file mode 100644 index 0000000000..ccf05c6124 --- /dev/null +++ b/fastvideo/mlx_runtime/sampling.py @@ -0,0 +1,129 @@ +# 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: + """Nearest-timestep ``sigma`` lookup, computed on the host.""" + 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": + """``x0`` prediction from a flow-match noise prediction (scalar ``sigma``). + + Equivalent to ``fastvideo/models/utils.py::pred_noise_to_pred_video`` for a + single scalar timestep: ``pred_video = noise_input - sigma * pred_noise``. + """ + 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": + """One on-device DMD update. + + Predicts the clean video from ``pred_noise`` at ``timestep`` and, when a + ``next_timestep`` is given, re-noises it to that level (the intermediate + DMD steps). The final step (``next_timestep is None``) returns the clean + prediction directly. + + Args: + latents: unused placeholder kept for call-site symmetry with the torch + loop; the update depends only on ``noise_input_latent`` and + ``pred_noise``. + noise_input_latent: the noisy latent fed to the DiT this step. + pred_noise: the DiT's velocity/noise prediction this step. + schedule: host-side flow-match schedule. + timestep: current DMD timestep. + next_timestep: the DMD timestep to re-noise toward, or ``None`` on the + last step. + noise: fresh Gaussian noise (same shape as ``pred_noise``) for the + re-noise; required when ``next_timestep`` is not ``None``. + + Returns: + The latents to carry into the next step (or the final clean latents). + """ + 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 0000000000..9ae376363c --- /dev/null +++ b/fastvideo/mlx_runtime/taehv_decode.py @@ -0,0 +1,118 @@ +# 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" + + +def _default_cache_dir() -> Path: + return Path.home() / ".cache" / "fastvideo" / "taehv" + + +def _sha256(path: Path) -> str: + 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: + 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: + """Return a verified ``taew2_1.pth``, downloading into the cache if needed. + + A caller-supplied ``checkpoint_path`` is treated as trusted user input and + is not hash-checked (it may legitimately be a newer or retrained decoder); + only the checkpoint this module downloads itself is pinned. + """ + 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}") + urllib.request.urlretrieve(TAEW2_1_CHECKPOINT_URL, checkpoint_path) # noqa: S310 - pinned public artifact, hash-verified below. + _verify_checkpoint(checkpoint_path) + return checkpoint_path + + +def _load_taehv_class(source_path: Path | None): + 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 0000000000..c3b52399f9 --- /dev/null +++ b/fastvideo/mlx_runtime/wan22.py @@ -0,0 +1,205 @@ +# 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 only after an approved VAE/image/mask parity +port. This module deliberately accepts *latent-only* I2V preparation: encode +the image, replace the first latent frame, and set that frame's timestep to 0. +It is not a public image-to-video pipeline or an image encoder. +""" + +from __future__ import annotations + +from pathlib import Path +from typing import TYPE_CHECKING + +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 + + +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) -> None: + 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)) + + 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 __call__(self, hidden_states, encoder_hidden_states, timestep, freqs_cis) -> mx.array: + cos, sin = freqs_cis + 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 mlx_wan22_dit_from_diffusers_safetensors( + checkpoint_path: str | Path, + config_path: str | Path, + *, + dtype: str = "fp16", + num_blocks: int | None = None, + quantization=None, +) -> 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) diff --git a/fastvideo/mlx_runtime/wan22_i2v.py b/fastvideo/mlx_runtime/wan22_i2v.py new file mode 100644 index 0000000000..2ecbd4b927 --- /dev/null +++ b/fastvideo/mlx_runtime/wan22_i2v.py @@ -0,0 +1,165 @@ +# SPDX-License-Identifier: Apache-2.0 +"""I2V (image-to-video) input construction for Wan2.2-TI2V-5B on MLX. + +TI2V-5B has no CLIP image embedder. Image conditioning is purely: + +1. VAE-encode the input image to a single latent frame (caller / torch-MPS). +2. **Replace** ``noise_latents[:, :, 0]`` with that image latent. +3. Build a **per-token** timestep ``[B, L]`` with frame-0 tokens at ``t=0`` + (clean) and the remaining tokens at the denoise level. + +Token order is **frame-major**, matching ``MLXWan22DiT._patch_embed`` and the +torch Wan patch embed: for each post-patch frame, tokens scan ``H_p × W_p``. +So the first ``tokens_per_frame`` entries of the ``[B, L]`` timestep are the +image-frame tokens. +""" + +from __future__ import annotations + +from collections.abc import Sequence +from typing import TYPE_CHECKING, Any + +import numpy as np + +if TYPE_CHECKING: + pass + + +def tokens_per_frame(height: int, width: int, patch_size: Sequence[int]) -> int: + """Number of patch tokens in one latent frame (frame-major layout).""" + _pt, ph, pw = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) + if height % ph != 0 or width % pw != 0: + raise ValueError(f"latent HxW ({height}x{width}) not divisible by patch {(ph, pw)}") + return (height // ph) * (width // pw) + + +def num_patch_tokens(frames: int, height: int, width: int, patch_size: Sequence[int]) -> int: + """Total sequence length ``L`` after patch embedding.""" + pt, _, _ = int(patch_size[0]), int(patch_size[1]), int(patch_size[2]) + if frames % pt != 0: + raise ValueError(f"latent frames {frames} not divisible by patch_t {pt}") + return (frames // pt) * tokens_per_frame(height, width, patch_size) + + +def replace_first_latent_frame(noise_latents: Any, image_latent_frame: Any) -> Any: + """Return latents with frame 0 replaced by ``image_latent_frame``. + + Args: + noise_latents: ``[B, C, T, H, W]`` (torch or numpy or mx.array). + image_latent_frame: ``[B, C, H, W]`` or ``[B, C, 1, H, W]``. + """ + # Dispatch without importing mlx/torch at module import time. + if not isinstance(noise_latents, np.ndarray): + try: + import torch + except ImportError: + torch = None # type: ignore[assignment] + if torch is not None and isinstance(noise_latents, torch.Tensor): + out = noise_latents.clone() + img = image_latent_frame + if img.dim() == 5: + img = img[:, :, 0] + if tuple(img.shape) != tuple(out[:, :, 0].shape): + raise ValueError(f"image frame shape {tuple(img.shape)} != " + f"latent frame 0 shape {tuple(out[:, :, 0].shape)}") + out[:, :, 0] = img.to(dtype=out.dtype, device=out.device) + return out + + import mlx.core as mx + + if isinstance(noise_latents, mx.array): + img = image_latent_frame + if not isinstance(img, mx.array): + # Torch tensors on MPS cannot be handed to NumPy directly. The + # I2V preparation path is intentionally CPU/host-mediated until + # the VAE parity gate is complete, so make that transfer + # explicit and fail nowhere later inside mx.array. + if hasattr(img, "detach") and hasattr(img, "cpu"): + img = img.detach().cpu().numpy() + img = mx.array(np.asarray(img)) + if img.ndim == 4: + img = img[:, :, None, :, :] + if int(img.shape[2]) != 1: + raise ValueError(f"image_latent_frame must be a single frame, got {img.shape}") + if tuple(img[:, :, 0].shape) != tuple(noise_latents[:, :, 0].shape): + raise ValueError( + f"image frame shape {tuple(img[:, :, 0].shape)} != " + f"latent frame 0 shape {tuple(noise_latents[:, :, 0].shape)}") + frame0_shape = tuple(noise_latents[:, :, :1, :, :].shape) + if tuple(img.shape) != frame0_shape: + raise ValueError(f"image frame shape {tuple(img.shape)} != latent frame 0 shape {frame0_shape}") + rest = noise_latents[:, :, 1:, :, :] + return mx.concatenate([img.astype(noise_latents.dtype), rest], axis=2) + + # NumPy path. + out = np.array(noise_latents, copy=True) + img = np.asarray(image_latent_frame) + if img.ndim == 5: + img = img[:, :, 0] + if img.shape != out[:, :, 0].shape: + raise ValueError(f"image frame shape {img.shape} != latent frame 0 shape {out[:, :, 0].shape}") + out[:, :, 0] = img + return out + + +def build_i2v_per_token_timestep( + *, + batch: int, + frames: int, + height: int, + width: int, + patch_size: Sequence[int], + video_timestep: float, + image_timestep: float = 0.0, + as_numpy: bool = True, +) -> Any: + """Build frame-major per-token timesteps for TI2V-style I2V. + + Frame 0's tokens get ``image_timestep`` (default 0 = clean image lock); + remaining frames get ``video_timestep``. + + Returns: + Array of shape ``[batch, L]`` as NumPy float32 by default, or an + ``mx.array`` when ``as_numpy=False``. + """ + tpf = tokens_per_frame(height, width, patch_size) + n_frames_patched = frames // int(patch_size[0]) + levels = [float(image_timestep)] + [float(video_timestep)] * (n_frames_patched - 1) + flat = [levels[i // tpf] for i in range(n_frames_patched * tpf)] + arr = np.array([flat] * batch, dtype=np.float32) + if as_numpy: + return arr + import mlx.core as mx + + return mx.array(arr) + + +def build_i2v_inputs( + noise_latents: Any, + image_latent_frame: Any, + *, + video_timestep: float, + image_timestep: float = 0.0, + patch_size: Sequence[int] = (1, 2, 2), +) -> tuple[Any, np.ndarray]: + """Construct ``(latents, per_token_timestep)`` for one I2V DiT forward. + + Works with NumPy / torch / mx latents; timestep is always returned as + NumPy float32 ``[B, L]`` (caller casts for the backend). + """ + shape = tuple(noise_latents.shape) if hasattr(noise_latents, "shape") else np.asarray(noise_latents).shape + if len(shape) != 5: + raise ValueError(f"noise_latents must be [B,C,T,H,W], got {shape}") + batch, _c, frames, height, width = shape + latents = replace_first_latent_frame(noise_latents, image_latent_frame) + timestep = build_i2v_per_token_timestep( + batch=batch, + frames=frames, + height=height, + width=width, + patch_size=patch_size, + video_timestep=video_timestep, + image_timestep=image_timestep, + as_numpy=True, + ) + return latents, timestep diff --git a/fastvideo/mlx_runtime/wan22_sample.py b/fastvideo/mlx_runtime/wan22_sample.py new file mode 100644 index 0000000000..f62ac3d843 --- /dev/null +++ b/fastvideo/mlx_runtime/wan22_sample.py @@ -0,0 +1,91 @@ +# 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]]: + """Return ``(schedule, continuous_timesteps)`` for Wan2.2 FastWan DMD.""" + 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) + scheduler.set_timesteps(1000, device="cpu") + 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: + """Run warped 3-step (default) DMD on ``noise_latents`` → clean 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 0000000000..782e574151 --- /dev/null +++ b/fastvideo/mlx_runtime/wan_vae.py @@ -0,0 +1,401 @@ +# 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 os +import tempfile +import urllib.request +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import numpy as np + +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/main/taew2_2.pth" +# Validated 2026-07-02 / 2026-07-09 against upstream madebyollin/taehv. +TAEW2_1_SHA256 = "d26151e76cdc2c9424bef988de874b33d9a53f30ef3060cd556c429c469c797e" +# Leave empty to skip pin for taew2_2 until a long-lived pin is recorded (still +# downloads; user-supplied paths are never hash-checked). +TAEW2_2_SHA256 = "" + +DecodeBackend = Literal["taehv", "taehv-torch", "wan-vae"] + + +def _cache_dir() -> Path: + return Path.home() / ".cache" / "fastvideo" / "taehv" + + +def _sha256(path: Path) -> str: + 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 _download_atomic(url: str, destination: Path) -> None: + """Download a checkpoint to a sibling temp file and atomically publish it.""" + destination.parent.mkdir(parents=True, exist_ok=True) + fd, temporary_name = tempfile.mkstemp(prefix=f".{destination.name}.", suffix=".partial", dir=destination.parent) + try: + with os.fdopen(fd, "wb") as output, urllib.request.urlopen(url) as response: # noqa: S310 + while chunk := response.read(1 << 20): + output.write(chunk) + os.replace(temporary_name, destination) + except BaseException: + Path(temporary_name).unlink(missing_ok=True) + raise + + +def ensure_taehv_checkpoint( + *, z_dim: int, checkpoint_path: Path | None = None, allow_download: bool = False +) -> Path: + """Return a TAEHV weight file for the given 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(): + if not allow_download: + raise FileNotFoundError( + f"TAEHV checkpoint is not staged at {path}. This future-only runtime never downloads artifacts " + "implicitly; pass a reviewed checkpoint_path or explicitly enable a pinned download.") + if not expect: + raise RuntimeError(f"No SHA256 pin is recorded for {name}; do not download it until PORT_STATUS is updated") + print(f"Downloading {url} -> {path}") + _download_atomic(url, path) + if expect: + actual = _sha256(path) + if actual != expect: + raise RuntimeError(f"TAEHV {path} sha256 mismatch (expected {expect}, got {actual})") + 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: + 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: + """Apply ``z / (1/std) + mean`` ⇔ ``z * std + mean`` (diffusers Wan path). + + Matches ``mlx_wan_prompt_to_video``:: + + latents = latents / latents_std + latents_mean + # where latents_std was stored as 1/std in the tensor view + """ + 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: + 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]: + 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: + import mlx.core as mx + + self.checkpoint_path = Path(checkpoint_path) + if z_dim not in (16, 48): + raise ValueError(f"unsupported TAEHV z_dim={z_dim}; expected 16 or 48") + self.patch_size = 2 if z_dim == 48 else 1 + self.latent_channels = z_dim + 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 ``[N, T, C, H, W]`` diffusion latents → ``[N, T_out, 3, H_p, W_p]`` in ~[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 decoder Sequential with MemBlock temporal memory (parallel mode).""" + import mlx.core as mx + + w = self.weights + + def memblock(base: int, xx: Any, past: Any) -> Any: + 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: + 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: + 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 DiT latents ``[B,C,T,H,W]`` with MLX TAEHV → pixels ``[B,T,H,W,3]`` in [0,1].""" + 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 and write MP4. Returns metrics dict including ``decode_s``.""" + 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) + 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/pipelines/stages/denoising.py b/fastvideo/pipelines/stages/denoising.py index 505f7e7f6d..a6a55a4f71 100644 --- a/fastvideo/pipelines/stages/denoising.py +++ b/fastvideo/pipelines/stages/denoising.py @@ -26,7 +26,7 @@ from fastvideo.pipelines.stages.validators import StageValidators as V from fastvideo.pipelines.stages.validators import VerificationResult from fastvideo.platforms import AttentionBackendEnum -from fastvideo.utils import dict_to_3d_list, masks_like +from fastvideo.utils import dict_to_3d_list, get_compute_dtype, masks_like try: from fastvideo.attention.backends.vmoba import VMOBAAttentionBackend @@ -44,6 +44,25 @@ logger = init_logger(__name__) + +def transformer_compute_dtype(transformer) -> torch.dtype: + """The dtype denoising inputs must match at the transformer's forwards. + + Under FSDP/HSDP mixed precision the raw parameters are fp32 masters that + the policy casts to ``param_dtype`` per-forward, so sniffing + ``next(parameters()).dtype`` returns fp32 -- which both leaves the inputs + fp32 and disables autocast, crashing against bf16-cast weights (observed + in training-time validation). Prefer the registered mixed-precision + policy; fall back to parameter sniffing when no policy is set (pure + inference, MPS fp16 path). + """ + policy_dtype = get_compute_dtype() + if policy_dtype != torch.float32: + return policy_dtype + module = transformer.module if hasattr(transformer, "module") else transformer + return next(module.parameters()).dtype + + class DenoisingStage(PipelineStage): """ Stage for running the denoising loop in diffusion pipelines. @@ -788,11 +807,7 @@ def forward( pipeline.add_module("transformer", self.transformer) fastvideo_args.model_loaded["transformer"] = True - if hasattr(self.transformer, "module"): - transformer_dtype = next(self.transformer.module.parameters()).dtype - else: - transformer_dtype = next(self.transformer.parameters()).dtype - target_dtype = transformer_dtype + target_dtype = transformer_compute_dtype(self.transformer) autocast_enabled = (target_dtype != torch.float32 and not fastvideo_args.disable_autocast) latents = batch.latents @@ -1264,10 +1279,9 @@ def forward( Returns: The batch with denoised latents. """ - # Setup precision and autocast settings - # TODO(will): make the precision configurable for inference - # target_dtype = PRECISION_TO_TYPE[fastvideo_args.precision] - target_dtype = torch.bfloat16 + # Match the loaded DiT instead of assuming CUDA BF16. MPS uses the + # FP16 compatibility path configured in FastVideoArgs. + target_dtype = transformer_compute_dtype(self.transformer) autocast_enabled = (target_dtype != torch.float32) and not fastvideo_args.disable_autocast # Get timesteps and calculate warmup steps diff --git a/fastvideo/platforms/mps.py b/fastvideo/platforms/mps.py index 3367c109ee..3771ec50d5 100644 --- a/fastvideo/platforms/mps.py +++ b/fastvideo/platforms/mps.py @@ -48,7 +48,12 @@ def get_current_memory_usage(cls, device: torch.types.Device | None = None) -> f @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 0000000000..01405867be --- /dev/null +++ b/fastvideo/tests/mlx/test_dmd_sampling.py @@ -0,0 +1,115 @@ +# 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. + 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: + 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: + 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_hardware_tier.py b/fastvideo/tests/mlx/test_hardware_tier.py new file mode 100644 index 0000000000..2543737eb6 --- /dev/null +++ b/fastvideo/tests/mlx/test_hardware_tier.py @@ -0,0 +1,167 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for hardware-adaptive MLX model tiering. + +Backend-agnostic: injects memory sizes so Metal is never required. Safe to run +under ``mlx[cpu]`` Linux CI and on Apple Silicon Metal boxes. +""" + +from __future__ import annotations + +from types import SimpleNamespace + +import pytest + +from fastvideo.mlx_runtime.hardware_tier import ( + DEFAULT_ASSUMED_MEMORY_GIB, + FIVE_B_MODEL_REPO, + MODEL_1_3B_REPO, + TIER_LARGE_MLX_CAP_GIB, + TIER_MEDIUM_MLX_CAP_GIB, + TIER_MEDIUM_MAX_GIB, + TIER_SMALL_MAX_GIB, + TIER_SMALL_MLX_CAP_GIB, + apply_tier_to_namespace, + detect_unified_memory_gib, + recommend_tier, +) + +# Fake 5B id used only to exercise the prefer_5b path (Track D not landed). +_FAKE_5B = "FastVideo/FakeWan2.2-TI2V-5B-Diffusers-NOT-PORTED" + + +@pytest.mark.parametrize( + ("memory_gib", "expect_name", "expect_repo", "expect_quant", "expect_cap", "expect_preset"), + [ + (8.0, "small", MODEL_1_3B_REPO, "int8", TIER_SMALL_MLX_CAP_GIB, "mac-16gb"), + (16.0, "small", MODEL_1_3B_REPO, "int8", TIER_SMALL_MLX_CAP_GIB, "mac-16gb"), + (18.0, "small", MODEL_1_3B_REPO, "int8", TIER_SMALL_MLX_CAP_GIB, "mac-16gb"), + # Just above small → medium, 1.3B fp16 fallback while FIVE_B_MODEL_REPO is None. + (18.01, "medium", MODEL_1_3B_REPO, "none", TIER_MEDIUM_MLX_CAP_GIB, "mac-32gb"), + (24.0, "medium", MODEL_1_3B_REPO, "none", TIER_MEDIUM_MLX_CAP_GIB, "mac-32gb"), + (32.0, "medium", MODEL_1_3B_REPO, "none", TIER_MEDIUM_MLX_CAP_GIB, "mac-32gb"), + (40.0, "medium", MODEL_1_3B_REPO, "none", TIER_MEDIUM_MLX_CAP_GIB, "mac-32gb"), + # Just above medium → large. + (40.01, "large", MODEL_1_3B_REPO, "none", TIER_LARGE_MLX_CAP_GIB, "mac-64gb"), + (48.0, "large", MODEL_1_3B_REPO, "none", TIER_LARGE_MLX_CAP_GIB, "mac-64gb"), + (64.0, "large", MODEL_1_3B_REPO, "none", TIER_LARGE_MLX_CAP_GIB, "mac-64gb"), + (128.0, "large", MODEL_1_3B_REPO, "none", TIER_LARGE_MLX_CAP_GIB, "mac-64gb"), + ], +) +def test_recommend_tier_injected_memory_1_3b_fallback( + memory_gib: float, + expect_name: str, + expect_repo: str, + expect_quant: str, + expect_cap: float, + expect_preset: str, +) -> None: + """Default path: FIVE_B_MODEL_REPO is unset → always 1.3B, quant by band.""" + assert FIVE_B_MODEL_REPO is None, "test assumes Track D 5B repo is not yet published" + tier = recommend_tier(memory_gib, prefer_5b=True) + assert tier.name == expect_name + assert tier.model_repo == expect_repo + assert tier.quantization == expect_quant + assert tier.mlx_memory_limit_gib == expect_cap + assert tier.benchmark_preset == expect_preset + assert tier.uses_5b is False + if expect_name == "small": + assert tier.decoder == "taehv" + assert tier.modes == "int8" + assert tier.decoders == "taehv" + assert tier.max_memory_gib == TIER_SMALL_MAX_GIB + elif expect_name == "medium": + assert tier.decoder == "taehv" + assert tier.modes == "fp16" + assert tier.max_memory_gib == TIER_MEDIUM_MAX_GIB + else: + assert tier.decoder == "wan-vae" + assert tier.modes == "fp16" + assert tier.max_memory_gib is None + + +@pytest.mark.parametrize( + ("memory_gib", "expect_name", "expect_quant", "expect_decoder"), + [ + (16.0, "small", "int8", "taehv"), # small never selects 5B + (32.0, "medium", "int8", "taehv"), + (64.0, "large", "none", "wan-vae"), + ], +) +def test_recommend_tier_prefer_5b_when_repo_known( + memory_gib: float, + expect_name: str, + expect_quant: str, + expect_decoder: str, +) -> None: + tier = recommend_tier(memory_gib, prefer_5b=True, five_b_model_repo=_FAKE_5B) + assert tier.name == expect_name + assert tier.quantization == expect_quant + assert tier.decoder == expect_decoder + if expect_name == "small": + assert tier.model_repo == MODEL_1_3B_REPO + assert tier.uses_5b is False + else: + assert tier.model_repo == _FAKE_5B + assert tier.uses_5b is True + + +def test_prefer_5b_false_keeps_1_3b_even_with_repo() -> None: + tier = recommend_tier(32.0, prefer_5b=False, five_b_model_repo=_FAKE_5B) + assert tier.name == "medium" + assert tier.model_repo == MODEL_1_3B_REPO + assert tier.quantization == "none" + assert tier.uses_5b is False + + +def test_non_positive_memory_falls_back_to_safe_small() -> None: + tier = recommend_tier(0.0) + assert tier.name == "small" + assert tier.model_repo == MODEL_1_3B_REPO + assert tier.quantization == "int8" + + +def test_detect_unified_memory_gib_positive_on_this_host() -> None: + """On the Mac under test this is ~36; on Linux CI it still returns a positive float.""" + gib = detect_unified_memory_gib() + assert isinstance(gib, float) + assert gib > 0.0 + # Sanity: not a nonsense petabyte reading. + assert gib < 1024.0 * 1024.0 + + +def test_detect_falls_back_to_default_when_probes_fail(monkeypatch: pytest.MonkeyPatch) -> None: + monkeypatch.setattr("fastvideo.mlx_runtime.hardware_tier._sysctl_memsize_bytes", lambda: None) + monkeypatch.setattr("fastvideo.mlx_runtime.hardware_tier._mlx_device_memory_bytes", lambda mx_module=None: None) + monkeypatch.setattr("fastvideo.mlx_runtime.hardware_tier._proc_meminfo_bytes", lambda: None) + assert detect_unified_memory_gib() == DEFAULT_ASSUMED_MEMORY_GIB + # Safe default maps to the small tier. + tier = recommend_tier(None) + # recommend_tier will re-call detect; still defaulted. + assert tier.name == "small" + + +def test_apply_tier_to_namespace_sets_modes_and_caps() -> None: + args = SimpleNamespace( + modes="fp16,bf16", + decoders="wan-vae", + mlx_memory_limit_gib=None, + mlx_disable_cache=False, + ) + tier = recommend_tier(16.0) + apply_tier_to_namespace(args, tier) + assert args.modes == "int8" + assert args.decoders == "taehv" + assert args.mlx_memory_limit_gib == TIER_SMALL_MLX_CAP_GIB + assert args.mlx_disable_cache is True + assert args.auto_tier_name == "small" + assert args.auto_tier_model_repo == MODEL_1_3B_REPO + assert args.auto_tier_quantization == "int8" + assert args.auto_tier_benchmark_preset == "mac-16gb" + + +def test_threshold_boundaries() -> None: + """Document the inclusive upper bounds: 18 → small, 40 → medium.""" + assert recommend_tier(TIER_SMALL_MAX_GIB).name == "small" + assert recommend_tier(TIER_SMALL_MAX_GIB + 1e-9).name == "medium" + assert recommend_tier(TIER_MEDIUM_MAX_GIB).name == "medium" + assert recommend_tier(TIER_MEDIUM_MAX_GIB + 1e-9).name == "large" diff --git a/fastvideo/tests/mlx/test_memory_limits.py b/fastvideo/tests/mlx/test_memory_limits.py new file mode 100644 index 0000000000..dd5276b2d9 --- /dev/null +++ b/fastvideo/tests/mlx/test_memory_limits.py @@ -0,0 +1,83 @@ +# 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: + self.calls.append(("memory", value)) + return 111 + + def set_cache_limit(self, value: int) -> int: + 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_affine_qat_parity.py b/fastvideo/tests/mlx/test_mlx_affine_qat_parity.py new file mode 100644 index 0000000000..f83664bf19 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_affine_qat_parity.py @@ -0,0 +1,153 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The M4 Phase A gate: torch fake-quant must match MLX's affine contract. + +The roadmap requires this parity to hold BEFORE any GPU is spent on a +quantization-aware training run: if train-time fake-quantization differs from +``mx.quantize``/``mx.dequantize`` at deploy time, the QAT gains evaporate. + +These tests compare the pure-torch reference in +``fastvideo/layers/quantization/mlx_affine_qat.py`` against the real MLX +implementation on both Metal and MLX CPU. Scales and biases are bit-exact. +The Metal quantizer and PyTorch can round a float16 division on opposite sides +of a code boundary despite the same stored scale/bias; codes are consequently +bounded to one adjacent bin and dequantized values to one source-dtype epsilon. +This is the precise portable contract for QAT, not a relaxed numerical-quality +test. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +mx = pytest.importorskip("mlx.core", reason="MLX is required for QAT numerics parity tests") + +from fastvideo.layers.quantization.mlx_affine_qat import ( # noqa: E402 + fake_quantize_mlx_affine, + mlx_affine_dequantize_reference, + mlx_affine_quantize_reference, +) + + +@pytest.fixture(params=[mx.gpu, mx.cpu], ids=["metal", "cpu"]) +def mlx_device(request: pytest.FixtureRequest): + """Run every parity check on both MLX backends and restore the default.""" + previous = mx.default_device() + mx.set_default_device(request.param) + try: + yield + finally: + mx.set_default_device(previous) + + +pytestmark = pytest.mark.usefixtures("mlx_device") + + +def _unpack_uint32_codes(packed: np.ndarray, *, bits: int, out_cols: int) -> np.ndarray: + """Unpack MLX's little-endian uint32 words into per-element integer codes.""" + el_per_word = 32 // bits + bitmask = (1 << bits) - 1 + words = packed.astype(np.uint64) + codes = np.zeros((*packed.shape[:-1], packed.shape[-1] * el_per_word), dtype=np.int32) + for k in range(el_per_word): + codes[..., k::el_per_word] = ((words >> (k * bits)) & bitmask).astype(np.int32) + return codes[..., :out_cols] + + +def _mlx_quantize(w_np: np.ndarray, *, group_size: int, bits: int): + q, scales, biases = mx.quantize(mx.array(w_np), group_size=group_size, bits=bits, mode="affine") + deq = mx.dequantize(q, scales, biases, group_size=group_size, bits=bits, mode="affine") + mx.eval(q, scales, biases, deq) + return (np.array(q), np.array(scales), np.array(biases), np.array(deq)) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +@pytest.mark.parametrize("bits", [8, 4]) +@pytest.mark.parametrize("shape", [(4, 128), (3, 64), (16, 256)]) +def test_codes_scales_biases_match_mlx(dtype: torch.dtype, bits: int, shape: tuple[int, int]) -> None: + torch.manual_seed(1234) + w = torch.randn(*shape, dtype=torch.float32).to(dtype) * 0.05 + + codes, scales, biases = mlx_affine_quantize_reference(w, group_size=64, bits=bits) + + w_np = w.float().numpy() if dtype == torch.float32 else w.numpy() + q_mlx, scales_mlx, biases_mlx, _ = _mlx_quantize(w_np, group_size=64, bits=bits) + + codes_mlx = _unpack_uint32_codes(q_mlx, bits=bits, out_cols=shape[-1]) + codes_flat = codes.reshape(shape[0], -1).numpy() + + # MLX's Metal kernel performs this division in Metal float while the + # training reference uses PyTorch float. With fp16 inputs one value can + # land on an adjacent integer bin; the serialized quantizer state remains + # bit-exact and no code can move by more than one bin. + np.testing.assert_array_less(np.abs(codes_flat - codes_mlx), 2) + np.testing.assert_array_equal(scales.float().numpy(), np.asarray(scales_mlx, dtype=np.float32)) + np.testing.assert_array_equal(biases.float().numpy(), np.asarray(biases_mlx, dtype=np.float32)) + + +@pytest.mark.parametrize("dtype", [torch.float32, torch.float16]) +def test_dequantized_weights_match_mlx_to_source_dtype_epsilon(dtype: torch.dtype) -> None: + torch.manual_seed(7) + w = torch.randn(8, 192, dtype=torch.float32).to(dtype) * 0.05 + + codes, scales, biases = mlx_affine_quantize_reference(w, group_size=64, bits=8) + deq = mlx_affine_dequantize_reference(codes, scales, biases, out_shape=w.shape) + + w_np = w.float().numpy() if dtype == torch.float32 else w.numpy() + _, _, _, deq_mlx = _mlx_quantize(w_np, group_size=64, bits=8) + + np.testing.assert_allclose( + deq.numpy(), + np.asarray(deq_mlx), + rtol=0, + atol=torch.finfo(dtype).eps, + ) + + +def test_fake_quantize_matches_deploy_pipeline_and_passes_gradients() -> None: + torch.manual_seed(99) + # bf16 master weights, exactly like a QAT training run. + w = (torch.randn(4, 128, dtype=torch.float32) * 0.05).to(torch.bfloat16).requires_grad_(True) + + fq = fake_quantize_mlx_affine(w, group_size=64, bits=8, simulate_dtype=torch.float16) + + # Forward matches the fp16 deployment representation within one storage + # epsilon; see the module contract for the Metal/PyTorch division boundary. + w_fp16 = w.detach().to(torch.float16) + _, _, _, deq_mlx = _mlx_quantize(w_fp16.numpy(), group_size=64, bits=8) + np.testing.assert_allclose( + fq.detach().to(torch.float16).numpy(), + np.asarray(deq_mlx), + rtol=0, + atol=torch.finfo(torch.float16).eps, + ) + + # Backward: straight-through — gradients reach the master weight unchanged. + fq.sum().backward() + assert w.grad is not None + np.testing.assert_array_equal(w.grad.float().numpy(), np.ones_like(w.grad.float().numpy())) + + +def test_quantized_matmul_close_to_fake_quant_linear() -> None: + torch.manual_seed(3) + w = (torch.randn(128, 256, dtype=torch.float32) * 0.05).to(torch.float16) + x = (torch.randn(2, 256, dtype=torch.float32) * 0.5).to(torch.float16) + + codes, scales, biases = mlx_affine_quantize_reference(w, group_size=64, bits=8) + deq = mlx_affine_dequantize_reference(codes, scales, biases, out_shape=w.shape) + y_torch = (x.float() @ deq.float().T) + + q, s, b = mx.quantize(mx.array(w.numpy()), group_size=64, bits=8, mode="affine") + y_mlx = mx.quantized_matmul( + mx.array(x.numpy()), q, s, b, transpose=True, group_size=64, bits=8, mode="affine") + mx.eval(y_mlx) + + # Accumulation order differs between frameworks, so this is a tolerance + # check (the serialized quantizer state is pinned exactly above). + np.testing.assert_allclose(y_torch.numpy(), np.array(y_mlx).astype(np.float32), atol=2e-2, rtol=2e-2) + + +def test_indivisible_group_size_raises() -> None: + with pytest.raises(ValueError, match="not divisible"): + mlx_affine_quantize_reference(torch.randn(4, 100), group_size=64, bits=8) diff --git a/fastvideo/tests/mlx/test_mlx_causal_attention.py b/fastvideo/tests/mlx/test_mlx_causal_attention.py new file mode 100644 index 0000000000..4eb54c082b --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_causal_attention.py @@ -0,0 +1,185 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Track C Rung 3: parity + KV-cache tests for MLX causal self-attention. + +Proves the porting insight the streaming runtime relies on — mask-free cached +decoding of one frame-block at a time equals a single block-causal *masked* pass +over the whole sequence — and checks the rolling eviction / sink-token +bookkeeping against explicit references. Backend-agnostic (Metal or mlx[cpu]). +""" + +from __future__ import annotations + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core", reason="MLX is required for the causal-attention tests") + +from fastvideo.mlx_runtime.causal import ( # noqa: E402 + MLXCausalKVCache, + causal_self_attention_step, + max_attention_size, +) +from fastvideo.mlx_runtime.fastwan import apply_rotary_emb # noqa: E402 + +RNG = np.random.default_rng(2026) + + +def _rand(*shape: int) -> "mx.array": + return mx.array(RNG.standard_normal(shape).astype(np.float32)) + + +def _qkv_cos_sin(num_tokens: int, num_heads: int, head_dim: int): + q = _rand(1, num_tokens, num_heads, head_dim) + k = _rand(1, num_tokens, num_heads, head_dim) + v = _rand(1, num_tokens, num_heads, head_dim) + # Rotary tables for absolute positions 0..num_tokens-1 (rope_dim == head_dim). + cos = _rand(num_tokens, head_dim) + sin = _rand(num_tokens, head_dim) + return q, k, v, cos, sin + + +def _block_causal_masked_reference(q, k, v, cos, sin, *, chunk_tokens: int, window: int, scale: float): + """Full-sequence attention with a block-causal (optionally windowed) mask. + + Token ``p`` attends every token strictly before the end of ``p``'s chunk and + no older than ``window`` tokens — the exact set the cached path accumulates. + """ + num_tokens = q.shape[1] + roped_q = apply_rotary_emb(q, cos, sin, is_neox_style=False) + roped_k = apply_rotary_emb(k, cos, sin, is_neox_style=False) + + ends = np.minimum(((np.arange(num_tokens) // chunk_tokens) + 1) * chunk_tokens, num_tokens) + q_idx = np.arange(num_tokens)[:, None] + kv_idx = np.arange(num_tokens)[None, :] + allowed = ((kv_idx < ends[:, None]) & (kv_idx >= (ends[:, None] - window))) | (q_idx == kv_idx) + mask = mx.array(np.where(allowed, 0.0, -np.inf).astype(np.float32)) + + out = mx.fast.scaled_dot_product_attention( + roped_q.transpose(0, 2, 1, 3), + roped_k.transpose(0, 2, 1, 3), + v.transpose(0, 2, 1, 3), + scale=scale, + mask=mask, + ).transpose(0, 2, 1, 3) + return out + + +def _run_cached(q, k, v, cos, sin, *, chunk_tokens, local_attn_size, frame_seqlen, kv_cache_size, sink_tokens=0): + num_tokens = q.shape[1] + num_heads, head_dim = q.shape[2], q.shape[3] + cache = MLXCausalKVCache.allocate( + batch=1, max_tokens=kv_cache_size, num_heads=num_heads, head_dim=head_dim, + sink_tokens=sink_tokens, dtype=mx.float32) + outputs = [] + for start in range(0, num_tokens, chunk_tokens): + end = start + chunk_tokens + out = causal_self_attention_step( + q[:, start:end], k[:, start:end], v[:, start:end], + cos[start:end], sin[start:end], cache, + current_start=start, local_attn_size=local_attn_size, frame_seqlen=frame_seqlen) + outputs.append(out) + return mx.concatenate(outputs, axis=1), cache + + +def test_cached_matches_block_causal_masked_no_eviction() -> None: + """local_attn_size=-1: chunked cached decode == block-causal masked full pass.""" + frame_seqlen, nfb, num_frames, num_heads, head_dim = 4, 2, 3, 2, 8 + chunk = frame_seqlen * nfb + n = num_frames * chunk + q, k, v, cos, sin = _qkv_cos_sin(n, num_heads, head_dim) + scale = head_dim**-0.5 + + cached, _ = _run_cached(q, k, v, cos, sin, chunk_tokens=chunk, local_attn_size=-1, + frame_seqlen=frame_seqlen, kv_cache_size=n) + ref = _block_causal_masked_reference(q, k, v, cos, sin, chunk_tokens=chunk, + window=max_attention_size(-1, frame_seqlen), scale=scale) + mx.eval(cached, ref) + np.testing.assert_allclose(np.array(cached), np.array(ref), atol=2e-4, rtol=2e-4) + + +def test_cached_matches_sliding_window_with_eviction() -> None: + """Limited local_attn_size: rolling-cache decode == sliding-window masked pass.""" + frame_seqlen, local_attn_size, num_frames, num_heads, head_dim = 4, 2, 5, 2, 8 + chunk = frame_seqlen # one frame per block + n = num_frames * chunk + window = local_attn_size * frame_seqlen # 8; cache holds exactly the window + q, k, v, cos, sin = _qkv_cos_sin(n, num_heads, head_dim) + scale = head_dim**-0.5 + + cached, cache = _run_cached(q, k, v, cos, sin, chunk_tokens=chunk, local_attn_size=local_attn_size, + frame_seqlen=frame_seqlen, kv_cache_size=window) + ref = _block_causal_masked_reference(q, k, v, cos, sin, chunk_tokens=chunk, window=window, scale=scale) + mx.eval(cached, ref) + np.testing.assert_allclose(np.array(cached), np.array(ref), atol=2e-4, rtol=2e-4) + # Indices advance to the full sequence / saturate at the window. + assert cache.global_end_index == n + assert cache.local_end_index == window + + +def test_cached_matches_sliding_window_overlapping_eviction() -> None: + """window > 2*chunk: eviction shifts overlapping source/dest regions. + + The adjacent-window case (window == 2*chunk) only moves non-overlapping + slices; production windows are larger, so this exercises the overlapping + shift path that would corrupt K/V if the rolled copy were not materialised. + """ + frame_seqlen, local_attn_size, num_frames, num_heads, head_dim = 4, 3, 6, 2, 8 + chunk = frame_seqlen # 4 + n = num_frames * chunk # 24 + window = local_attn_size * frame_seqlen # 12 (> 2*chunk, overlapping shift) + q, k, v, cos, sin = _qkv_cos_sin(n, num_heads, head_dim) + scale = head_dim**-0.5 + + cached, cache = _run_cached(q, k, v, cos, sin, chunk_tokens=chunk, local_attn_size=local_attn_size, + frame_seqlen=frame_seqlen, kv_cache_size=window) + ref = _block_causal_masked_reference(q, k, v, cos, sin, chunk_tokens=chunk, window=window, scale=scale) + mx.eval(cached, ref) + np.testing.assert_allclose(np.array(cached), np.array(ref), atol=2e-4, rtol=2e-4) + assert cache.global_end_index == n + assert cache.local_end_index == window + + +def test_chunk_exceeding_cache_capacity_raises() -> None: + """Reject chunks larger than the non-sink cache capacity (would clobber sinks).""" + frame_seqlen, num_heads, head_dim, sink_tokens = 4, 1, 8, 2 + # Capacity after sinks is 2; a 4-token chunk cannot fit without overwriting sinks. + kv_cache_size = sink_tokens + 2 + cache = MLXCausalKVCache.allocate( + batch=1, max_tokens=kv_cache_size, num_heads=num_heads, head_dim=head_dim, + sink_tokens=sink_tokens, dtype=mx.float32) + # Fill past the cache so the next write triggers overflow eviction. + q0, k0, v0, cos0, sin0 = _qkv_cos_sin(2, num_heads, head_dim) + causal_self_attention_step( + q0, k0, v0, cos0, sin0, cache, current_start=0, local_attn_size=2, frame_seqlen=frame_seqlen) + q1, k1, v1, cos1, sin1 = _qkv_cos_sin(4, num_heads, head_dim) + with pytest.raises(ValueError, match="exceeds available cache capacity"): + causal_self_attention_step( + q1, k1, v1, cos1, sin1, cache, current_start=2, local_attn_size=2, frame_seqlen=frame_seqlen) + + +def test_sink_tokens_preserved_across_eviction() -> None: + """The first ``sink_tokens`` cache slots survive rolling eviction unchanged.""" + frame_seqlen, local_attn_size, num_frames, num_heads, head_dim = 4, 3, 6, 1, 8 + chunk = frame_seqlen + n = num_frames * chunk + window = local_attn_size * frame_seqlen # 12 + sink_tokens = frame_seqlen # keep the first frame as attention sinks + q, k, v, cos, sin = _qkv_cos_sin(n, num_heads, head_dim) + + cache = MLXCausalKVCache.allocate(batch=1, max_tokens=window, num_heads=num_heads, head_dim=head_dim, + sink_tokens=sink_tokens, dtype=mx.float32) + sink_k_after_first = None + sink_v_after_first = None + for i, start in enumerate(range(0, n, chunk)): + end = start + chunk + causal_self_attention_step( + q[:, start:end], k[:, start:end], v[:, start:end], cos[start:end], sin[start:end], cache, + current_start=start, local_attn_size=local_attn_size, frame_seqlen=frame_seqlen) + if i == 0: + sink_k_after_first = np.array(cache.k[:, :sink_tokens]) + sink_v_after_first = np.array(cache.v[:, :sink_tokens]) + mx.eval(cache.k, cache.v) + # After many chunks (and at least one eviction), the sink region is untouched. + assert cache.global_end_index == n + np.testing.assert_array_equal(np.array(cache.k[:, :sink_tokens]), sink_k_after_first) + np.testing.assert_array_equal(np.array(cache.v[:, :sink_tokens]), sink_v_after_first) diff --git a/fastvideo/tests/mlx/test_mlx_causal_dit_parity.py b/fastvideo/tests/mlx/test_mlx_causal_dit_parity.py new file mode 100644 index 0000000000..d48ce8c633 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_causal_dit_parity.py @@ -0,0 +1,129 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Track C Rung 4: MLX causal DiT full-forward parity vs the torch reference. + +Feeds the same latent frame-blocks chunk-by-chunk through the torch +``CausalWanTransformer3DModel._forward_inference`` (dense-SDPA KV-cache path, +CPU) and the MLX ``MLXCausalWanDiT`` and asserts the streaming outputs match. +This is the model-level version of the mask-free-cached == block-causal +equivalence proven at the attention level in ``test_mlx_causal_attention.py``. +Tiny random-weight config; no real checkpoint needed. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +mx = pytest.importorskip("mlx.core", reason="MLX is required for the causal DiT parity test") + +from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig # noqa: E402 +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.causal_dit import MLXCausalWanDiT, MLXCausalWanTransformerBlock # noqa: E402 +from fastvideo.mlx_runtime.fastwan import mlx_block_weights_from_torch # noqa: E402 +from fastvideo.models.dits.causal_wanvideo import CausalWanTransformer3DModel # noqa: E402 +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch # noqa: E402 +from fastvideo.tests.mlx.tiny_wan import TOP_LEVEL_KEY_MAP # noqa: E402 + +SEED = 2026 +NUM_HEADS, HEAD_DIM = 4, 16 +NUM_LAYERS = 2 +NUM_FRAMES = 4 # latent frames; num_frames_per_block = 1 -> 4 chunks +HEIGHT, WIDTH = 8, 8 # latent HxW; patch (1,2,2) -> 4x4 = 16 tokens/frame +TEXT_LEN = 8 +ARCH = dict( + num_attention_heads=NUM_HEADS, attention_head_dim=HEAD_DIM, in_channels=16, out_channels=16, + text_dim=64, freq_dim=64, ffn_dim=128, num_layers=NUM_LAYERS, patch_size=(1, 2, 2), + rope_max_seq_len=64, local_attn_size=-1, sink_size=0, num_frames_per_block=1) +HF = dict( + num_attention_heads=NUM_HEADS, attention_head_dim=HEAD_DIM, in_channels=16, out_channels=16, + text_dim=64, freq_dim=64, ffn_dim=128, num_layers=NUM_LAYERS, patch_size=(1, 2, 2), text_len=TEXT_LEN, + rope_max_seq_len=64, eps=1e-6) + + +def _build_torch_model() -> CausalWanTransformer3DModel: + cfg = WanVideoConfig(arch_config=WanVideoArchConfig(**ARCH)) + model = CausalWanTransformer3DModel(config=cfg, hf_config=HF).to("cpu", torch.float32).eval() + torch.manual_seed(SEED + 3) + with torch.no_grad(): + for name, param in model.named_parameters(): + if param.ndim <= 1: + param.fill_(1.0) if (name.endswith("weight") and "norm" in name) else param.normal_(0.0, 0.02) + else: + torch.nn.init.xavier_uniform_(param) + return model + + +def _mlx_from_torch(model: CausalWanTransformer3DModel) -> MLXCausalWanDiT: + state = {name: value.detach().float() for name, value in model.state_dict().items()} + inner_dim = NUM_HEADS * 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 = [ + MLXCausalWanTransformerBlock(mlx_block_weights_from_torch(tb), dim=inner_dim, ffn_dim=HF["ffn_dim"], + num_heads=NUM_HEADS, eps=HF["eps"]) for tb in model.blocks + ] + config = dict(HF) + config["text_len"] = model.text_len # match the model's padding target (WanVideoConfig default 512) + return MLXCausalWanDiT(weights, blocks, config, local_attn_size=-1, sink_size=0, num_frames_per_block=1) + + +def _full_rotary(): + d = (NUM_HEADS * HEAD_DIM) // NUM_HEADS + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + cos, sin = get_rotary_pos_embed( + (NUM_FRAMES, HEIGHT // 2, WIDTH // 2), NUM_HEADS * HEAD_DIM, NUM_HEADS, rope_dim_list, + dtype=torch.float64, rope_theta=10000) + return cos, sin, mx.array(cos.float().numpy()), mx.array(sin.float().numpy()) + + +@pytest.mark.usefixtures("distributed_setup") +def test_mlx_causal_dit_matches_torch_streaming() -> None: + torch_model = _build_torch_model() + mlx_model = _mlx_from_torch(torch_model) + frame_seqlen = (HEIGHT // 2) * (WIDTH // 2) + + gen = torch.Generator().manual_seed(SEED + 1) + latents = torch.randn(1, ARCH["in_channels"], NUM_FRAMES, HEIGHT, WIDTH, generator=gen, dtype=torch.float32) + text = torch.randn(1, TEXT_LEN, ARCH["text_dim"], generator=gen, dtype=torch.float32) + _, _, cos_mx, sin_mx = _full_rotary() + + # torch reference: chunked _forward_inference with per-block kv/crossattn caches. + window = 21 * frame_seqlen + kv_cache = [{ + "k": torch.zeros(1, window, NUM_HEADS, HEAD_DIM), "v": torch.zeros(1, window, NUM_HEADS, HEAD_DIM), + "global_end_index": torch.tensor([0]), "local_end_index": torch.tensor([0]) + } for _ in range(NUM_LAYERS)] + crossattn_cache = [{ + "k": torch.zeros(1, TEXT_LEN, NUM_HEADS, HEAD_DIM), "v": torch.zeros(1, TEXT_LEN, NUM_HEADS, HEAD_DIM), + "is_init": False + } for _ in range(NUM_LAYERS)] + + torch_outs, mlx_outs = [], [] + mlx_kv, mlx_cx = mlx_model.allocate_caches(batch=1, frame_seqlen=frame_seqlen, dtype=mx.float32) + with torch.no_grad(), set_forward_context( + current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")): + for i in range(NUM_FRAMES): + chunk = latents[:, :, i:i + 1] + timestep = torch.tensor([[10]], dtype=torch.long) + torch_outs.append(torch_model( + hidden_states=chunk, encoder_hidden_states=text, timestep=timestep, + kv_cache=kv_cache, crossattn_cache=crossattn_cache, + current_start=i * frame_seqlen, cache_start=0, start_frame=i).detach().float().numpy()) + + cos_i = cos_mx[i * frame_seqlen:(i + 1) * frame_seqlen] + sin_i = sin_mx[i * frame_seqlen:(i + 1) * frame_seqlen] + out = mlx_model.forward_chunk( + mx.array(chunk.numpy()), mx.array(text.numpy()), mx.array(timestep.float().numpy()), + cos_i, sin_i, mlx_kv, mlx_cx, current_start=i * frame_seqlen) + mx.eval(out) + mlx_outs.append(np.array(out.astype(mx.float32))) + + torch_all = np.concatenate(torch_outs, axis=2) + mlx_all = np.concatenate(mlx_outs, axis=2) + np.testing.assert_allclose(mlx_all, torch_all, atol=2e-3, rtol=2e-3) diff --git a/fastvideo/tests/mlx/test_mlx_causal_dit_real_weights.py b/fastvideo/tests/mlx/test_mlx_causal_dit_real_weights.py new file mode 100644 index 0000000000..134000c733 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_causal_dit_real_weights.py @@ -0,0 +1,78 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Track C rung 4: real-weight smoke for MLXCausalWanDiT. + +Loads the released Self-Forcing causal checkpoint +(``wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers``) into the MLX causal DiT and runs a +few streaming chunks, asserting finite, correctly-shaped output. This exercises +the Diffusers loader + full 30-layer forward at real scale/dtype — complementary +to the numeric parity test, which runs on a tiny random-weight config. + +Gated on the checkpoint being present locally (set ``FASTVIDEO_SFWAN_ROOT`` or +place it at ``~/models/sfwan_t2v_1.3b``); skipped otherwise so CI and other +machines stay green. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core", reason="MLX is required for the real-weight smoke") + +_ROOT = Path(os.environ.get("FASTVIDEO_SFWAN_ROOT", str(Path.home() / "models" / "sfwan_t2v_1.3b"))) +_CHECKPOINT = _ROOT / "transformer" / "diffusion_pytorch_model.safetensors" +_CONFIG = _ROOT / "transformer" / "config.json" + +pytestmark = pytest.mark.skipif( + not (_CHECKPOINT.exists() and _CONFIG.exists()), + reason=f"SFWan causal checkpoint not found under {_ROOT} (set FASTVIDEO_SFWAN_ROOT)") + + +def test_real_causal_dit_streams_finite_output() -> None: + from fastvideo.mlx_runtime.causal_dit import mlx_causal_dit_from_diffusers_safetensors + + model = mlx_causal_dit_from_diffusers_safetensors( + _CHECKPOINT, _CONFIG, dtype="fp16", local_attn_size=-1, sink_size=0, num_frames_per_block=1) + + # Small latent shape for a fast smoke: 32x32 latent (patch 1x2x2 -> 16x16 tokens/frame). + in_channels = int(model.config["in_channels"]) + text_dim = int(model.config["text_dim"]) + height = width = 32 + frame_seqlen = (height // 2) * (width // 2) + num_chunks = 3 + + rng = np.random.default_rng(0) + text = mx.array((rng.standard_normal((1, 24, text_dim)) * 0.1).astype(np.float16)) + kv_caches, crossattn_caches = model.allocate_caches(batch=1, frame_seqlen=frame_seqlen, dtype=mx.float16) + + # Rotary tables for the full clip (frame-major), sliced per chunk. + from fastvideo.layers.rotary_embedding import get_rotary_pos_embed + import torch + + head_dim = int(model.config["attention_head_dim"]) + num_heads = int(model.config["num_attention_heads"]) + rope_dim_list = [head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)] + cos, sin = get_rotary_pos_embed( + (num_chunks, height // 2, width // 2), num_heads * head_dim, num_heads, rope_dim_list, + dtype=torch.float32, rope_theta=10000) + cos = mx.array(cos.float().numpy()) + sin = mx.array(sin.float().numpy()) + + for i in range(num_chunks): + chunk = mx.array((rng.standard_normal((1, in_channels, 1, height, width)) * 0.5).astype(np.float16)) + out = model.forward_chunk( + chunk, text, mx.array([[900.0]]), + cos[i * frame_seqlen:(i + 1) * frame_seqlen], sin[i * frame_seqlen:(i + 1) * frame_seqlen], + kv_caches, crossattn_caches, current_start=i * frame_seqlen) + mx.eval(out) + arr = np.array(out.astype(mx.float32)) + assert arr.shape == (1, in_channels, 1, height, width) + assert np.isfinite(arr).all(), f"chunk {i} produced non-finite output" + assert np.abs(arr).max() < 1e3, f"chunk {i} output magnitude {np.abs(arr).max()} unreasonable" + + # The KV cache accumulated all chunks; cross-attn cache was populated once. + assert kv_caches[0].global_end_index == num_chunks * frame_seqlen + assert crossattn_caches[0]["is_init"] is True diff --git a/fastvideo/tests/mlx/test_mlx_causal_long_rollout.py b/fastvideo/tests/mlx/test_mlx_causal_long_rollout.py new file mode 100644 index 0000000000..bcda7aa1bb --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_causal_long_rollout.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Long causal rollout: bounded KV memory under rolling eviction + sink tokens. + +Proves the north-star property of the streaming runtime: with a positive +``local_attn_size``, generating many more blocks than the attention window keeps +the K/V tensor size fixed at ``local_attn_size * frame_seqlen`` while +``global_end_index`` advances and ``local_end_index`` saturates at the window. +Backend-agnostic (Metal or mlx[cpu]); uses the tiny random-weight builders from +``test_mlx_causal_dit_parity``. +""" + +from __future__ import annotations + +import numpy as np +import pytest +import torch + +mx = pytest.importorskip("mlx.core", reason="MLX is required for the long-rollout test") + +from fastvideo.layers.rotary_embedding import get_rotary_pos_embed # noqa: E402 +from fastvideo.mlx_runtime.causal import max_attention_size # noqa: E402 +from fastvideo.mlx_runtime.causal_sampler import build_dmd_schedule, stream_causal_latents # noqa: E402 +from fastvideo.tests.mlx.test_mlx_causal_dit_parity import ( # noqa: E402 + ARCH, + HEAD_DIM, + HEIGHT, + NUM_HEADS, + WIDTH, + _build_torch_model, + _mlx_from_torch, +) + +# Many more blocks than the window so eviction + sinks are exercised end-to-end. +NUM_BLOCKS = 28 +LOCAL_ATTN_SIZE = 2 # frames kept in the rolling window +SINK_SIZE = 1 # frames protected as attention sinks +# Fewer DMD steps than production — control-flow / memory bound is the goal. +DMD_STEPS = [1000, 500] + + +def _rotary_for_frames(num_frames: int): + """Frame-major cos/sin tables long enough for the full rollout.""" + d = (NUM_HEADS * HEAD_DIM) // NUM_HEADS + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + cos, sin = get_rotary_pos_embed( + (num_frames, HEIGHT // 2, WIDTH // 2), + NUM_HEADS * HEAD_DIM, + NUM_HEADS, + rope_dim_list, + dtype=torch.float64, + rope_theta=10000, + ) + return mx.array(cos.float().numpy()), mx.array(sin.float().numpy()) + + +@pytest.mark.usefixtures("distributed_setup") +def test_long_rollout_kv_cache_bounded_independent_of_length() -> None: + """Stream many blocks; KV length stays at the window, local_end saturates.""" + model = _mlx_from_torch(_build_torch_model()) + # Windowed attention — the product path for videos longer than ~5s. + model.local_attn_size = LOCAL_ATTN_SIZE + model.sink_size = SINK_SIZE + + frame_seqlen = (HEIGHT // 2) * (WIDTH // 2) + window = max_attention_size(LOCAL_ATTN_SIZE, frame_seqlen) + assert window == LOCAL_ATTN_SIZE * frame_seqlen + + cos_full, sin_full = _rotary_for_frames(NUM_BLOCKS) + schedule, timesteps = build_dmd_schedule(DMD_STEPS, flow_shift=8.0, warp_denoising_step=True) + + rng = np.random.default_rng(7) + noise = mx.array( + rng.standard_normal((1, ARCH["in_channels"], NUM_BLOCKS, HEIGHT, WIDTH)).astype(np.float32)) + text = mx.array((rng.standard_normal((1, 24, ARCH["text_dim"])) * 0.1).astype(np.float32)) + + kv_caches, crossattn_caches = model.allocate_caches(batch=1, frame_seqlen=frame_seqlen, dtype=mx.float32) + # Pre-stream: allocated to exactly the attention window (bounded memory). + assert kv_caches[0].k.shape[1] == window + assert kv_caches[0].v.shape[1] == window + assert kv_caches[0].sink_tokens == SINK_SIZE * frame_seqlen + + blocks = list( + stream_causal_latents( + model, + text, + noise, + cos_full, + sin_full, + schedule, + timesteps, + frame_seqlen=frame_seqlen, + seed=0, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + )) + + assert len(blocks) == NUM_BLOCKS + for index, (block_index, latent) in enumerate(blocks): + assert block_index == index + arr = np.array(latent.astype(mx.float32)) + assert arr.shape == (1, ARCH["out_channels"], 1, HEIGHT, WIDTH) + assert np.isfinite(arr).all(), f"block {index} produced non-finite latents" + + # Memory bound holds after every block — independent of how far we are. + assert kv_caches[0].k.shape[1] == window + assert kv_caches[0].v.shape[1] == window + for layer_cache in kv_caches: + assert layer_cache.k.shape[1] == window + assert layer_cache.v.shape[1] == window + + # global_end advances with the full sequence; local_end saturates at window. + expected_global = NUM_BLOCKS * frame_seqlen # nfb == 1 + assert kv_caches[0].global_end_index == expected_global + assert kv_caches[0].local_end_index == window + # Cross-attn was populated once and reused across the long rollout. + assert crossattn_caches[0]["is_init"] is True + # Many more tokens generated than the window → eviction path was used. + assert expected_global > window diff --git a/fastvideo/tests/mlx/test_mlx_causal_long_rollout_real.py b/fastvideo/tests/mlx/test_mlx_causal_long_rollout_real.py new file mode 100644 index 0000000000..308e91ea04 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_causal_long_rollout_real.py @@ -0,0 +1,164 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Metal real-weight long stream: peak memory plateaus under rolling KV cache. + +Loads SFWan-1.3B when present, streams 21+ latent frames with a bounded +``local_attn_size``, and records time-to-first-block, steady per-block latency, +and peak MLX memory after each block. Asserts the K/V tensor length equals the +attention window for the whole run (memory bound independent of length). + +Skipped without Metal or without the local SFWan root so Linux ``mlx[cpu]`` CI +stays green. Numbers from a real Mac run are pasted into +``docs/design/apple_silicon_benchmark_baseline.md`` (Long-video streaming). +""" + +from __future__ import annotations + +import json +import os +import statistics +import time +from pathlib import Path + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core", reason="MLX is required for the real long-rollout test") + +_ROOT = Path(os.environ.get("FASTVIDEO_SFWAN_ROOT", str(Path.home() / "models" / "sfwan_t2v_1.3b"))) +_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() + +pytestmark = [ + pytest.mark.skipif(not _HAS_METAL, reason="Metal required for the real-weight long stream"), + pytest.mark.skipif(not _HAS_WEIGHTS, reason=f"SFWan checkpoint not found under {_ROOT}"), +] + +# Modest latent spatial size keeps the run tractable while still using real weights. +# Product shape (480×832) is exercised by the streaming demo; this gate proves the +# memory plateau property end-to-end on the released checkpoint. +LATENT_H = LATENT_W = 32 +NUM_FRAMES = 24 # > local_attn_size so eviction is exercised +LOCAL_ATTN_SIZE = 6 # frames +SINK_SIZE = 1 +DMD_STEPS = [1000, 750, 500, 250] + + +def test_real_long_stream_peak_memory_plateaus(tmp_path: Path) -> None: + from fastvideo.layers.rotary_embedding import get_rotary_pos_embed + from fastvideo.mlx_runtime.causal import max_attention_size + from fastvideo.mlx_runtime.causal_dit import mlx_causal_dit_from_diffusers_safetensors + from fastvideo.mlx_runtime.causal_sampler import build_dmd_schedule, stream_causal_latents + import torch + + model = mlx_causal_dit_from_diffusers_safetensors( + _CHECKPOINT, + _CONFIG, + dtype="fp16", + local_attn_size=LOCAL_ATTN_SIZE, + sink_size=SINK_SIZE, + num_frames_per_block=1, + ) + config = model.config + in_channels = int(config["in_channels"]) + text_dim = int(config["text_dim"]) + text_len = int(config.get("text_len", 512)) + frame_seqlen = (LATENT_H // 2) * (LATENT_W // 2) + window = max_attention_size(LOCAL_ATTN_SIZE, frame_seqlen) + + head_dim = int(config["attention_head_dim"]) + num_heads = int(config["num_attention_heads"]) + rope_dim_list = [head_dim - 4 * (head_dim // 6), 2 * (head_dim // 6), 2 * (head_dim // 6)] + cos_t, sin_t = get_rotary_pos_embed( + (NUM_FRAMES, LATENT_H // 2, LATENT_W // 2), + num_heads * head_dim, + num_heads, + rope_dim_list, + dtype=torch.float32, + rope_theta=10000, + ) + cos_full = mx.array(cos_t.float().numpy()) + sin_full = mx.array(sin_t.float().numpy()) + + rng = np.random.default_rng(0) + text = mx.array((rng.standard_normal((1, min(24, text_len), text_dim)) * 0.1).astype(np.float16)) + noise = mx.array( + rng.standard_normal((1, in_channels, NUM_FRAMES, LATENT_H, LATENT_W)).astype(np.float16)) + + schedule, timesteps = build_dmd_schedule(DMD_STEPS, flow_shift=8.0, warp_denoising_step=True) + kv_caches, crossattn_caches = model.allocate_caches(batch=1, frame_seqlen=frame_seqlen, dtype=mx.float16) + assert kv_caches[0].k.shape[1] == window + + mx.clear_cache() + mx.reset_peak_memory() + stream_start = time.perf_counter() + block_latencies: list[float] = [] + peak_gib_by_block: list[float] = [] + prev = stream_start + time_to_first = None + + for block_index, latent in stream_causal_latents( + model, + text, + noise, + cos_full, + sin_full, + schedule, + timesteps, + frame_seqlen=frame_seqlen, + seed=0, + kv_caches=kv_caches, + crossattn_caches=crossattn_caches, + ): + now = time.perf_counter() + block_latencies.append(now - prev) + prev = now + if time_to_first is None: + time_to_first = now - stream_start + peak_gib_by_block.append(mx.get_peak_memory() / (1024**3)) + arr = np.array(latent.astype(mx.float32)) + assert arr.shape[2] == 1 + assert np.isfinite(arr).all(), f"block {block_index} non-finite" + assert kv_caches[0].k.shape[1] == window + + total_s = time.perf_counter() - stream_start + assert kv_caches[0].global_end_index == NUM_FRAMES * frame_seqlen + assert kv_caches[0].local_end_index == window + assert len(block_latencies) == NUM_FRAMES + + steady = statistics.median(block_latencies[1:]) if len(block_latencies) > 1 else block_latencies[0] + # Peak memory after the window fills should not keep climbing with length. + # Allow small allocator noise; the plateau is the property under test. + mid = len(peak_gib_by_block) // 2 + late_max = max(peak_gib_by_block[mid:]) + early_at_window = peak_gib_by_block[min(LOCAL_ATTN_SIZE + SINK_SIZE, len(peak_gib_by_block) - 1)] + # Late peak within 25% of post-window peak (plateau, not O(T) growth). + assert late_max <= early_at_window * 1.25 + 0.05, ( + f"peak memory grew with length: early={early_at_window:.3f} late_max={late_max:.3f} " + f"series={peak_gib_by_block}") + + metrics = { + "model_root": str(_ROOT), + "num_frames": NUM_FRAMES, + "latent_hw": f"{LATENT_H}x{LATENT_W}", + "local_attn_size": LOCAL_ATTN_SIZE, + "sink_size": SINK_SIZE, + "frame_seqlen": frame_seqlen, + "kv_window_tokens": window, + "dmd_steps": len(timesteps), + "time_to_first_block_s": round(time_to_first, 3), + "block_latency_steady_s": round(steady, 3), + "block_latencies_s": [round(x, 3) for x in block_latencies], + "peak_gib_by_block": [round(x, 3) for x in peak_gib_by_block], + "peak_gib": round(max(peak_gib_by_block), 3), + "total_stream_s": round(total_s, 2), + "global_end_index": kv_caches[0].global_end_index, + "local_end_index": kv_caches[0].local_end_index, + } + out = tmp_path / "long_stream_metrics.json" + out.write_text(json.dumps(metrics, indent=2)) + # Also print so a local `pytest -s` run captures pasteable numbers for the baseline doc. + print("\n=== long-video streaming metrics (real SFWan) ===") + print(json.dumps(metrics, indent=2)) diff --git a/fastvideo/tests/mlx/test_mlx_causal_sampler.py b/fastvideo/tests/mlx/test_mlx_causal_sampler.py new file mode 100644 index 0000000000..5b278b40b8 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_causal_sampler.py @@ -0,0 +1,51 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Track C rung 5: unit test for the streaming causal DMD sampler. + +Runs the block-autoregressive sampler on the tiny random-weight causal model and +checks it streams one finite latent block per frame-block while the KV cache +advances across the whole clip. Numerics quality is out of scope here (that is a +visual check vs the CUDA reference); this guards the control flow and shapes. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core", reason="MLX is required for the streaming sampler test") + +from fastvideo.mlx_runtime.causal_sampler import build_dmd_schedule, stream_causal_latents # noqa: E402 +from fastvideo.tests.mlx.test_mlx_causal_dit_parity import ( # noqa: E402 + ARCH, + NUM_FRAMES, + _build_torch_model, + _full_rotary, + _mlx_from_torch, +) + + +@pytest.mark.usefixtures("distributed_setup") +def test_streaming_sampler_yields_finite_blocks() -> None: + model = _mlx_from_torch(_build_torch_model()) + height = width = 8 + frame_seqlen = (height // 2) * (width // 2) + _, _, cos_full, sin_full = _full_rotary() + + schedule, timesteps = build_dmd_schedule([1000, 750, 500, 250], flow_shift=8.0, warp_denoising_step=True) + assert len(timesteps) == 4 + + rng = np.random.default_rng(0) + noise = mx.array(rng.standard_normal((1, ARCH["in_channels"], NUM_FRAMES, height, width)).astype(np.float32)) + text = mx.array((rng.standard_normal((1, 24, ARCH["text_dim"])) * 0.1).astype(np.float32)) + + blocks = list( + stream_causal_latents( + model, text, noise, cos_full, sin_full, schedule, timesteps, frame_seqlen=frame_seqlen, seed=0)) + + # One block per latent frame (num_frames_per_block == 1), each finite and shaped right. + assert len(blocks) == NUM_FRAMES + for index, (block_index, latent) in enumerate(blocks): + assert block_index == index + arr = np.array(latent.astype(mx.float32)) + assert arr.shape == (1, ARCH["out_channels"], 1, height, width) + assert np.isfinite(arr).all(), f"block {index} produced non-finite latents" diff --git a/fastvideo/tests/mlx/test_mlx_checkpoint.py b/fastvideo/tests/mlx/test_mlx_checkpoint.py new file mode 100644 index 0000000000..98ee9e65d7 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_checkpoint.py @@ -0,0 +1,98 @@ +# 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: + 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_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 0000000000..62562bb7aa --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_compile_parity.py @@ -0,0 +1,87 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Regression coverage for the MLX FastWan ``mx.compile`` path. + +The dense MLX runtime may opt into ``mx.compile`` with +``FASTVIDEO_MLX_COMPILE=1``. Python scalar constants must stay on the MLX trace: +NumPy scalar arithmetic evaluates traced arrays and causes the runtime to fall +back to eager execution. +""" + +from __future__ import annotations + +import numpy as np +import pytest + +mx = pytest.importorskip("mlx.core", reason="MLX is required for compile-parity tests") + +from fastvideo.mlx_runtime.fastwan import ( # noqa: E402 + MLXWanTransformerBlock, + gelu_tanh, + timestep_embedding, +) + + +_RNG = np.random.default_rng(0) + + +def _rand(*shape: int, scale: float = 1.0) -> "mx.array": + return mx.array((_RNG.standard_normal(shape) * scale).astype(np.float32)) + + +def test_gelu_tanh_compiles_and_matches_eager() -> None: + x = _rand(1, 120, 64) + + eager = gelu_tanh(x) + compiled = mx.compile(gelu_tanh)(x) + mx.eval(eager, compiled) + + np.testing.assert_array_equal(np.array(eager), np.array(compiled)) + + +def test_timestep_embedding_compiles_and_matches_eager() -> None: + timestep = mx.array(np.array([0.0, 250.0, 500.0, 1000.0], dtype=np.float32)) + dim = 64 + + eager = timestep_embedding(timestep, dim) + compiled = mx.compile(lambda steps: timestep_embedding(steps, dim))(timestep) + mx.eval(eager, compiled) + + np.testing.assert_allclose(np.array(eager), np.array(compiled), rtol=1e-4, atol=1e-4) + + +def _tiny_block_weights(dim: int, ffn_dim: int) -> dict[str, "mx.array"]: + square = ("to_q", "to_k", "to_v", "to_out", "attn2.to_q", "attn2.to_k", "attn2.to_v", "attn2.to_out") + weights = {f"{name}.weight": _rand(dim, dim, scale=0.05) for name in square} + weights.update({f"{name}.bias": _rand(dim, scale=0.05) for name 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, dim, scale=0.05), + "ffn.fc_in.bias": _rand(ffn_dim, scale=0.05), + "ffn.fc_out.weight": _rand(dim, ffn_dim, scale=0.05), + "ffn.fc_out.bias": _rand(dim, scale=0.05), + }) + return weights + + +def test_transformer_block_compiles_and_matches_eager() -> None: + dim, num_heads, head_dim, ffn_dim, sequence, context = 64, 4, 16, 128, 120, 32 + block = MLXWanTransformerBlock( + _tiny_block_weights(dim, ffn_dim), dim=dim, ffn_dim=ffn_dim, num_heads=num_heads, eps=1e-6) + hidden_states = _rand(1, sequence, dim, scale=0.05) + encoder_hidden_states = _rand(1, context, dim, scale=0.05) + temb = _rand(1, 6, dim, scale=0.05) + cos = _rand(sequence, head_dim) + sin = _rand(sequence, head_dim) + + eager = block(hidden_states, encoder_hidden_states, temb, (cos, sin)) + compiled = mx.compile(lambda hidden, encoder, embedding, cosine, sine: block( + hidden, encoder, embedding, (cosine, sine)))(hidden_states, encoder_hidden_states, temb, cos, sin) + mx.eval(eager, compiled) + + np.testing.assert_allclose(np.array(eager), np.array(compiled), rtol=1e-5, atol=1e-5) 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 0000000000..9fb91a5f96 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_dit_parity.py @@ -0,0 +1,84 @@ +# 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: + 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_fastwan_benchmark.py b/fastvideo/tests/mlx/test_mlx_fastwan_benchmark.py new file mode 100644 index 0000000000..4e23cd0405 --- /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_wan22_i2v.py b/fastvideo/tests/mlx/test_mlx_wan22_i2v.py new file mode 100644 index 0000000000..f44d0396d9 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_wan22_i2v.py @@ -0,0 +1,118 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Track D Rung 4: I2V input-path parity for Wan2.2-TI2V-5B (DiT-level). + +Proves that first-latent-frame replacement + per-token timestep (frame0=0, +rest=video_t) yields matching outputs between torch ``WanTransformer3DModel`` +and ``MLXWan22DiT``. No VAE required — the "image" is a fixed latent frame. +Backend-agnostic (Metal or mlx[cpu]). +""" + +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 I2V parity test") + +from fastvideo.forward_context import set_forward_context # noqa: E402 +from fastvideo.mlx_runtime.wan22_i2v import ( # noqa: E402 + build_i2v_inputs, + build_i2v_per_token_timestep, + num_patch_tokens, + replace_first_latent_frame, + tokens_per_frame, +) +from fastvideo.pipelines.pipeline_batch_info import ForwardBatch # noqa: E402 +from fastvideo.tests.mlx.test_mlx_wan22_parity import _mlx_wan22_from_torch # noqa: E402 +from fastvideo.tests.mlx.tiny_wan import ( # noqa: E402 + TINY_ARCH, + build_hf_config, + build_tiny_wan_config, + build_torch_model, + mlx_rotary_embeddings, +) + + +def test_tokens_per_frame_frame_major_layout() -> None: + """Document frame-major token order used by patch_embed / I2V timesteps.""" + frames, height, width = 4, 8, 8 + patch = TINY_ARCH["patch_size"] + tpf = tokens_per_frame(height, width, patch) + assert tpf == (height // patch[1]) * (width // patch[2]) == 16 + assert num_patch_tokens(frames, height, width, patch) == (frames // patch[0]) * tpf == 64 + + ts = build_i2v_per_token_timestep( + batch=1, frames=frames, height=height, width=width, patch_size=patch, video_timestep=500.0) + assert ts.shape == (1, 64) + # First frame's tokens are clean (0); remaining frames noised (500). + assert np.all(ts[0, :tpf] == 0.0) + assert np.all(ts[0, tpf:] == 500.0) + + +def test_replace_first_latent_frame_numpy_and_torch() -> None: + rng = np.random.default_rng(0) + noise = rng.standard_normal((2, 16, 4, 8, 8)).astype(np.float32) + image = rng.standard_normal((2, 16, 8, 8)).astype(np.float32) + out = replace_first_latent_frame(noise, image) + np.testing.assert_array_equal(out[:, :, 0], image) + np.testing.assert_array_equal(out[:, :, 1:], noise[:, :, 1:]) + + noise_t = torch.from_numpy(noise) + image_t = torch.from_numpy(image) + out_t = replace_first_latent_frame(noise_t, image_t) + assert torch.equal(out_t[:, :, 0], image_t) + assert torch.equal(out_t[:, :, 1:], noise_t[:, :, 1:]) + + # The MLX branch must safely stage a torch image latent through CPU rather + # than relying on NumPy's unsupported direct conversion for accelerator + # tensors. CPU is used here so the regression runs on all MLX backends. + out_mx = replace_first_latent_frame(mx.array(noise), image_t) + np.testing.assert_array_equal(np.array(out_mx[:, :, 0]), image) + np.testing.assert_array_equal(np.array(out_mx[:, :, 1:]), noise[:, :, 1:]) + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan22_i2v_inputs_match_torch_per_token_path() -> None: + """I2V latents + timestep: MLXWan22DiT matches torch expand_timesteps path.""" + torch_model = build_torch_model() + hf_config = build_hf_config(build_tiny_wan_config()) + mlx_model = _mlx_wan22_from_torch(torch_model, hf_config) + + frames, height, width = 4, 8, 8 + patch = TINY_ARCH["patch_size"] + gen = torch.Generator().manual_seed(42) + noise = torch.randn(1, TINY_ARCH["in_channels"], frames, height, width, generator=gen, dtype=torch.float32) + # Distinct "image" content in frame 0 so a missed replacement would fail parity. + image_frame = torch.randn(1, TINY_ARCH["in_channels"], height, width, generator=gen, dtype=torch.float32) + text = torch.randn(1, 8, TINY_ARCH["text_dim"], generator=gen, dtype=torch.float32) + + latents_t, timestep_np = build_i2v_inputs( + noise, image_frame, video_timestep=500.0, image_timestep=0.0, patch_size=patch) + assert isinstance(latents_t, torch.Tensor) + assert torch.equal(latents_t[:, :, 0], image_frame) + timestep = torch.from_numpy(timestep_np).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=latents_t, encoder_hidden_states=text, timestep=timestep) + ref_np = ref.detach().float().numpy() + + # MLX path: same I2V construction via numpy/mx. + latents_mx, ts_np = build_i2v_inputs( + mx.array(noise.numpy()), + mx.array(image_frame.numpy()), + video_timestep=500.0, + patch_size=patch, + ) + freqs_cis = mlx_rotary_embeddings(latents_t) + out = mlx_model( + latents_mx if isinstance(latents_mx, mx.array) else mx.array(latents_mx), + mx.array(text.numpy()), + mx.array(ts_np), + freqs_cis, + ) + mx.eval(out) + mlx_np = np.array(out.astype(mx.float32)) + + np.testing.assert_allclose(mlx_np, ref_np, atol=2e-3, rtol=2e-3) 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 0000000000..3066c6af97 --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_wan22_parity.py @@ -0,0 +1,84 @@ +# 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) -> 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)) + + +@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) diff --git a/fastvideo/tests/mlx/test_mlx_wan22_real_weights.py b/fastvideo/tests/mlx/test_mlx_wan22_real_weights.py new file mode 100644 index 0000000000..701584dcad --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_wan22_real_weights.py @@ -0,0 +1,166 @@ +# 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 +import hashlib +from pathlib import Path + +import numpy as np +import pytest + +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" +_REVISION = os.environ.get("FASTVIDEO_WAN22_5B_REVISION") +_CHECKPOINT_SHA256 = os.environ.get("FASTVIDEO_WAN22_5B_SHA256") + +_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 +_HAS_PIN = bool(_REVISION and _CHECKPOINT_SHA256 and len(_CHECKPOINT_SHA256) == 64) + +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_PIN, + reason="Wan2.2-5B official source revision and transformer SHA256 are not pinned; set " + "FASTVIDEO_WAN22_5B_REVISION and FASTVIDEO_WAN22_5B_SHA256 after recording them in PORT_STATUS.md", + ), +] + + +def _sha256(path: Path) -> str: + 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 _load_torch_wan22_from_diffusers(checkpoint: Path, config_path: Path, *, dtype): + """Instantiate ``WanTransformer3DModel`` and load Diffusers-format 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) + model.load_state_dict(custom_sd, strict=True) + return model.to(dtype=dtype) + + +@pytest.mark.usefixtures("distributed_setup") +def test_wan22_real_weights_mlx_matches_torch_per_token_timestep() -> None: + 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 + + assert _sha256(_CHECKPOINT) == _CHECKPOINT_SHA256 + 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 ~ max 9e-2, + # mean 6e-3, cosine 0.99994 — assert with recorded headroom. + assert max_abs <= 0.15, f"max|Δ|={max_abs} over budget" + assert mean_abs <= 0.02, f"mean|Δ|={mean_abs} over budget" + assert cosine >= 0.999, f"cosine={cosine} too low" diff --git a/fastvideo/tests/mlx/test_mlx_wan_vae.py b/fastvideo/tests/mlx/test_mlx_wan_vae.py new file mode 100644 index 0000000000..ecb119898b --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_wan_vae.py @@ -0,0 +1,84 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan VAE / TAEHV decode tests — MLX TAEHV parity + denormalize helpers. + +- TAEHV MLX vs torch is bit-close (atol 1e-5) for z_dim=16 and 48. +- Full AutoencoderKLWan mean/std denormalize is checked against the torch path + used in ``mlx_wan_prompt_to_video`` (Metal-gated wall-clock optional). +""" + +from __future__ import annotations + +import os +from pathlib import Path + +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 + WanVAEConfigView, + decode_latents_taehv_mlx, + denormalize_latents_np, + ensure_taehv_checkpoint, +) + +_HAS_METAL = bool(getattr(mx, "metal", None) and mx.metal.is_available()) + + +@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 + + env_name = "FASTVIDEO_TAEW2_1_CHECKPOINT" if z_dim == 16 else "FASTVIDEO_TAEW2_2_CHECKPOINT" + checkpoint = os.environ.get(env_name) + if not checkpoint: + pytest.skip(f"{env_name} is not staged; this future-only test never downloads decoder artifacts") + ckpt = ensure_taehv_checkpoint(z_dim=z_dim, checkpoint_path=Path(checkpoint)) + 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] + tmin = min(torch_np.shape[0], mlx_np.shape[0]) + np.testing.assert_allclose(mlx_np[:tmin], torch_np[:tmin], atol=1e-5, rtol=1e-5) + + +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 (CPU).""" + import torch + + vae = os.environ.get("FASTVIDEO_WAN22_5B_VAE_DIR") + if not vae: + pytest.skip("FASTVIDEO_WAN22_5B_VAE_DIR is not staged; source/revision/checksum remain blocked in PORT_STATUS") + vae_dir = Path(vae) + 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_capability.py b/fastvideo/tests/mlx/test_quant_capability.py new file mode 100644 index 0000000000..5cdd965d2c --- /dev/null +++ b/fastvideo/tests/mlx/test_quant_capability.py @@ -0,0 +1,71 @@ +# 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: + 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 0000000000..25ed9ca9c6 --- /dev/null +++ b/fastvideo/tests/mlx/test_taehv_decode.py @@ -0,0 +1,15 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Offline contracts for the vendored TAEHV decoder helper.""" + +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 diff --git a/fastvideo/tests/mlx/test_wan22_sample.py b/fastvideo/tests/mlx/test_wan22_sample.py new file mode 100644 index 0000000000..052ce42fad --- /dev/null +++ b/fastvideo/tests/mlx/test_wan22_sample.py @@ -0,0 +1,19 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Unit tests for Wan2.2 warped DMD schedule construction.""" + +from __future__ import annotations + +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] diff --git a/fastvideo/tests/mlx/tiny_wan.py b/fastvideo/tests/mlx/tiny_wan.py new file mode 100644 index 0000000000..0550b4d6e4 --- /dev/null +++ b/fastvideo/tests/mlx/tiny_wan.py @@ -0,0 +1,217 @@ +# 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: + return WanVideoConfig(arch_config=WanVideoArchConfig(**TINY_ARCH)) + + +def build_hf_config(config: WanVideoConfig) -> dict[str, object]: + 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. + 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: + 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]: + 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: + 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): + """The rotary table the torch model builds internally, converted to MLX.""" + 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: + 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: + 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/modal/causal_cuda_reference.py b/fastvideo/tests/modal/causal_cuda_reference.py new file mode 100644 index 0000000000..67da972326 --- /dev/null +++ b/fastvideo/tests/modal/causal_cuda_reference.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CUDA reference dump for the MLX causal-DiT port (Track C verification). + +Two modes: + +``dump`` (run on a CUDA GPU, e.g. via ``launch_l40s_job.py``): builds the tiny +causal Wan model, runs ``_forward_inference`` chunk-by-chunk on the GPU with +deterministic (NumPy-seeded, platform-independent) inputs, and writes the model +weights + inputs + per-chunk outputs to an ``.npz`` on the Modal volume. + +``compare`` (run on the Mac): loads that ``.npz``, rebuilds the torch model from +the dumped weights, converts it to ``MLXCausalWanDiT``, replays the same chunked +inference on Metal, and asserts the MLX outputs match the CUDA reference. This +closes the one Track-C gate the Mac session could not: MLX-Metal vs real CUDA. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np + +# Tiny config shared by dump and compare (mirrors test_mlx_causal_dit_parity). +NUM_HEADS, HEAD_DIM, NUM_LAYERS = 4, 16, 2 +NUM_FRAMES, HEIGHT, WIDTH, TEXT_DIM = 4, 8, 8, 64 +ARCH = dict( + num_attention_heads=NUM_HEADS, attention_head_dim=HEAD_DIM, in_channels=16, out_channels=16, + text_dim=TEXT_DIM, freq_dim=64, ffn_dim=128, num_layers=NUM_LAYERS, patch_size=(1, 2, 2), + rope_max_seq_len=64, local_attn_size=-1, sink_size=0, num_frames_per_block=1) +HF = dict( + num_attention_heads=NUM_HEADS, attention_head_dim=HEAD_DIM, in_channels=16, out_channels=16, + text_dim=TEXT_DIM, freq_dim=64, ffn_dim=128, num_layers=NUM_LAYERS, patch_size=(1, 2, 2), text_len=512, + rope_max_seq_len=64, eps=1e-6) +SEED = 2026 + + +def _build_torch_model(): + import torch + + from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig + from fastvideo.models.dits.causal_wanvideo import CausalWanTransformer3DModel + + cfg = WanVideoConfig(arch_config=WanVideoArchConfig(**ARCH)) + model = CausalWanTransformer3DModel(config=cfg, hf_config=HF).eval() + torch.manual_seed(SEED + 3) + with torch.no_grad(): + for name, param in model.named_parameters(): + if param.ndim <= 1: + param.fill_(1.0) if (name.endswith("weight") and "norm" in name) else param.normal_(0.0, 0.02) + else: + torch.nn.init.xavier_uniform_(param) + return model + + +def _deterministic_inputs(): + rng = np.random.default_rng(SEED + 1) + latents = rng.standard_normal((1, ARCH["in_channels"], NUM_FRAMES, HEIGHT, WIDTH)).astype(np.float32) + text = rng.standard_normal((1, 24, TEXT_DIM)).astype(np.float32) + return latents, text + + +def _rotary(): + import torch + + from fastvideo.layers.rotary_embedding import get_rotary_pos_embed + + d = HEAD_DIM + rope_dim_list = [d - 4 * (d // 6), 2 * (d // 6), 2 * (d // 6)] + cos, sin = get_rotary_pos_embed( + (NUM_FRAMES, HEIGHT // 2, WIDTH // 2), NUM_HEADS * HEAD_DIM, NUM_HEADS, rope_dim_list, + dtype=torch.float32, rope_theta=10000) + return cos.float().numpy(), sin.float().numpy() + + +def dump(out_path: Path) -> None: + import torch + + from fastvideo.distributed.parallel_state import maybe_init_distributed_environment_and_model_parallel + from fastvideo.forward_context import set_forward_context + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + maybe_init_distributed_environment_and_model_parallel(1, 1) + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"dumping on device={device}", flush=True) + model = _build_torch_model().to(device, torch.float32) + latents_np, text_np = _deterministic_inputs() + frame_seqlen = (HEIGHT // 2) * (WIDTH // 2) + window = 21 * frame_seqlen + + latents = torch.from_numpy(latents_np).to(device) + text = torch.from_numpy(text_np).to(device) + kv = [{ + "k": torch.zeros(1, window, NUM_HEADS, HEAD_DIM, device=device), + "v": torch.zeros(1, window, NUM_HEADS, HEAD_DIM, device=device), + "global_end_index": torch.tensor([0], device=device), "local_end_index": torch.tensor([0], device=device) + } for _ in range(NUM_LAYERS)] + cx = [{ + "k": torch.zeros(1, 512, NUM_HEADS, HEAD_DIM, device=device), + "v": torch.zeros(1, 512, NUM_HEADS, HEAD_DIM, device=device), "is_init": False + } for _ in range(NUM_LAYERS)] + + outs = [] + with torch.no_grad(), set_forward_context( + current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")): + for i in range(NUM_FRAMES): + out = model( + hidden_states=latents[:, :, i:i + 1], encoder_hidden_states=text, + timestep=torch.tensor([[10]], device=device), kv_cache=kv, crossattn_cache=cx, + current_start=i * frame_seqlen, cache_start=0, start_frame=i) + outs.append(out.detach().float().cpu().numpy()) + + weights = {k: v.detach().float().cpu().numpy() for k, v in model.state_dict().items()} + out_path.parent.mkdir(parents=True, exist_ok=True) + np.savez( + out_path, device=np.array(device), latents=latents_np, text=text_np, + outputs=np.concatenate(outs, axis=2), **{f"w::{k}": v for k, v in weights.items()}) + print(f"wrote {out_path} ({len(weights)} weight tensors, device={device})", flush=True) + + +def compare(npz_path: Path, *, atol: float = 5e-3, rtol: float = 5e-3) -> None: + import mlx.core as mx + import torch + + from fastvideo.distributed.parallel_state import maybe_init_distributed_environment_and_model_parallel + + maybe_init_distributed_environment_and_model_parallel(1, 1) + data = np.load(npz_path) + ref_device = str(data["device"]) + ref_outputs = data["outputs"] + weights = {k[len("w::"):]: data[k] for k in data.files if k.startswith("w::")} + + # Rebuild the torch model with the *dumped* weights, then convert to MLX. + model = _build_torch_model() + model.load_state_dict({k: torch.from_numpy(v) for k, v in weights.items()}) + from fastvideo.tests.mlx.test_mlx_causal_dit_parity import _mlx_from_torch + + mlx_model = _mlx_from_torch(model) + + cos_np, sin_np = _rotary() + cos, sin = mx.array(cos_np), mx.array(sin_np) + frame_seqlen = (HEIGHT // 2) * (WIDTH // 2) + kv_caches, crossattn_caches = mlx_model.allocate_caches(batch=1, frame_seqlen=frame_seqlen, dtype=mx.float32) + + latents = data["latents"] + text = mx.array(data["text"]) + outs = [] + for i in range(NUM_FRAMES): + out = mlx_model.forward_chunk( + mx.array(latents[:, :, i:i + 1]), text, mx.array([[10.0]]), + cos[i * frame_seqlen:(i + 1) * frame_seqlen], sin[i * frame_seqlen:(i + 1) * frame_seqlen], + kv_caches, crossattn_caches, current_start=i * frame_seqlen) + mx.eval(out) + outs.append(np.array(out.astype(mx.float32))) + mlx_outputs = np.concatenate(outs, axis=2) + + max_abs = float(np.abs(mlx_outputs - ref_outputs).max()) + print(f"MLX-Metal vs torch-{ref_device}: max|Δ|={max_abs:.3e} (atol={atol})") + np.testing.assert_allclose(mlx_outputs, ref_outputs, atol=atol, rtol=rtol) + print("PASS: MLX causal port matches the CUDA reference.") + + +def main() -> None: + parser = argparse.ArgumentParser(description="CUDA reference dump / compare for the MLX causal DiT.") + parser.add_argument("mode", choices=("dump", "compare")) + parser.add_argument("--path", type=Path, default=Path("/root/data/causal_ref/causal_cuda_ref.npz")) + args = parser.parse_args() + if args.mode == "dump": + dump(args.path) + else: + compare(args.path) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/tests/modal/make_synth_t2v_parquet.py b/fastvideo/tests/modal/make_synth_t2v_parquet.py new file mode 100644 index 0000000000..2a5da5f048 --- /dev/null +++ b/fastvideo/tests/modal/make_synth_t2v_parquet.py @@ -0,0 +1,72 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Write a tiny synthetic T2V parquet in the ``pyarrow_schema_t2v`` format. + +For the SF+QAD training smoke: the values are random, but the schema/shapes are +exactly what ``collate_rows_from_parquet_schema`` and the Wan T2V dataloader +expect, so the training loop runs end-to-end. Not for quality — only to validate +that the recipe assembles, ``mlx_qat`` arms, and a few steps produce finite loss. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np +import pyarrow as pa +import pyarrow.parquet as pq + +from fastvideo.dataset.dataloader.schema import pyarrow_schema_t2v + + +def _tensor_fields(name: str, array: np.ndarray) -> dict: + return { + f"{name}_bytes": array.tobytes(), + f"{name}_shape": list(array.shape), + f"{name}_dtype": str(array.dtype), + } + + +def main() -> None: + parser = argparse.ArgumentParser(description="Synthesize a tiny T2V parquet for smoke training.") + parser.add_argument("--out-dir", type=Path, required=True) + parser.add_argument("--rows", type=int, default=8) + parser.add_argument("--channels", type=int, default=16) + parser.add_argument("--latent-t", type=int, default=3) + parser.add_argument("--latent-h", type=int, default=60) + parser.add_argument("--latent-w", type=int, default=104) + parser.add_argument("--text-seqlen", type=int, default=512) + parser.add_argument("--text-dim", type=int, default=4096) + args = parser.parse_args() + + rng = np.random.default_rng(0) + rows = [] + for i in range(args.rows): + latent = (rng.standard_normal((args.channels, args.latent_t, args.latent_h, args.latent_w)) * 0.5).astype( + np.float32) + text = (rng.standard_normal((args.text_seqlen, args.text_dim)) * 0.1).astype(np.float32) + row = { + "id": f"synth-{i:04d}", + "file_name": f"synth-{i:04d}.mp4", + "caption": "a synthetic smoke-test clip", + "media_type": "video", + "width": args.latent_w * 8, + "height": args.latent_h * 8, + "num_frames": (args.latent_t - 1) * 4 + 1, + "duration_sec": 5.0, + "fps": 16.0, + } + row.update(_tensor_fields("vae_latent", latent)) + row.update(_tensor_fields("text_embedding", text)) + rows.append(row) + + table = pa.Table.from_pylist(rows, schema=pyarrow_schema_t2v) + args.out_dir.mkdir(parents=True, exist_ok=True) + out_path = args.out_dir / "synth_00000.parquet" + pq.write_table(table, out_path) + print(f"wrote {out_path} ({args.rows} rows, latent {args.channels}x{args.latent_t}x{args.latent_h}x{args.latent_w}, " + f"text {args.text_seqlen}x{args.text_dim})", flush=True) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/tests/modal/sf_qad_smoke.yaml b/fastvideo/tests/modal/sf_qad_smoke.yaml new file mode 100644 index 0000000000..1dbaf8f4c2 --- /dev/null +++ b/fastvideo/tests/modal/sf_qad_smoke.yaml @@ -0,0 +1,95 @@ +# SF + QAD (INT8 mlx_qat) training smoke — Track C run-5 path validation. +# +# NOT a real recipe: the 14B teacher is swapped for 1.3B, data is a tiny +# synthetic parquet, validation/EMA are stripped, and only a couple of steps +# run. Goal: prove the self-forcing + mlx_qat recipe assembles, the QAD callback +# arms on the student, and a few steps produce finite loss on CUDA before a real +# run-5 launch. See fastvideo/tests/modal/make_synth_t2v_parquet.py. + +models: + student: + _target_: fastvideo.train.models.wan.WanCausalModel + init_from: wlsaidhi/SFWan2.1-T2V-1.3B-Diffusers + trainable: true + teacher: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers # smoke: 1.3B instead of 14B + trainable: false + disable_custom_init_weights: true + critic: + _target_: fastvideo.train.models.wan.WanModel + init_from: Wan-AI/Wan2.1-T2V-1.3B-Diffusers + trainable: true + disable_custom_init_weights: true + +method: + _target_: fastvideo.train.methods.distribution_matching.self_forcing.SelfForcingMethod + rollout_mode: simulate + generator_update_interval: 5 + real_score_guidance_scale: 4.0 + dmd_denoising_steps: [1000, 750, 500, 250] + warp_denoising_step: true + chunk_size: 3 + student_sample_type: sde + context_noise: 0.0 + same_step_across_blocks: true + enable_gradient_in_rollout: true + start_gradient_frame: 0 + fake_score_learning_rate: 4.0e-7 + fake_score_betas: [0.0, 0.999] + fake_score_lr_scheduler: constant + +training: + distributed: + num_gpus: 1 + sp_size: 1 + tp_size: 1 + hsdp_replicate_dim: 1 + hsdp_shard_dim: 1 + + data: + data_path: data/synth_t2v_smoke + dataloader_num_workers: 1 + train_batch_size: 1 + training_cfg_rate: 0.0 + seed: 1000 + num_latent_t: 3 + num_height: 480 + num_width: 832 + num_frames: 9 + + optimizer: + learning_rate: 2.0e-6 + betas: [0.0, 0.999] + weight_decay: 0.01 + lr_scheduler: constant + lr_warmup_steps: 0 + + loop: + max_train_steps: 2 + gradient_accumulation_steps: 1 + + checkpoint: + output_dir: outputs/sf_qad_smoke + training_state_checkpointing_steps: 0 + checkpoints_total_limit: 1 + + tracker: + project_name: sf-qad-smoke + run_name: sf_qad_int8_smoke + + model: + enable_gradient_checkpointing_type: full + +callbacks: + grad_clip: + max_grad_norm: 1.0 + mlx_qat: + group_size: 64 + bits: 8 + +pipeline: + flow_shift: 5 + dit_config: + local_attn_size: -1 + sink_size: 0 diff --git a/fastvideo/tests/modal/wan22_5b_qad_arming.py b/fastvideo/tests/modal/wan22_5b_qad_arming.py new file mode 100644 index 0000000000..85a08ebb1f --- /dev/null +++ b/fastvideo/tests/modal/wan22_5b_qad_arming.py @@ -0,0 +1,175 @@ +# SPDX-License-Identifier: Apache-2.0 +"""5B QAD arming check — run-6 launch gate (cheap, no training loop). + +Instantiates the Wan2.2-TI2V-5B student transformer from the FullAttn Diffusers +checkpoint (local path or HF id), applies +``MLXQuantizationAwareCallback`` (group_size=64, bits=8), and asserts it +fake-quantizes the expected weight count (~30 blocks × 10 + head ≈ 300+). + +Run on Modal (CUDA) so the HF download can land on the volume: + + modal run fastvideo/tests/modal/launch_l40s_job.py \\ + --command "python fastvideo/tests/modal/wan22_5b_qad_arming.py" \\ + --gpu-type L40S --num-gpus 1 --install-extra dev --pr-number \\ + --env-vars "MASTER_ADDR=localhost,MASTER_PORT=29561,FASTVIDEO_ATTENTION_BACKEND=TORCH_SDPA" \\ + --commit-volume + +Or locally if the transformer is already at ``FASTVIDEO_WAN22_5B_ROOT`` / +``~/models/fastwan22_5b``. +""" + +from __future__ import annotations + +import json +import os +import sys +from pathlib import Path + + +_MIN_EXPECTED_WEIGHTS = 300 # 30 blocks × ~10 + head/embedders + + +def _resolve_checkpoint() -> tuple[Path, Path]: + root = Path(os.environ.get("FASTVIDEO_WAN22_5B_ROOT", str(Path.home() / "models" / "fastwan22_5b"))) + ckpt = root / "transformer" / "diffusion_pytorch_model.safetensors" + cfg = root / "transformer" / "config.json" + if ckpt.exists() and cfg.exists() and ckpt.stat().st_size > 1_000_000_000: + print(f"[arming] using local checkpoint under {root}", flush=True) + return ckpt, cfg + + raise FileNotFoundError( + f"5B transformer is not staged under {root}. This internal arming gate never downloads artifacts; " + "set FASTVIDEO_WAN22_5B_ROOT to a reviewed local checkout whose source revision and SHA256 are " + "recorded in tests/local_tests/wan2_2_ti2v_5b/PORT_STATUS.md.") + + +def _load_transformer(checkpoint: Path, config_path: Path): + import torch + from safetensors.torch import load_file + + from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig + from fastvideo.distributed.parallel_state import maybe_init_distributed_environment_and_model_parallel + from fastvideo.models.dits.wanvideo import WanTransformer3DModel + from fastvideo.models.loader.utils import get_param_names_mapping, hf_to_custom_state_dict + + maybe_init_distributed_environment_and_model_parallel(1, 1) + + hf_config = json.loads(config_path.read_text()) + 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") + custom_sd, _ = hf_to_custom_state_dict(raw.items(), get_param_names_mapping(model.param_names_mapping)) + model.load_state_dict(custom_sd, strict=True) + # CPU fp32 is fine for arming (no forward needed). + return model.to(dtype=torch.float32) + + +def _import_mlx_qat_callback(): + """Load MLXQuantizationAwareCallback without ``fastvideo.train`` package init. + + ``fastvideo.train.__init__`` pulls the full Trainer stack (torchdata, pyarrow, + ...). Modal ``install-extra dev`` has those; a lean arming smoke may not. + """ + import importlib.util + import types as _types + + # Parent packages without running their heavy __init__ side effects. + for pkg in ("fastvideo.train", "fastvideo.train.callbacks"): + if pkg not in sys.modules: + sys.modules[pkg] = _types.ModuleType(pkg) + + if "fastvideo.train.callbacks.callback" not in sys.modules: + callback_mod = _types.ModuleType("fastvideo.train.callbacks.callback") + + class Callback: # noqa: D101 + pass + + callback_mod.Callback = Callback + sys.modules["fastvideo.train.callbacks.callback"] = callback_mod + + # parents: modal -> tests -> fastvideo + path = Path(__file__).resolve().parents[2] / "train" / "callbacks" / "mlx_qat.py" + spec = importlib.util.spec_from_file_location("fastvideo.train.callbacks.mlx_qat", path) + if spec is None or spec.loader is None: + raise RuntimeError(f"cannot load mlx_qat from {path}") + mod = importlib.util.module_from_spec(spec) + sys.modules["fastvideo.train.callbacks.mlx_qat"] = mod + spec.loader.exec_module(mod) + return mod.MLXQuantizationAwareCallback + + +def main() -> int: + from types import SimpleNamespace + + try: + from fastvideo.train.callbacks.mlx_qat import MLXQuantizationAwareCallback + except Exception as exc: # noqa: BLE001 - missing optional train deps on Mac + print(f"[arming] package import failed ({type(exc).__name__}: {exc}); " + "loading mlx_qat via file stub", flush=True) + MLXQuantizationAwareCallback = _import_mlx_qat_callback() + + ckpt, cfg = _resolve_checkpoint() + transformer = _load_transformer(ckpt, cfg) + n_params = sum(p.numel() for p in transformer.parameters()) + print(f"[arming] transformer params={n_params / 1e9:.2f}B layers={len(transformer.blocks)}", flush=True) + + method = SimpleNamespace(student=SimpleNamespace(transformer=transformer)) + callback = MLXQuantizationAwareCallback(group_size=64, bits=8, simulate_dtype="fp16") + try: + callback.on_train_start(method, iteration=0) + except ValueError as exc: + # Surface the silent-no-op failure mode explicitly. + print(f"FAIL: mlx_qat arming raised: {exc}", flush=True) + return 2 + except Exception as exc: # noqa: BLE001 + print(f"FAIL: unexpected error during mlx_qat arming: {type(exc).__name__}: {exc}", flush=True) + import traceback + traceback.print_exc() + return 3 + + n = len(callback.quantized_module_names) + print( + f"mlx_qat: fake-quantizing {n} weights (int8, group_size=64, simulate=torch.float16), " + f"e.g. {callback.quantized_module_names[:3]}", + flush=True, + ) + if n < _MIN_EXPECTED_WEIGHTS: + print(f"FAIL: expected >= {_MIN_EXPECTED_WEIGHTS} quantized weights, got {n}", flush=True) + return 1 + # Spot-check: modules still look vanilla outside forwards (FSDP safety). + sample_name = callback.quantized_module_names[0] + sample = dict(transformer.named_modules())[sample_name] + if not isinstance(sample.weight, type(next(transformer.parameters()))): + # Parameter check + pass + if "weight" not in sample._parameters: + print("FAIL: weight missing from _parameters outside forward (parametrization leak?)", flush=True) + return 4 + print(f"PASS: mlx_qat armed on 5B with {n} weights (>= {_MIN_EXPECTED_WEIGHTS})", flush=True) + return 0 + + +if __name__ == "__main__": + sys.exit(main()) diff --git a/fastvideo/tests/modal/wan22_cuda_reference.py b/fastvideo/tests/modal/wan22_cuda_reference.py new file mode 100644 index 0000000000..4caffea369 --- /dev/null +++ b/fastvideo/tests/modal/wan22_cuda_reference.py @@ -0,0 +1,186 @@ +# SPDX-License-Identifier: Apache-2.0 +"""CUDA reference dump/compare for the MLX Wan2.2 (per-token timestep) port. + +Two modes: + +``dump`` (CUDA via Modal L40S): builds a tiny dense Wan model, runs a single +forward with a 2-D per-token timestep (frame 0 at t=0, rest at t=500), and +writes weights + inputs + output to an ``.npz`` on the Modal volume. + +``compare`` (Mac): rebuilds the torch model from dumped weights, converts to +``MLXWan22DiT``, replays the same forward on Metal, and asserts a match. +""" + +from __future__ import annotations + +import argparse +from pathlib import Path + +import numpy as np + +# Tiny config shared by dump and compare (mirrors test_mlx_wan22_parity). +NUM_HEADS, HEAD_DIM, NUM_LAYERS = 4, 16, 2 +NUM_FRAMES, HEIGHT, WIDTH, TEXT_DIM, TEXT_LEN = 4, 8, 8, 64, 8 +ARCH = dict( + num_attention_heads=NUM_HEADS, + attention_head_dim=HEAD_DIM, + in_channels=16, + out_channels=16, + text_dim=TEXT_DIM, + freq_dim=64, + ffn_dim=128, + num_layers=NUM_LAYERS, + patch_size=(1, 2, 2), + rope_max_seq_len=64, +) +HF = dict( + num_attention_heads=NUM_HEADS, + attention_head_dim=HEAD_DIM, + in_channels=16, + out_channels=16, + text_dim=TEXT_DIM, + freq_dim=64, + ffn_dim=128, + num_layers=NUM_LAYERS, + patch_size=(1, 2, 2), + text_len=TEXT_LEN, + rope_max_seq_len=64, + eps=1e-6, +) +SEED = 2026 + + +def _build_torch_model(): + import torch + + from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig + from fastvideo.models.dits.wanvideo import WanTransformer3DModel + + cfg = WanVideoConfig(arch_config=WanVideoArchConfig(**ARCH)) + model = WanTransformer3DModel(config=cfg, hf_config=HF).eval() + 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_(0.0, 0.02) + else: + torch.nn.init.xavier_uniform_(param) + return model + + +def _deterministic_inputs(): + rng = np.random.default_rng(SEED + 1) + latents = rng.standard_normal((1, ARCH["in_channels"], NUM_FRAMES, HEIGHT, WIDTH)).astype(np.float32) + text = rng.standard_normal((1, TEXT_LEN, TEXT_DIM)).astype(np.float32) + p_t, p_h, p_w = ARCH["patch_size"] + tokens_per_frame = (HEIGHT // p_h) * (WIDTH // p_w) + num_tokens = (NUM_FRAMES // p_t) * tokens_per_frame + per_frame = [0] + [500] * (NUM_FRAMES // p_t - 1) + timestep = np.array([[per_frame[i // tokens_per_frame] for i in range(num_tokens)]], dtype=np.int64) + return latents, text, timestep + + +def dump(out_path: Path) -> None: + import torch + + from fastvideo.distributed.parallel_state import maybe_init_distributed_environment_and_model_parallel + from fastvideo.forward_context import set_forward_context + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + + maybe_init_distributed_environment_and_model_parallel(1, 1) + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"dumping wan22 reference on device={device}", flush=True) + model = _build_torch_model().to(device, torch.float32) + latents_np, text_np, timestep_np = _deterministic_inputs() + latents = torch.from_numpy(latents_np).to(device) + text = torch.from_numpy(text_np).to(device) + timestep = torch.from_numpy(timestep_np).to(device) + + with torch.no_grad(), set_forward_context( + current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")): + out = model(hidden_states=latents, encoder_hidden_states=text, timestep=timestep) + out_np = out.detach().float().cpu().numpy() + + weights = {k: v.detach().float().cpu().numpy() for k, v in model.state_dict().items()} + out_path.parent.mkdir(parents=True, exist_ok=True) + np.savez( + out_path, + device=np.array(device), + latents=latents_np, + text=text_np, + timestep=timestep_np, + output=out_np, + **{f"w::{k}": v for k, v in weights.items()}, + ) + print(f"wrote {out_path} ({len(weights)} weight tensors, device={device}, out_shape={out_np.shape})", flush=True) + + +def compare(npz_path: Path, *, atol: float = 5e-3, rtol: float = 5e-3) -> None: + import mlx.core as mx + import torch + + from fastvideo.distributed.parallel_state import maybe_init_distributed_environment_and_model_parallel + from fastvideo.mlx_runtime.fastwan import mlx_block_weights_from_torch + from fastvideo.mlx_runtime.wan22 import MLXWan22DiT, MLXWan22TransformerBlock + from fastvideo.tests.mlx.tiny_wan import TOP_LEVEL_KEY_MAP, mlx_rotary_embeddings + + maybe_init_distributed_environment_and_model_parallel(1, 1) + data = np.load(npz_path) + ref_device = str(data["device"]) + ref_output = data["output"] + weights = {k[len("w::"):]: data[k] for k in data.files if k.startswith("w::")} + + model = _build_torch_model() + model.load_state_dict({k: torch.from_numpy(v) for k, v in weights.items()}) + + state = {name: value.detach().float() for name, value in model.state_dict().items()} + inner_dim = NUM_HEADS * HEAD_DIM + mlx_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) + mlx_weights[mlx_name] = mx.array(tensor.numpy()) + blocks = [ + MLXWan22TransformerBlock( + mlx_block_weights_from_torch(tb), + dim=inner_dim, + ffn_dim=ARCH["ffn_dim"], + num_heads=NUM_HEADS, + eps=1e-6, + ) for tb in model.blocks + ] + mlx_model = MLXWan22DiT(mlx_weights, blocks, dict(HF)) + + latents = mx.array(data["latents"]) + text = mx.array(data["text"]) + timestep = mx.array(data["timestep"].astype(np.float32)) + freqs_cis = mlx_rotary_embeddings(torch.from_numpy(data["latents"])) + out = mlx_model(latents, text, timestep, freqs_cis) + mx.eval(out) + mlx_output = np.array(out.astype(mx.float32)) + + max_abs = float(np.abs(mlx_output - ref_output).max()) + print(f"MLX-Metal vs torch-{ref_device}: max|Δ|={max_abs:.3e} (atol={atol})") + np.testing.assert_allclose(mlx_output, ref_output, atol=atol, rtol=rtol) + print("PASS: MLX Wan2.2 per-token-timestep port matches the CUDA reference.") + + +def main() -> None: + parser = argparse.ArgumentParser(description="CUDA reference dump/compare for MLX Wan2.2 DiT.") + parser.add_argument("mode", choices=("dump", "compare")) + parser.add_argument("--path", type=Path, default=Path("/root/data/wan22_ref/ref.npz")) + parser.add_argument("--atol", type=float, default=5e-3) + parser.add_argument("--rtol", type=float, default=5e-3) + args = parser.parse_args() + if args.mode == "dump": + dump(args.path) + else: + compare(args.path, atol=args.atol, rtol=args.rtol) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/tests/modal/wan22_sampler_reference.py b/fastvideo/tests/modal/wan22_sampler_reference.py new file mode 100644 index 0000000000..d93bb53281 --- /dev/null +++ b/fastvideo/tests/modal/wan22_sampler_reference.py @@ -0,0 +1,208 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Dump / compare final DMD latents: torch FastWan2.2 pipeline vs MLX sampler. + +``dump`` (CUDA via Modal): run the torch densen DMD path with a fixed prompt/seed +and write final latents to an ``.npz`` on the volume. + +``compare`` (Mac): reload the dump, re-run MLX ``sample_wan22_dmd`` with the same +noise + text embeds, report max|Δ| / cosine. +""" + +from __future__ import annotations + +import argparse +import json +from pathlib import Path + +import numpy as np + +SEED = 1234 +RENOISE_SEED = 0 +PROMPT = "A red fox trotting through a snowy pine forest at golden hour, cinematic" +# Small enough for a fast dump but exercises multi-step DMD. +HEIGHT, WIDTH, NUM_FRAMES = 256, 448, 33 # latent 16x28x9 after /16 /4 +DMD_STEPS = [1000, 757, 522] +FLOW_SHIFT = 5.0 + + +def dump(out_path: Path, model_id: str = "FastVideo/FastWan2.2-TI2V-5B-FullAttn-Diffusers") -> None: + """Run torch VideoGenerator / pipeline DMD and save final latents.""" + import torch + + from fastvideo.distributed.parallel_state import maybe_init_distributed_environment_and_model_parallel + + maybe_init_distributed_environment_and_model_parallel(1, 1) + device = "cuda" if torch.cuda.is_available() else "cpu" + print(f"dumping torch 5B DMD latents on {device}", flush=True) + + # Prefer the public VideoGenerator API when available; fall back to a + # lightweight DiT-only loop matching MLX (same schedule math). + try: + from fastvideo import VideoGenerator + from fastvideo.configs.pipelines.wan import FastWan2_2_TI2V_5B_Config + + gen = VideoGenerator.from_pretrained( + model_id, + num_gpus=1, + ) + # VideoGenerator.generate returns pixels; we need latents — use internal DiT path below. + del gen + except Exception as exc: # noqa: BLE001 + print(f"VideoGenerator path unavailable ({exc}); using DiT-only dump", flush=True) + + # DiT-only dump: load transformer, run warped DMD in torch float32. + from diffusers import AutoencoderKLWan + from transformers import AutoTokenizer, UMT5EncoderModel + from huggingface_hub import snapshot_download + + from fastvideo.configs.models.dits.wanvideo import WanVideoArchConfig, WanVideoConfig + from fastvideo.forward_context import set_forward_context + from fastvideo.models.dits.wanvideo import WanTransformer3DModel + from fastvideo.models.loader.utils import get_param_names_mapping, hf_to_custom_state_dict + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + from fastvideo.models.utils import pred_noise_to_pred_video + from fastvideo.pipelines.pipeline_batch_info import ForwardBatch + from safetensors.torch import load_file + + root = Path(snapshot_download(model_id, allow_patterns=["transformer/*", "text_encoder/*", "tokenizer/*", "vae/*"])) + cfg = json.loads((root / "transformer" / "config.json").read_text()) + arch = WanVideoConfig(arch_config=WanVideoArchConfig( + num_attention_heads=int(cfg["num_attention_heads"]), + attention_head_dim=int(cfg["attention_head_dim"]), + in_channels=int(cfg["in_channels"]), + out_channels=int(cfg["out_channels"]), + text_dim=int(cfg["text_dim"]), + freq_dim=int(cfg["freq_dim"]), + ffn_dim=int(cfg["ffn_dim"]), + num_layers=int(cfg["num_layers"]), + patch_size=tuple(cfg["patch_size"]), + cross_attn_norm=bool(cfg.get("cross_attn_norm", True)), + qk_norm=cfg.get("qk_norm", "rms_norm_across_heads"), + eps=float(cfg.get("eps", 1e-6)), + rope_max_seq_len=int(cfg.get("rope_max_seq_len", 1024)), + )) + model = WanTransformer3DModel(config=arch, hf_config=cfg).to(device, torch.float16).eval() + raw = load_file(str(root / "transformer" / "diffusion_pytorch_model.safetensors"), device="cpu") + custom, _ = hf_to_custom_state_dict(raw.items(), get_param_names_mapping(model.param_names_mapping)) + model.load_state_dict(custom, strict=False) + model = model.to(device) + + tokenizer = AutoTokenizer.from_pretrained(root / "tokenizer") + text_enc = UMT5EncoderModel.from_pretrained(root / "text_encoder", torch_dtype=torch.float16).to(device).eval() + tokens = tokenizer([PROMPT], return_tensors="pt", padding="max_length", max_length=512, truncation=True) + with torch.no_grad(): + ehs = text_enc(tokens.input_ids.to(device)).last_hidden_state + + lat_h, lat_w = HEIGHT // 16, WIDTH // 16 + lat_t = (NUM_FRAMES - 1) // 4 + 1 + gen = torch.Generator(device="cpu").manual_seed(SEED) + noise = torch.randn(1, int(cfg["in_channels"]), lat_t, lat_h, lat_w, generator=gen, dtype=torch.float32) + latents = noise.to(device=device, dtype=torch.float16) + + scheduler = FlowMatchEulerDiscreteScheduler(shift=FLOW_SHIFT) + scheduler.set_timesteps(1000, device="cpu") + step_idx = torch.tensor(DMD_STEPS, dtype=torch.long) + warped = torch.cat((scheduler.timesteps.cpu(), torch.tensor([0.0]))) + timesteps = warped[1000 - step_idx] + # sigmas + def sigma_for(t: float) -> float: + idx = int(torch.argmin(torch.abs(scheduler.timesteps.cpu() - t)).item()) + return float(scheduler.sigmas.cpu()[idx]) + + torch.manual_seed(RENOISE_SEED) + # Use NumPy RNG for renoise so MLX can reproduce the exact sequence. + renoise_rng = np.random.default_rng(RENOISE_SEED) + pt, ph, pw = tuple(cfg["patch_size"]) + tokens_n = (lat_t // pt) * (lat_h // ph) * (lat_w // pw) + from fastvideo.models.utils import pred_noise_to_pred_video as torch_p2v + + with torch.no_grad(), set_forward_context( + current_timestep=0, attn_metadata=None, forward_batch=ForwardBatch(data_type="dummy")): + for i, t in enumerate(timesteps): + t_val = float(t.item()) + ts = torch.full((1, tokens_n), t_val, device=device, dtype=torch.long) + pred = model(hidden_states=latents, encoder_hidden_states=ehs, timestep=ts) + # Match MLX: pred_video = noise_input - sigma * pred_noise + sigma = sigma_for(t_val) + pred_video = latents.float() - sigma * pred.float() + if i < len(timesteps) - 1: + sigma_next = sigma_for(float(timesteps[i + 1].item())) + noise_r = torch.from_numpy( + renoise_rng.standard_normal(tuple(latents.shape)).astype(np.float32)).to(device) + latents = ((1.0 - sigma_next) * pred_video + sigma_next * noise_r).to(latents.dtype) + else: + latents = pred_video.to(latents.dtype) + + out_path = Path(out_path) + out_path.parent.mkdir(parents=True, exist_ok=True) + np.savez( + out_path, + latents=latents.float().cpu().numpy(), + noise=noise.numpy(), + text=ehs.float().cpu().numpy(), + prompt=np.array(PROMPT), + seed=np.array(SEED), + renoise_seed=np.array(RENOISE_SEED), + dmd_steps=np.array(DMD_STEPS), + flow_shift=np.array(FLOW_SHIFT), + height=np.array(HEIGHT), + width=np.array(WIDTH), + num_frames=np.array(NUM_FRAMES), + device=np.array(device), + ) + print(f"wrote {out_path} latents {tuple(latents.shape)}", flush=True) + + +def compare(npz_path: Path, dit_ckpt: Path, dit_config: Path) -> None: + import mlx.core as mx + import torch + + from examples.inference.basic.mlx_wan_prompt_to_video import make_rotary_embeddings + from fastvideo.mlx_runtime.wan22 import mlx_wan22_dit_from_diffusers_safetensors + from fastvideo.mlx_runtime.wan22_sample import sample_wan22_dmd + + data = np.load(npz_path, allow_pickle=True) + noise = mx.array(data["noise"]).astype(mx.float16) + text = mx.array(data["text"]).astype(mx.float16) + ref = data["latents"] + config = json.loads(Path(dit_config).read_text()) + lat_t, lat_h, lat_w = int(noise.shape[2]), int(noise.shape[3]), int(noise.shape[4]) + freqs = make_rotary_embeddings(config, latent_frames=lat_t, latent_height=lat_h, latent_width=lat_w) + model = mlx_wan22_dit_from_diffusers_safetensors(dit_ckpt, dit_config, dtype="fp16") + out = sample_wan22_dmd( + model, + text, + noise, + freqs, + dmd_denoising_steps=list(data["dmd_steps"]), + flow_shift=float(data["flow_shift"]), + warp_denoising_step=True, + seed=int(data["renoise_seed"]), + ) + mx.eval(out) + mlx_np = np.array(out.astype(mx.float32)) + max_abs = float(np.abs(mlx_np - ref).max()) + mean_abs = float(np.abs(mlx_np - ref).mean()) + cos = float(np.dot(mlx_np.ravel(), ref.ravel()) / (np.linalg.norm(mlx_np) * np.linalg.norm(ref) + 1e-12)) + print(f"MLX vs torch-{data['device']}: max|Δ|={max_abs:.3e} mean|Δ|={mean_abs:.3e} cosine={cos:.6f}") + # Full multi-step fp16 can drift; cosine is the decisive structural check. + if cos < 0.99: + raise SystemExit(f"FAIL: cosine {cos} < 0.99 — sampler likely wrong") + print("PASS: MLX sampler matches torch DMD (cosine ≥ 0.99)") + + +def main() -> None: + p = argparse.ArgumentParser() + p.add_argument("mode", choices=("dump", "compare")) + p.add_argument("--path", type=Path, default=Path("/root/data/wan22_sampler/ref.npz")) + p.add_argument("--dit-checkpoint", type=Path, default=Path.home() / "models/fastwan22_5b/transformer/diffusion_pytorch_model.safetensors") + p.add_argument("--dit-config", type=Path, default=Path.home() / "models/fastwan22_5b/transformer/config.json") + args = p.parse_args() + if args.mode == "dump": + dump(args.path) + else: + compare(args.path, args.dit_checkpoint, args.dit_config) + + +if __name__ == "__main__": + main() 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 0000000000..7bb1444a69 --- /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/tests/training/test_denoising_dtype.py b/fastvideo/tests/training/test_denoising_dtype.py new file mode 100644 index 0000000000..1141a50aaf --- /dev/null +++ b/fastvideo/tests/training/test_denoising_dtype.py @@ -0,0 +1,68 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Denoising stages must target the transformer's *compute* dtype. + +Regression for the DGX training-validation crash: under FSDP/HSDP mixed +precision the raw parameters are fp32 masters cast to bf16 per-forward, so +``next(parameters()).dtype`` said fp32 — leaving validation latents fp32 and +disabling autocast against bf16-cast weights +(``Input type (float) and bias type (c10::BFloat16) should be the same``). +""" + +from __future__ import annotations + +import torch + +from fastvideo.pipelines.stages.denoising import transformer_compute_dtype +from fastvideo.utils import _mixed_precision_state, set_mixed_precision_policy + + +def _with_policy_cleared(): + previous = getattr(_mixed_precision_state, "state", None) + if previous is not None: + del _mixed_precision_state.state + return previous + + +def _restore_policy(previous) -> None: + if previous is not None: + _mixed_precision_state.state = previous + elif hasattr(_mixed_precision_state, "state"): + del _mixed_precision_state.state + + +def test_policy_param_dtype_wins_over_fp32_masters() -> None: + previous = _with_policy_cleared() + try: + model = torch.nn.Linear(8, 8, dtype=torch.float32) # fp32 masters, like FSDP storage + set_mixed_precision_policy(param_dtype=torch.bfloat16, reduce_dtype=torch.float32) + assert transformer_compute_dtype(model) == torch.bfloat16 + finally: + _restore_policy(previous) + + +def test_falls_back_to_parameter_sniffing_without_policy() -> None: + previous = _with_policy_cleared() + try: + assert transformer_compute_dtype(torch.nn.Linear(8, 8).to(torch.float16)) == torch.float16 + assert transformer_compute_dtype(torch.nn.Linear(8, 8)) == torch.float32 + finally: + _restore_policy(previous) + + +def test_fp32_policy_defers_to_the_model() -> None: + previous = _with_policy_cleared() + try: + set_mixed_precision_policy(param_dtype=torch.float32, reduce_dtype=torch.float32) + assert transformer_compute_dtype(torch.nn.Linear(8, 8).to(torch.bfloat16)) == torch.bfloat16 + finally: + _restore_policy(previous) + + +def test_unwraps_module_attribute() -> None: + previous = _with_policy_cleared() + try: + wrapper = torch.nn.Module() + wrapper.module = torch.nn.Linear(8, 8).to(torch.float16) + assert transformer_compute_dtype(wrapper) == torch.float16 + finally: + _restore_policy(previous) diff --git a/fastvideo/tests/training/test_ema_portable_state.py b/fastvideo/tests/training/test_ema_portable_state.py new file mode 100644 index 0000000000..4e10200222 --- /dev/null +++ b/fastvideo/tests/training/test_ema_portable_state.py @@ -0,0 +1,105 @@ +# SPDX-License-Identifier: Apache-2.0 +"""EMA checkpoint state must be world-size-portable. + +Regression for the corrupted `--ema` export: the EMA shadow is a dict of +per-rank local shards keyed by live module names (with +activation-checkpointing wrapper prefixes). Checkpointed directly, DCP saved +only rank 0's quarter shards from a 4-GPU run, and a 1-GPU export loaded +them into full-shape buffers — producing noise weights whose uniform +INT8-vs-FP16 SSIM masqueraded as a QAT win. The callback now saves full +tensors with normalized names and re-slices on load. + +CPU tests cover the world-size-1 path and name normalization; the multi-rank +gather/scatter runs through the same DTensor helpers and is exercised by the +next DGX smoke run. +""" + +from __future__ import annotations + +import types + +import torch +import torch.distributed.checkpoint.stateful # noqa: F401 (import-order guard, see test_mlx_qat_callback) + +from fastvideo.train.callbacks.ema import EMACallback + + +class _Tiny(torch.nn.Module): + + def __init__(self) -> None: + super().__init__() + self.to_q = torch.nn.Linear(16, 16, bias=False) + self.proj_out = torch.nn.Linear(16, 8, bias=True) + + +def _method(transformer: torch.nn.Module): + return types.SimpleNamespace(student=types.SimpleNamespace(transformer=transformer)) + + +def _armed_callback(transformer: torch.nn.Module) -> EMACallback: + cb = EMACallback(decay=0.9, start_iter=0) + cb.on_train_start(_method(transformer)) + cb._ema_started = True + return cb + + +def test_state_dict_uses_clean_names_and_dcp_key() -> None: + torch.manual_seed(0) + model = _Tiny() + cb = _armed_callback(model) + # Simulate activation-checkpointing wrapper names in the live shadow. + cb.student_ema.shadow = { + f"{name.rsplit('.', 1)[0]}._checkpoint_wrapped_module.{name.rsplit('.', 1)[1]}": value + for name, value in cb.student_ema.shadow.items() + } + + state = cb.state_dict() + assert state["ema_started"] is True + # New DCP-native key; the legacy plain-shard key must be gone so old + # loaders cannot silently accept the new format's shapeless shards. + assert "student_ema" not in state + sharded = state["student_ema_sharded"] + # Names are normalized (AC wrapper stripped) so they match live params. + assert set(sharded) == {"to_q.weight", "proj_out.weight", "proj_out.bias"} + # On a non-distributed (world_size=1) model the params are plain tensors, + # so the shards round-trip at full shape. + for name, param in model.named_parameters(): + assert sharded[name].shape == param.shape + + +def test_round_trip_restores_shadow_for_current_param_names() -> None: + torch.manual_seed(1) + source_model = _Tiny() + source = _armed_callback(source_model) + for shard in source.student_ema.shadow.values(): + shard.add_(torch.randn_like(shard)) # make shadow distinct from init + saved = source.state_dict() + + target_model = _Tiny() + target = _armed_callback(target_model) + target._ema_started = False + target.load_state_dict(saved) + + assert target._ema_started is True + for name in dict(target_model.named_parameters()): + torch.testing.assert_close( + target.student_ema.shadow[name], source.student_ema.shadow[name], atol=0, rtol=0) + + # And the swapped-in weights match the source EMA, dtype-preserved. + with target.ema_context(target_model): + torch.testing.assert_close( + target_model.to_q.weight.detach().float(), + source.student_ema.shadow["to_q.weight"], + atol=1e-6, rtol=1e-6) + + +def test_legacy_per_shard_state_is_refused() -> None: + model = _Tiny() + cb = _armed_callback(model) + legacy = {"student_ema": {"to_q.weight": torch.zeros(4, 16)}, "ema_started": True} + try: + cb.load_state_dict(legacy) + except ValueError as exc: + assert "legacy" in str(exc).lower() or "per-shard" in str(exc) + else: + raise AssertionError("legacy per-shard EMA state must be refused, not silently loaded") diff --git a/fastvideo/tests/training/test_mlx_qat_callback.py b/fastvideo/tests/training/test_mlx_qat_callback.py new file mode 100644 index 0000000000..bbafbbb249 --- /dev/null +++ b/fastvideo/tests/training/test_mlx_qat_callback.py @@ -0,0 +1,154 @@ +# SPDX-License-Identifier: Apache-2.0 +"""The MLX QAT callback: fake-quantized forwards, straight-through training. + +Torch-only (no MLX needed): the underlying quantizer numerics are pinned +bitwise against MLX in ``test_mlx_affine_qat_parity.py``; these tests cover +the callback mechanics — module targeting, the forward-scoped weight swap, +gradient flow to the real weights, and requantization after optimizer steps. + +The forward-scoped swap (rather than ``torch.nn.utils.parametrize``) is +deliberate: under FSDP2 the parameters are sharded DTensors outside forwards +and unsharded plain tensors inside them, so the module must look completely +vanilla except during its own forward call. +""" + +from __future__ import annotations + +import types + +import pytest +import torch +# fastvideo.dataset.preprocessing_datasets references +# torch.distributed.checkpoint.stateful without importing the submodule and +# relies on an earlier import having loaded it; make that explicit here so +# this test does not depend on import order elsewhere in the suite. +import torch.distributed.checkpoint.stateful # noqa: F401 + +from fastvideo.layers.quantization.mlx_affine_qat import fake_quantize_mlx_affine +from fastvideo.train.callbacks.mlx_qat import MLXQuantizationAwareCallback + + +class _TinyStudentTransformer(torch.nn.Module): + """Names mirror the Wan layout the exclude patterns are written against.""" + + def __init__(self) -> None: + super().__init__() + self.to_q = torch.nn.Linear(128, 128, bias=False) + self.ffn_fc_in = torch.nn.Linear(128, 256, bias=True) + self.norm_q = torch.nn.Linear(128, 128, bias=False) # excluded by name + self.tiny = torch.nn.Linear(100, 32, bias=False) # indivisible -> skipped + self.patch_embedding = torch.nn.Conv3d(16, 128, kernel_size=(1, 2, 2)) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.ffn_fc_in(self.to_q(x)) + + +def _method_with_student(transformer: torch.nn.Module): + student = types.SimpleNamespace(transformer=transformer) + return types.SimpleNamespace(student=student) + + +def _expected_fq(weight: torch.Tensor) -> torch.Tensor: + flat = weight.reshape(weight.shape[0], -1) if weight.dim() > 2 else weight + fq = fake_quantize_mlx_affine(flat, group_size=64, bits=8) + return fq.reshape(weight.shape).to(weight.dtype) + + +def test_callback_targets_matrix_weights_and_skips_norms_and_indivisible() -> None: + transformer = _TinyStudentTransformer() + callback = MLXQuantizationAwareCallback(group_size=64, bits=8) + + callback.on_train_start(_method_with_student(transformer)) + + assert set(callback.quantized_module_names) == {"to_q", "ffn_fc_in", "patch_embedding"} + assert getattr(transformer.to_q, "_mlx_qat_wrapped", False) + assert not getattr(transformer.norm_q, "_mlx_qat_wrapped", False) + assert not getattr(transformer.tiny, "_mlx_qat_wrapped", False) + + +def test_module_stays_vanilla_outside_forward() -> None: + transformer = _TinyStudentTransformer() + MLXQuantizationAwareCallback(group_size=64, bits=8).on_train_start(_method_with_student(transformer)) + + # Outside a forward call the module must look untouched: `weight` is the + # real Parameter in _parameters (what FSDP, optimizers, checkpointing, + # and dtype sniffing all see). + assert isinstance(transformer.to_q.weight, torch.nn.Parameter) + assert "weight" in transformer.to_q._parameters + assert "weight" not in transformer.to_q.__dict__ + assert next(transformer.parameters()).dtype == torch.float32 + + +def test_forward_computes_with_the_deploy_grid() -> None: + torch.manual_seed(2) + transformer = _TinyStudentTransformer() + MLXQuantizationAwareCallback(group_size=64, bits=8).on_train_start(_method_with_student(transformer)) + + x = torch.randn(4, 128) + out = transformer.to_q(x) + expected = x @ _expected_fq(transformer.to_q.weight.detach()).T + torch.testing.assert_close(out, expected, atol=0, rtol=0) + + # Conv weights quantize over the flattened non-output dims. + latent = torch.randn(1, 16, 4, 8, 8) + conv = transformer.patch_embedding + out_conv = conv(latent) + expected_conv = torch.nn.functional.conv3d( + latent, _expected_fq(conv.weight.detach()), conv.bias, stride=conv.stride) + torch.testing.assert_close(out_conv, expected_conv, atol=0, rtol=0) + + +def test_gradients_flow_to_real_weights_and_requantize_after_step() -> None: + torch.manual_seed(11) + transformer = _TinyStudentTransformer() + MLXQuantizationAwareCallback(group_size=64, bits=8).on_train_start(_method_with_student(transformer)) + + weight = transformer.to_q.weight + optimizer = torch.optim.SGD(transformer.parameters(), lr=0.5) + + out = transformer(torch.randn(4, 128)) + out.square().mean().backward() + assert weight.grad is not None and weight.grad.abs().sum() > 0 + + before_weight = weight.detach().clone() + optimizer.step() + assert not torch.equal(weight.detach(), before_weight) + + # The next forward requantizes from the updated weight and stays on the + # deploy grid. + x = torch.randn(2, 128) + torch.testing.assert_close( + transformer.to_q(x), x @ _expected_fq(weight.detach()).T, atol=0, rtol=0) + + +def test_wrapping_is_idempotent() -> None: + transformer = _TinyStudentTransformer() + method = _method_with_student(transformer) + first = MLXQuantizationAwareCallback(group_size=64, bits=8) + first.on_train_start(method) + with pytest.raises(ValueError, match="matched no weights"): + # A second callback finds nothing left to wrap instead of + # double-quantizing. + MLXQuantizationAwareCallback(group_size=64, bits=8).on_train_start(method) + + +def test_no_matching_weights_raises() -> None: + with pytest.raises(ValueError, match="matched no weights"): + MLXQuantizationAwareCallback().on_train_start(_method_with_student(torch.nn.Module())) + + +def test_missing_student_raises() -> None: + with pytest.raises(ValueError, match="No student transformer"): + MLXQuantizationAwareCallback().on_train_start(types.SimpleNamespace(student=None)) + + +def test_weight_is_restored_even_when_forward_raises() -> None: + transformer = _TinyStudentTransformer() + MLXQuantizationAwareCallback(group_size=64, bits=8).on_train_start(_method_with_student(transformer)) + + with pytest.raises(RuntimeError): + transformer.to_q(torch.randn(2, 64)) # wrong input dim -> matmul error + + assert isinstance(transformer.to_q.weight, torch.nn.Parameter) + assert "weight" in transformer.to_q._parameters + assert "weight" not in transformer.to_q.__dict__ diff --git a/fastvideo/third_party/taehv/LICENSE b/fastvideo/third_party/taehv/LICENSE new file mode 100644 index 0000000000..6dc1edfbb0 --- /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 0000000000..0c9143c850 --- /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 0000000000..9feb522c3a --- /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/fastvideo/train/callbacks/callback.py b/fastvideo/train/callbacks/callback.py index ca057c151b..5ee614e06c 100644 --- a/fastvideo/train/callbacks/callback.py +++ b/fastvideo/train/callbacks/callback.py @@ -24,6 +24,7 @@ "grad_clip": "fastvideo.train.callbacks.grad_clip.GradNormClipCallback", "validation": "fastvideo.train.callbacks.validation.ValidationCallback", "ema": "fastvideo.train.callbacks.ema.EMACallback", + "mlx_qat": "fastvideo.train.callbacks.mlx_qat.MLXQuantizationAwareCallback", } diff --git a/fastvideo/train/callbacks/ema.py b/fastvideo/train/callbacks/ema.py index 2bd9f01f09..3b0ac56a0e 100644 --- a/fastvideo/train/callbacks/ema.py +++ b/fastvideo/train/callbacks/ema.py @@ -76,6 +76,9 @@ def on_train_start( decay=self._decay, mode="local_shard", ) + # Kept for checkpoint (de)serialization: converting between local + # shards and full tensors needs the live params' DTensor placements. + self._transformer = student.transformer logger.info( "EMA callback enabled (decay=%s, " "start_iter=%d).", @@ -136,11 +139,64 @@ def ema_context( # Checkpoint state # ---------------------------------------------------------- + # The EMA shadow is a dict of each rank's *local parameter shards* keyed + # by live module names (including activation-checkpointing wrapper + # prefixes). Checkpointing the plain shards directly is broken twice + # over: DCP deduplicates plain tensors as "replicated" so only rank 0's + # shard survives a multi-GPU save, and a checkpoint written at one world + # size cannot be loaded at another. Gathering full tensors at save time + # is also wrong: DTensor.full_tensor() is a collective, and issuing + # collectives from inside DCP's save path deadlocks (observed as a + # 10-minute NCCL timeout at the first checkpoint of the run-2 smoke). + # + # The correct mechanism is DCP's own: present each shard *as a DTensor* + # (DTensor.from_local is metadata-only, no communication). DCP then saves + # every rank's shard and reshards natively on load at any world size — + # exactly how the model weights under ``roles.*`` are handled. + + _AC_WRAPPER = "._checkpoint_wrapped_module" + + @classmethod + def _clean_name(cls, name: str) -> str: + return name.replace(cls._AC_WRAPPER, "") + + @staticmethod + def _as_dcp_tensor(shard: torch.Tensor, param: torch.Tensor) -> torch.Tensor: + from torch.distributed.tensor import DTensor + + if isinstance(param, DTensor): + return DTensor.from_local( + shard.to(device=param.device), + device_mesh=param.device_mesh, + placements=param.placements, + ) + return shard.detach().clone() + + @staticmethod + def _to_local_cpu(value: torch.Tensor) -> torch.Tensor: + from torch.distributed.tensor import DTensor + + if isinstance(value, DTensor): + value = value.to_local() + return value.detach().float().cpu() + def state_dict(self) -> dict[str, Any]: if self.student_ema is None: return {} + params = { + self._clean_name(name): param + for name, param in self._transformer.named_parameters() + } + shadow_dcp: dict[str, torch.Tensor] = {} + for name, shard in self.student_ema.shadow.items(): + clean = self._clean_name(name) + param = params.get(clean) + if param is None: + logger.warning("EMA shadow key %r has no matching parameter; dropping from checkpoint.", name) + continue + shadow_dcp[clean] = self._as_dcp_tensor(shard, param) return { - "student_ema": self.student_ema.state_dict(), + "student_ema_sharded": shadow_dcp, "ema_started": self._ema_started, } @@ -148,7 +204,22 @@ def load_state_dict( self, state_dict: dict[str, Any], ) -> None: - ema_state = state_dict.get("student_ema") - if (ema_state is not None and self.student_ema is not None): - self.student_ema.load_state_dict(ema_state) + if self.student_ema is not None: + sharded = state_dict.get("student_ema_sharded") + if sharded is not None: + shadow: dict[str, torch.Tensor] = {} + for name, _param in self._transformer.named_parameters(): + clean = self._clean_name(name) + if clean in sharded: + shadow[name] = self._to_local_cpu(sharded[clean]) + self.student_ema.shadow = shadow + elif state_dict.get("student_ema") is not None: + # Legacy plain-shard state: world-size-dependent and, on + # multi-GPU saves, missing every rank but 0. Refuse to load + # silently-corrupt weights. + raise ValueError( + "This checkpoint holds legacy per-shard EMA state, which is only valid " + "on the exact world size that wrote it and loses all non-rank-0 shards " + "on multi-GPU saves. The EMA in this checkpoint cannot be trusted; " + "resume without the EMA state or re-train with the portable format.") self._ema_started = bool(state_dict.get("ema_started", False), ) diff --git a/fastvideo/train/callbacks/mlx_qat.py b/fastvideo/train/callbacks/mlx_qat.py new file mode 100644 index 0000000000..0cdb489f1e --- /dev/null +++ b/fastvideo/train/callbacks/mlx_qat.py @@ -0,0 +1,162 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Quantization-aware training callback targeting the Apple/MLX runtime. + +Wraps the student transformer's target modules so every forward computes with +MLX-affine fake-quantized weights (the exact deploy-time grid of +``mx.quantize``/``mx.dequantize``; see +``fastvideo/layers/quantization/mlx_affine_qat.py`` and its bitwise parity +tests), while gradients flow straight-through to the real weights. Composes +with any ``TrainingMethod`` (DMD2, KD, fine-tune) via YAML: + +.. code-block:: yaml + + callbacks: + mlx_qat: + group_size: 64 + bits: 8 + +Mechanism: a forward-scoped weight swap, NOT ``torch.nn.utils.parametrize``. +Under FSDP2/HSDP the parameters are sharded DTensors *outside* module +forwards and unsharded plain tensors (already cast to the compute dtype) +*inside* them. Parametrizations restructure the parameter into a submodule +and compute from the raw master, which breaks both worlds: dtype sniffing +sees fp32 masters, and the forward mixes sharded DTensors with plain tensors +(``aten.convolution.default: got mixed torch.Tensor and DTensor``, observed +on a DGX B200). Swapping the weight for its fake-quantized version only for +the duration of each wrapped ``forward`` call means: outside forwards the +module is untouched (FSDP, optimizers, checkpointing, and export see vanilla +parameters), and inside forwards the fake-quant operates on exactly the +unsharded compute-dtype weight the matmul would have used. + +Targeting mirrors ``mlx_dit_from_diffusers_safetensors``: 2-D (or reshapable +conv) ``.weight`` tensors whose grouped dim divides ``group_size``, excluding +norms and modulation tables. The quantization *decisions* (codes, scales, +biases) bit-match deploy time; under bf16 compute the dequantized values are +rounded once more to bf16 for the matmul, like any bf16 arithmetic. +""" + +from __future__ import annotations + +import re +from typing import TYPE_CHECKING + +import torch + +from fastvideo.layers.quantization.mlx_affine_qat import fake_quantize_mlx_affine +from fastvideo.logger import init_logger +from fastvideo.train.callbacks.callback import Callback + +if TYPE_CHECKING: + from fastvideo.train.methods.base import TrainingMethod + +logger = init_logger(__name__) + +# Weights the MLX loader never quantizes: norms and modulation tables. +DEFAULT_EXCLUDE_PATTERNS = (r"norm", r"scale_shift_table") + +_SIMULATE_DTYPES = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} + +_WRAPPED_MARKER = "_mlx_qat_wrapped" + + +def _fake_quantize_weight(weight: torch.Tensor, *, group_size: int, bits: int, + simulate_dtype: torch.dtype) -> torch.Tensor: + original_shape = weight.shape + # Conv-style weights (e.g. Wan's patch embedding) quantize over the + # flattened non-output dims, matching the MLX loader's (out, -1) reshape. + weight2d = weight.reshape(original_shape[0], -1) if weight.dim() > 2 else weight + fq = fake_quantize_mlx_affine(weight2d, group_size=group_size, bits=bits, simulate_dtype=simulate_dtype) + # Keep the dtype the module was about to compute with (inside FSDP + # forwards that is the unsharded compute-dtype weight). + return fq.reshape(original_shape).to(weight.dtype) + + +def _install_qat_forward(module: torch.nn.Module, *, group_size: int, bits: int, + simulate_dtype: torch.dtype) -> None: + inner_forward = module.forward # bound method of this instance + + def qat_forward(*args, **kwargs): + original = module._parameters.pop("weight") + try: + # Plain-attribute shadow: getattr finds it before _parameters. + module.weight = _fake_quantize_weight( + original, group_size=group_size, bits=bits, simulate_dtype=simulate_dtype) + return inner_forward(*args, **kwargs) + finally: + if "weight" in module.__dict__: + del module.weight + module._parameters["weight"] = original + + module.forward = qat_forward + setattr(module, _WRAPPED_MARKER, True) + + +class MLXQuantizationAwareCallback(Callback): + """Apply MLX-affine fake quantization to the student's weights. + + Args (all YAML-configurable): + group_size: quantization group size along the input dim (MLX default 64). + bits: 8 (deploy target) or 4 (evaluated after INT8 proves out). + simulate_dtype: precision the deploy path casts weights to before + quantizing ("fp16" matches the MLX loader). + exclude_patterns: regex fragments; a weight is skipped when any + matches its module name. + """ + + def __init__( + self, + *, + group_size: int = 64, + bits: int = 8, + simulate_dtype: str = "fp16", + exclude_patterns: tuple[str, ...] | list[str] = DEFAULT_EXCLUDE_PATTERNS, + ) -> None: + self._group_size = int(group_size) + self._bits = int(bits) + self._simulate_dtype = _SIMULATE_DTYPES[simulate_dtype] + self._exclude = [re.compile(pattern) for pattern in exclude_patterns] + self.quantized_module_names: list[str] = [] + + def _is_target(self, module_name: str, module: torch.nn.Module) -> bool: + if getattr(module, _WRAPPED_MARKER, False): + return False + weight = module._parameters.get("weight") + if weight is None or weight.dim() < 2: + return False + if any(pattern.search(module_name) for pattern in self._exclude): + return False + grouped_dim = weight.shape[1:].numel() + if grouped_dim % self._group_size != 0: + logger.warning( + "mlx_qat: skipping %s — grouped dim %d is not divisible by group_size %d " + "(the MLX runtime could not quantize this weight either).", + module_name, grouped_dim, self._group_size) + return False + return True + + def on_train_start(self, method: TrainingMethod, iteration: int = 0) -> None: + student = getattr(method, "student", None) + if student is None or student.transformer is None: + raise ValueError("No student transformer found on method; cannot apply MLX QAT") + + self.quantized_module_names = [] + for module_name, module in student.transformer.named_modules(): + if not self._is_target(module_name, module): + continue + _install_qat_forward( + module, + group_size=self._group_size, + bits=self._bits, + simulate_dtype=self._simulate_dtype, + ) + self.quantized_module_names.append(module_name) + + if not self.quantized_module_names: + raise ValueError( + "mlx_qat matched no weights on the student transformer — check exclude_patterns " + "and group_size against the model architecture.") + logger.info( + "mlx_qat: fake-quantizing %d weights (int%d, group_size=%d, simulate=%s), e.g. %s", + len(self.quantized_module_names), self._bits, self._group_size, self._simulate_dtype, + self.quantized_module_names[:3], + ) diff --git a/fastvideo/train/entrypoint/dcp_to_diffusers.py b/fastvideo/train/entrypoint/dcp_to_diffusers.py index f2ef982817..ab849065c8 100644 --- a/fastvideo/train/entrypoint/dcp_to_diffusers.py +++ b/fastvideo/train/entrypoint/dcp_to_diffusers.py @@ -202,9 +202,15 @@ def convert( config_path: str | None = None, role: str = "student", overwrite: bool = False, + use_ema: bool = False, ) -> str: """Load a DCP checkpoint and export as a diffusers model. + With ``use_ema`` the EMA callback's shadow weights (checkpointed under + the DCP ``callbacks`` entry) are swapped into the role's transformer for + the export, so the written model carries the EMA weights instead of the + raw ones. + Returns the path to the exported model directory. """ _ensure_distributed() @@ -265,11 +271,33 @@ def convert( # -- Load DCP weights into the model -- states = method.checkpoint_state() + ema_cb = None + if use_ema: + from fastvideo.train.callbacks.callback import CallbackDict + from fastvideo.train.callbacks.ema import EMACallback + from fastvideo.train.utils.checkpoint import _CallbackStateWrapper + + # Rebuild the callbacks so their state namespace matches what the + # CheckpointManager saved, then let the EMA callback allocate its + # shadow container before dcp.load fills it. + callbacks = CallbackDict(cfg.callbacks or {}, tc) + ema_cb = next( + (cb for cb in callbacks._callbacks.values() if isinstance(cb, EMACallback)), + None, + ) + if ema_cb is None: + raise ValueError("--ema requested but the run config declares no EMA callback.") + ema_cb.on_train_start(method) + states["callbacks"] = _CallbackStateWrapper(callbacks) logger.info( "Loading DCP checkpoint from %s", resolved, ) dcp.load(states, checkpoint_id=str(dcp_dir)) + if ema_cb is not None and not ema_cb._ema_started: + raise ValueError( + "--ema requested but the checkpoint contains no started EMA state " + "(ema_started is false). Export without --ema instead.") # -- Export to diffusers format -- model = method._role_models[role] @@ -285,13 +313,18 @@ def convert( output_dir, base_model_path, ) - result = _save_role_pretrained( - role=role, - base_model_path=base_model_path, - output_dir=output_dir, - overwrite=overwrite, - model=model, - ) + import contextlib + + ema_ctx = (ema_cb.ema_context(model.transformer) + if ema_cb is not None else contextlib.nullcontext()) + with ema_ctx: + result = _save_role_pretrained( + role=role, + base_model_path=base_model_path, + output_dir=output_dir, + overwrite=overwrite, + model=model, + ) logger.info("Export complete: %s", result) return result @@ -401,6 +434,11 @@ def main() -> None: action="store_true", help="Overwrite output-dir if it exists.", ) + parser.add_argument( + "--ema", + action="store_true", + help="Export the EMA shadow weights instead of the raw role weights.", + ) args = parser.parse_args(sys.argv[1:]) convert( @@ -408,6 +446,7 @@ def main() -> None: output_dir=args.output_dir, config_path=args.config, role=args.role, + use_ema=args.ema, overwrite=args.overwrite, ) diff --git a/mkdocs.yml b/mkdocs.yml index df1d9b2990..91f0dfd313 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: @@ -192,6 +193,7 @@ nav: - Debugging: utilities/debugging.md - Design: - Overview: design/overview.md + - Apple Silicon QAD Runbook: design/apple_silicon_qad_runbook.md - Training Architecture: design/training_architecture.md - Server Contracts: - Overview: design/server_contracts/index.md diff --git a/pyproject.toml b/pyproject.toml index 305240156a..68e1219d7b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -33,7 +33,9 @@ dependencies = [ "torch==2.12.0", "torchvision", "torchaudio", - "flashinfer-python", + # FlashInfer ships CUDA kernels and is not needed by the Apple-native MLX + # release path (which uses MLX for the DiT and PyTorch MPS for encode/decode). + "flashinfer-python; platform_system != 'Darwin'", # Acceleration & Optimization "accelerate==1.0.1", @@ -71,7 +73,10 @@ dependencies = [ "remote-pdb", # Kernel & Packaging - "fastvideo-kernel==0.3.2", + # The MLX release does not load FastVideo's CUDA kernel package. Keeping it + # out of macOS resolution makes `uv pip install -e '.[mlx]'` a source-only + # Apple Silicon install rather than a CUDA-wheel lookup. + "fastvideo-kernel==0.3.2; platform_system != 'Darwin'", "wheel", # Training Dependencies @@ -120,6 +125,10 @@ flash-attn-4 = { git = "https://github.com/Dao-AILab/flash-attention.git", rev = # flash-attn: uv pip install flash-attn==2.8.1 --no-cache-dir --no-build-isolation +# Apple Silicon FastWan runtime and QAT contract. MLX 0.31.2 is the exact +# Metal/CPU version covered by the checked-in affine parity tests. +mlx = ["mlx==0.31.2; platform_system == 'Darwin' and platform_machine == 'arm64'"] + swanlab = [ "swanlab>=0.6.7", ] diff --git a/tests/local_tests/wan2_2_ti2v_5b/PORT_STATUS.md b/tests/local_tests/wan2_2_ti2v_5b/PORT_STATUS.md new file mode 100644 index 0000000000..caa8c17545 --- /dev/null +++ b/tests/local_tests/wan2_2_ti2v_5b/PORT_STATUS.md @@ -0,0 +1,79 @@ +# Wan2.2-TI2V-5B Port Status + +## Summary +- model_family: `wan2_2_ti2v_5b` +- workload_types: T2V and I2V (future-only; neither is public) +- official_ref: unpinned historical candidate only; no approved source/revision/checksum +- official_ref_dir: none +- hf_weights_path: unknown +- local_weights_dir: none +- source_layout: unknown +- local_tests_readme: `tests/local_tests/wan2_2_ti2v_5b/README.md` + +## Current Phase +- phase: 0 — preparation +- status: blocked +- owner: prep +- last_updated: 2026-07-09 + +## Component Matrix +| Component | Type | Reuse/Port | Official Definition | Official Instantiation | FastVideo Target | Prototype | Conversion | Parity | Open Issues | +|---|---|---|---|---|---|---|---|---|---| +| Transformer | DiT | internal MLX prototype | unpinned historical candidate | unknown | `fastvideo.mlx_runtime.wan22.MLXWan22DiT` | tiny pass | blocked | real scaffold skip | Q001, Q002, Q005 | +| VAE | video VAE | internal decode helper | unpinned historical candidate | unknown | `fastvideo.mlx_runtime.wan_vae` | helper only | blocked | scaffold skip | Q001, Q003, Q005 | +| Text conditioner | encoder | undecided | unknown | unknown | undecided | blocked | blocked | scaffold skip | Q001 | +| Image conditioner/mask | encoder/conditioner | latent-only helper | unknown | unknown | `fastvideo.mlx_runtime.wan22_i2v` | latent replacement only | blocked | scaffold skip | Q001, Q004, Q005 | +| Scheduler | scheduler | undecided | unknown | unknown | undecided | blocked | blocked | scaffold skip | Q001 | + +## Conversion State +- conversion_script: none +- converted_weights_dir: none +- source_layout: unknown +- strict_load_status: not started +- passthrough_components: unknown +- retry_history: none + +## Parity Commands +| Scope | Command | Last Result | Notes | +|---|---|---|---| +| Scaffold | `pytest tests/local_tests/wan2_2_ti2v_5b/test_wan2_2_ti2v_5b_parity_scaffold.py -v -s` | skip expected | Requires Q001/Q005; not parity evidence. | +| Transformer tiny | `pytest fastvideo/tests/mlx/test_mlx_wan22_parity.py -v -s` | local MLX gate | Per-token timesteps only; not real-weight evidence. | +| Transformer real | `pytest fastvideo/tests/mlx/test_mlx_wan22_real_weights.py -v -s` | precise skip expected | Requires staged weights plus `FASTVIDEO_WAN22_5B_REVISION` and `FASTVIDEO_WAN22_5B_SHA256`. | +| I2V latent-only | `pytest fastvideo/tests/mlx/test_mlx_wan22_i2v.py -v -s` | local MLX gate | Frame replacement/timestep only; no VAE, image encoder, mask, or public I2V claim. | +| VAE | Recorded after official selection | not started | 48-channel latent normalization and decode. | +| Pipeline | Recorded after components pass | not started | Compare denoised latents and decoded media. | + +## Open Questions +| ID | Question | Owner | Needed By Phase | Status | Resolution | +|---|---|---|---|---|---| +| Q001 | Which official source, immutable revision, model-card license, and SHA256 manifest match the intended 5B release? | product/model owner | 0 | open | None; historical `FastWan2.2-TI2V-5B-FullAttn` names are not sufficient. Do not download or activate real tests before selection. | +| Q002 | What is the official per-token timestep input shape, dtype, and transformer call path? | reference-study owner | 1 | open | None. | +| Q003 | What are the 48-channel VAE latent mean/std and normalization/decode semantics? | reference-study owner | 1 | open | None. | +| Q004 | What image conditioner, image-latent layout, and I2V mask semantics does the official pipeline use? | reference-study owner | 1 | open | None. | +| Q005 | Is the official checkpoint Diffusers-compatible or is a conversion script required? | conversion owner | 0 | open | None. | + +## Issues And Blockers +| ID | Phase | Component | Severity | Issue | Evidence | Owner | Status | Resolution | +|---|---|---|---|---|---|---|---|---| +| I001 | 0 | all | blocker | Official reference and matching weights are not selected. | Q001 | product/model owner | open | None. | +| I002 | 1 | transformer/VAE | blocker | No native classes or converted weights exist. | Component matrix | port owner | open | Expected until Q001–Q005 close. | +| I003 | future dependency | self-forcing | blocker | PRs #1307/#1042/#814 lack this port's required causal/KV-cache/training/media gates. | Future dependency policy | self-forcing owner | open | Keep separate. | +| I004 | future candidate | FullAttn config | medium | PR #1563 is OPEN and BLOCKED; it changes Wan config/registry/runtime surfaces. | `gh pr view 1563` on 2026-07-09 | 5B owner | open | Rebase and test independently before selection. | +| I005 | 0 | artifact contract | blocker | Real-weight parity, QAD arming, VAE decode, and benchmarks lack an approved source revision/checksum manifest. | Q001; tests intentionally skip without explicit environment pins | product/model owner | open | Record source URL/revision, transformer/VAE/decoder SHA256 values, license, and staging path. | + +## Escape Hatches +| ID | Phase | Decision Type | Question | Recommended Option | Status | Resolution | +|---|---|---|---|---|---|---| +| E001 | 0 | model source and cost | Choose official source/weights and authorize any gated download before implementation. | Select the weights whose architecture matches the desired 5B T2V/I2V release, then run preparation. | open | None. | + +## Decisions +| Date | Decision | Rationale | Impact | +|---|---|---|---| +| 2026-07-09 | Keep this branch preparation-only. | Initial release scope is 1.3B T2V; no 5B/I2V public commitment is authorized. | No model code, weights, training, or API on this branch. | +| 2026-07-09 | Do not merge PR #1563 yet. | It is blocked and needs a fresh focused 5B config test. | Candidate remains external. | +| 2026-07-09 | Preserve fork PRs #4/#6/#7/#8/#9/#10/#11 as internal future code and tests, not release surfaces. | Avoid losing reviewed fork work while preventing unsupported launch claims. | No 5B registry, CLI, public I2V, automatic downloads, or model publication. | +| 2026-07-09 | Require strict state-dict loading in every parity/arming path. | Missing keys invalidate a parity claim. | Tests fail rather than printing missing keys or using `strict=False`. | + +## Handoff Notes +- First authorized next step: resolve E001/I005 with one immutable official source and checksum manifest, then run the add-model preparation workflow before activating real component parity tests. +- Required completion gates: real official imports in the FastVideo environment, real weights, native component prototypes, strict loading/conversion, non-skip component parity, pipeline parity, and approved 32 GB+ Mac quality/benchmark evidence. diff --git a/tests/local_tests/wan2_2_ti2v_5b/README.md b/tests/local_tests/wan2_2_ti2v_5b/README.md new file mode 100644 index 0000000000..5036837b19 --- /dev/null +++ b/tests/local_tests/wan2_2_ti2v_5b/README.md @@ -0,0 +1,78 @@ +# Wan2.2-TI2V-5B Port Preparation + +This directory is the required preparation handoff for a future Apple-Silicon +Wan2.2-TI2V-5B port. It contains internal MLX prototypes and local tests +preserved from the fork's Track C/D PR stack; it is not a released model, +training recipe, checkpoint converter, or public image-to-video API. + +## Preparation contract + +| Field | Current value | +|---|---| +| Model family | `wan2_2_ti2v_5b` | +| Planned workloads | T2V and I2V after separate approval | +| Official reference | Candidate artifact names exist in historical PRs, but no official source/revision/checksum is pinned; see Q001 | +| Official checkout | Not cloned | +| HF weights | Not selected or downloaded | +| Source layout | Unknown | +| Conversion required | Unknown | +| Official environment | Blocked on Q001, not an installation failure | +| FastVideo implementation | Internal MLX DiT/latent helpers only; no registry, pipeline, CLI, or public I2V surface | + +Do not download weights, clone an unofficial reference, or create a production +loader until Q001 is resolved. If the chosen weights are gated, use one of +`HF_TOKEN`, `HUGGINGFACE_HUB_TOKEN`, or `HF_API_KEY`; never record its value. + +## Preserved fork PR map + +| Fork PR | Preserved work on this branch | Current gate | +|---|---|---| +| #4 | Causal MLX attention/DiT/sampler, compile regression, CUDA reference and QAD smoke fixtures | Internal self-forcing track; no public streaming example; needs pinned SFWan artifact, parity, and media-quality review | +| #6 | Hardware-tier recommendation and backend-agnostic tests | 5B auto-selection remains disabled until Q001 is pinned | +| #7 | Long-rollout cache injection and bounded-KV regressions | Real-weight stream skips until reviewed SFWan artifacts are staged | +| #8 | Per-token-timestep `MLXWan22DiT` and tiny parity test | Tiny implementation parity only; not a 5B release claim | +| #9 | Local-only real-weight parity scaffold, CUDA reference, and 5B benchmark | Requires exact source revision plus transformer SHA256 before the test activates | +| #10 | Latent-only I2V preparation and strict QAD arming scaffold | Image/VAE/mask parity is still blocked; arming gate never downloads artifacts | +| #11 | 48-channel TAEHV/Wan-VAE helpers, DMD sampler, and local/reference tests | Decoder/VAE artifacts require explicit reviewed paths; no public 5B generation example retained | + +Historical benchmark values and example launch commands from those PRs are not +release evidence. Reproduce them only after the official artifact contract is +recorded below and all non-skip gates pass. + +## Required parity activation + +The placeholder test is deliberately a specific dependency skip: it activates +only after the official import path, revision, and weight layout are recorded +in `PORT_STATUS.md`. Its completed replacement must load real official weights +and a FastVideo-native component, use deterministic inputs, and compare actual +tensors. Shape-only checks, unconditional skips, and external upstream virtual +environments are not acceptable evidence. + +| Gate | Planned evidence | Status | +|---|---|---| +| Transformer | Per-token timestep conditioning and text/I2V conditioning outputs | tiny local gate only; real gate blocked on Q001/Q002 | +| VAE | 48-channel latent shape and mean/std normalization encode/decode outputs | blocked on Q001/Q003 | +| Pipeline | T2V denoised-latent parity and I2V mask/image conditioning | I2V intentionally latent-only until component parity | +| Conversion | Official to FastVideo state-dict strict load and pre-quantized MLX artifact round trip | blocked on source layout | +| Quality | Fixed prompt/image packet on a 32 GB+ Mac, valid MP4s, memory/latency, human review | blocked on implementation | + +## Future-only operating gates + +- A 5B implementation, model download, public I2V surface, GPU training run, + model upload, and quality-reference upload each require explicit approval. +- The 32 GB+ Mac benchmark must report exact machine, macOS, MLX, PyTorch, + peak memory, cold/steady latency, and generated-media validation. +- Self-forcing PRs #1307, #1042, and #814 remain independent blocked + dependencies. They need their own causal runtime, KV-cache, training, and + media-quality evidence before they can be considered here. +- PR #1563 is an open, blocked `FastWan FullAttn` configuration candidate. It + has not been rebased or merged into this branch; reconsider it only with a + fresh rebase and focused 5B configuration test. +- PRs #1557, #1496, #1488, #1494, and #1344 are explicitly out of scope. + +## Commands once preparation is complete + +```bash +pytest tests/local_tests/wan2_2_ti2v_5b/test_wan2_2_ti2v_5b_parity_scaffold.py -v -s +pytest tests/local_tests -k 'wan2_2_ti2v_5b and parity' -v -s +``` diff --git a/tests/local_tests/wan2_2_ti2v_5b/test_wan2_2_ti2v_5b_parity_scaffold.py b/tests/local_tests/wan2_2_ti2v_5b/test_wan2_2_ti2v_5b_parity_scaffold.py new file mode 100644 index 0000000000..498c9601f5 --- /dev/null +++ b/tests/local_tests/wan2_2_ti2v_5b/test_wan2_2_ti2v_5b_parity_scaffold.py @@ -0,0 +1,29 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Preparation-only parity scaffold for the future Wan2.2-TI2V-5B port. + +Coverage scope: both. This becomes a real official-versus-FastVideo component +and pipeline parity test only after PORT_STATUS Q001–Q005 identify the source, +weights, component APIs, and conversion route. Until then its dependency skip +prevents accidental claims that the unimplemented port has parity evidence. +""" + +from __future__ import annotations + +import os +from pathlib import Path + +import pytest + + +def test_wan2_2_ti2v_5b_parity_requires_preparation_handoff() -> None: + """Make missing official inputs a precise skip rather than false evidence.""" + official_ref = os.environ.get("WAN2_2_TI2V_5B_OFFICIAL_REF_DIR") + weights = os.environ.get("WAN2_2_TI2V_5B_WEIGHTS") + if not official_ref or not weights: + pytest.skip("PORT_STATUS Q001/Q005: official reference and weights are not selected") + + missing = [path for path in (official_ref, weights) if not Path(path).exists()] + if missing: + pytest.skip(f"official parity inputs are unavailable: {', '.join(missing)}") + + pytest.skip("PORT_STATUS Q002–Q004: official component call paths are not recorded")