[feat] Add FastMetal 14B and 5B MLX serving, harden test coverage

Author: Ishxn20

examples/serving/mlx_wan21_14b.yaml

@@ -0,0 +1,23 @@
+# Native MLX FastMetal 14B server. Run from the repository root after
+# downloading FastVideo/FastMetal-14B-QAD (it ships a pre-packed MLX DiT,
+# so model_root and mlx_checkpoint are the same directory).
+#   python -m fastvideo.entrypoints.openai.mlx_wan_server --config examples/serving/mlx_wan21_14b.yaml
+runtime: mlx
+generator:
+  model_path: FastVideo/FastMetal-14B-QAD
+  model_root: ./FastMetal-14B-QAD
+  mlx_checkpoint: ./FastMetal-14B-QAD
+server:
+  host: 127.0.0.1
+  port: 8000
+  served_model_name: fastwan21-14b-mlx
+  output_dir: outputs/mlx_wan21_14b
+default_request:
+  sampling:
+    height: 480
+    width: 832
+    num_frames: 81
+    fps: 16
+    num_inference_steps: 3
+    guidance_scale: 1.0
+    seed: 1024

examples/serving/mlx_wan22_5b.yaml

@@ -0,0 +1,23 @@
+# Native MLX FastMetal 5B (Wan2.2-TI2V) server. Run from the repository root
+# after downloading FastVideo/FastMetal-5B-QAD (it ships a pre-packed MLX DiT,
+# so model_root and mlx_checkpoint are the same directory).
+#   python -m fastvideo.entrypoints.openai.mlx_wan_server --config examples/serving/mlx_wan22_5b.yaml
+runtime: mlx
+generator:
+  model_path: FastVideo/FastMetal-5B-QAD
+  model_root: ./FastMetal-5B-QAD
+  mlx_checkpoint: ./FastMetal-5B-QAD
+server:
+  host: 127.0.0.1
+  port: 8000
+  served_model_name: fastwan22-5b-mlx
+  output_dir: outputs/mlx_wan22_5b
+default_request:
+  sampling:
+    height: 704
+    width: 1280
+    num_frames: 81
+    fps: 24
+    num_inference_steps: 3
+    guidance_scale: 1.0
+    seed: 1234

fastvideo/entrypoints/openai/mlx_wan_server.py

@@ -21,16 +21,24 @@
 from fastvideo.entrypoints.openai.api_server import create_app
 from fastvideo.entrypoints.openai.protocol import VideoGenerationRequest
 
-MODEL: Literal["FastVideo/FastMetal-1.3B-QAD"] = "FastVideo/FastMetal-1.3B-QAD"
-# The DMD-distilled step ladder the validated 1.3B recipe uses (fixed count,
-# same reason H3 MLX serving pins num_inference_steps to its own ladder size).
+# The DMD-distilled step ladder the validated recipes use (fixed count, same
+# reason H3 MLX serving pins num_inference_steps to its own ladder size).
 _DMD_STEP_COUNT = 3
 
+# Maps a served model id to the pipeline class that generates it: 1.3B/14B
+# share Wan2.1's architecture (MLXWanPipeline), 5B is Wan2.2-TI2V instead
+# (MLXWan22Pipeline, 48-channel latents, a different DiT and sampler).
+_PIPELINE_CLASS_NAMES = {
+    "FastVideo/FastMetal-1.3B-QAD": "MLXWanPipeline",
+    "FastVideo/FastMetal-14B-QAD": "MLXWanPipeline",
+    "FastVideo/FastMetal-5B-QAD": "MLXWan22Pipeline",
+}
+
 
 class MLXWanGeneratorConfig(BaseModel):
     """Where the two FastMetal checkpoint halves live on disk."""
     model_config = ConfigDict(extra="forbid")
-    model_path: Literal["FastVideo/FastMetal-1.3B-QAD"] = MODEL
+    model_path: Literal["FastVideo/FastMetal-1.3B-QAD", "FastVideo/FastMetal-14B-QAD", "FastVideo/FastMetal-5B-QAD"]
     model_root: str
     mlx_checkpoint: str
 
@@ -107,9 +115,10 @@ def _load(config: MLXWanGeneratorConfig):
             raise RuntimeError("Wan MLX serving requires an Apple Silicon Mac.")
         if shutil.which("ffmpeg") is None:
             raise RuntimeError("Install ffmpeg before starting the Wan MLX server.")
-        from fastvideo.mlx_runtime.wan_pipeline import MLXWanPipeline
+        import fastvideo.mlx_runtime.wan_pipeline as wan_pipeline_module
 
-        return MLXWanPipeline(
+        pipeline_cls = getattr(wan_pipeline_module, _PIPELINE_CLASS_NAMES[config.model_path])
+        return pipeline_cls(
             model_root=Path(config.model_root).expanduser(),
             mlx_checkpoint=Path(config.mlx_checkpoint).expanduser(),
         )
@@ -166,7 +175,7 @@ def create_mlx_wan_app(config: MLXWanServeConfig):
         raise ValueError("Wan MLX default_request must set: " + ", ".join(sorted(required - set(explicit))))
     validate_wan_video_request(VideoGenerationRequest(prompt="validate config", **explicit))
     # Transport admission uses the registered Wan family, not CUDA engine options.
-    args = SimpleNamespace(model_path=MODEL,
+    args = SimpleNamespace(model_path=config.generator.model_path,
                            lora_path=None,
                            lora_nickname="default",
                            lora_strength=1.0,

fastvideo/mlx_runtime/wan_pipeline.py

@@ -1,20 +1,26 @@
 # SPDX-License-Identifier: Apache-2.0
-"""Text-to-video generation for Wan2.1/FastMetal through the native MLX runtime.
-
-Scoped to the validated cookbook path only: text-to-video, DMD-distilled
-denoising, a packed MLX DiT checkpoint, and TAEHV decode. Refine, fast-spatial,
-RIFE fast mode, and prompt enrichment stay in the CLI script
-(examples/inference/basic/mlx_wan_prompt_to_video.py) -- this module holds
-only what a resident server needs to call repeatedly.
-
-Every step below reuses the same helpers the CLI script already runs (prompt
-encoding, checkpoint loading, DMD scheduling, VAE decode) so this pipeline and
-the script cannot silently drift into two different implementations of the
-same math.
+"""Text-to-video generation for Wan2.1 and Wan2.2-TI2V (FastMetal) through the
+native MLX runtime.
+
+Scoped to each family's validated cookbook path only: text-to-video,
+DMD-distilled denoising, a packed MLX DiT checkpoint, and TAEHV decode.
+Refine, fast-spatial, RIFE fast mode, and prompt enrichment stay in the CLI
+scripts (examples/inference/basic/mlx_wan_prompt_to_video.py and
+mlx_wan22_generate.py) -- this module holds only what a resident server needs
+to call repeatedly.
+
+Every step below reuses the same helpers the CLI scripts already run (prompt
+encoding, checkpoint loading, DMD scheduling, VAE decode) so these pipelines
+and the scripts cannot silently drift into different implementations of the
+same math. MLXWanPipeline (Wan2.1: 1.3B/14B) and MLXWan22Pipeline (Wan2.2:
+5B) share the UMT5 prompt encoder and rotary-embedding builder below, since
+that piece is identical across both families; everything DiT/VAE-shaped is
+not, because the two are genuinely different architectures.
 """
 from __future__ import annotations
 
 from dataclasses import dataclass, field
+import json
 from pathlib import Path
 import time
 from typing import Any
@@ -31,11 +37,16 @@
 
 logger = init_logger(__name__)
 
-# Wan2.1's VAE compresses 4x temporally and 8x spatially; this is fixed for
-# every Wan2.1/FastMetal checkpoint this pipeline supports (see the guard in
-# MLXWanPipeline.__init__ that rejects Wan2.2-TI2V's 48-channel checkpoints).
+# Wan2.1's VAE compresses 4x temporally and 8x spatially; Wan2.2-TI2V's
+# compresses 4x temporally and 16x spatially. The two families are not
+# interchangeable -- MLXWanPipeline/MLXWan22Pipeline each guard against being
+# pointed at the other's checkpoint (see _packed_dit_channels below).
 _WAN21_TEMPORAL_COMPRESSION = 4
 _WAN21_SPATIAL_COMPRESSION = 8
+_WAN21_CHANNELS = 16
+_WAN22_TEMPORAL_COMPRESSION = 4
+_WAN22_SPATIAL_COMPRESSION = 16
+_WAN22_CHANNELS = 48
 _DEFAULT_DMD_STEPS = (1000, 757, 522)
 
 
@@ -133,6 +144,24 @@ def _make_wan_rotary_embeddings(config: dict[str, Any], *, latent_frames: int, l
     return mx.array(freqs_cos.numpy()).astype(mx.float32), mx.array(freqs_sin.numpy()).astype(mx.float32)
 
 
+def _packed_dit_channels(mlx_checkpoint: Path) -> int | None:
+    """Read in_channels from a packed mlx_dit.json, or None if unreadable.
+
+    A best-effort check: an unpacked/diffusers-style or missing checkpoint is
+    left for generate() to fail on when it actually loads the weights.
+    """
+    manifest_path = mlx_checkpoint / "mlx_dit.json"
+    if not manifest_path.is_file():
+        return None
+    try:
+        manifest = json.loads(manifest_path.read_text())
+    except (json.JSONDecodeError, OSError):
+        return None
+    config = manifest.get("config", manifest)
+    channels = config.get("in_channels")
+    return int(channels) if channels is not None else None
+
+
 class MLXWanPipeline:
     """Text-to-video generation through the native MLX runtime (Wan2.1/FastMetal)."""
 
@@ -145,6 +174,11 @@ def __init__(self, *, model_root: str | Path, mlx_checkpoint: str | Path) -> Non
             raise ValueError(str(error)) from error
         if not (self.model_root / "tokenizer").exists() or not (self.model_root / "text_encoder").exists():
             raise FileNotFoundError(f"Missing tokenizer/ or text_encoder/ under {self.model_root}.")
+        channels = _packed_dit_channels(self.mlx_checkpoint)
+        if channels == _WAN22_CHANNELS:
+            raise ValueError(f"{self.mlx_checkpoint} is a {channels}-channel Wan2.2-TI2V checkpoint "
+                             "(e.g. FastMetal-5B-QAD); MLXWanPipeline only supports Wan2.1's "
+                             f"{_WAN21_CHANNELS}-channel checkpoints (1.3B/14B). Use MLXWan22Pipeline instead.")
 
     def generate(
         self,
@@ -263,3 +297,132 @@ def generate(
                                     k: v
                                     for k, v in timings.items() if k.endswith("_gib")
                                 })
+
+
+class MLXWan22Pipeline:
+    """Text-to-video generation through the native MLX runtime (Wan2.2-TI2V/FastMetal-5B)."""
+
+    def __init__(self, *, model_root: str | Path, mlx_checkpoint: str | Path) -> None:
+        self.model_root = Path(model_root)
+        self.mlx_checkpoint = Path(mlx_checkpoint)
+        try:
+            raise_if_unsupported_mlx_checkpoint(self.mlx_checkpoint)
+        except UnsupportedMLXCheckpointError as error:
+            raise ValueError(str(error)) from error
+        if not (self.model_root / "tokenizer").exists() or not (self.model_root / "text_encoder").exists():
+            raise FileNotFoundError(f"Missing tokenizer/ or text_encoder/ under {self.model_root}.")
+        channels = _packed_dit_channels(self.mlx_checkpoint)
+        if channels is not None and channels != _WAN22_CHANNELS:
+            raise ValueError(f"{self.mlx_checkpoint} is a {channels}-channel checkpoint; MLXWan22Pipeline only "
+                             f"supports Wan2.2-TI2V's {_WAN22_CHANNELS}-channel checkpoints (FastMetal-5B-QAD). "
+                             "Use MLXWanPipeline for 1.3B/14B.")
+
+    def generate(
+        self,
+        prompt: str,
+        *,
+        output_path: str | Path,
+        # Defaults match the validated FastMetal-5B-QAD cookbook recipe, not
+        # mlx_wan22_generate.py's own argparse defaults (448x832x121), which
+        # were never the evidence-backed shape for this checkpoint.
+        height: int = 704,
+        width: int = 1280,
+        num_frames: int = 81,
+        seed: int = 1234,
+        dmd_denoising_steps: tuple[int, ...] = _DEFAULT_DMD_STEPS,
+        flow_shift: float = 5.0,
+        fps: int = 24,
+        max_sequence_length: int = 512,
+    ) -> GenerationResult:
+        import mlx.core as mx
+        import torch
+
+        from fastvideo.mlx_runtime.wan22 import mlx_wan22_dit_from_mlx_checkpoint
+        from fastvideo.mlx_runtime.wan22_sample import sample_wan22_dmd
+        from fastvideo.mlx_runtime.wan_vae import decode_latents_to_video
+
+        timings: dict[str, float] = {}
+        mx.random.seed(seed)
+
+        started = time.perf_counter()
+        prompt_embeds = _encode_wan_prompt(model_root=self.model_root,
+                                           prompt=prompt,
+                                           max_sequence_length=max_sequence_length)
+        timings["encode_s"] = time.perf_counter() - started
+
+        plan = plan_refine_resolutions(
+            height=height,
+            width=width,
+            num_frames=num_frames,
+            vae_spatial_compression=_WAN22_SPATIAL_COMPRESSION,
+            vae_temporal_compression=_WAN22_TEMPORAL_COMPRESSION,
+            enabled=False,
+        )
+
+        started = time.perf_counter()
+        mx.clear_cache()
+        mx.reset_peak_memory()
+        dit = mlx_wan22_dit_from_mlx_checkpoint(self.mlx_checkpoint, compile=True)
+        timings["load_s"] = time.perf_counter() - started
+        timings["load_peak_gib"] = _peak_memory_gib()
+
+        latents_seed = torch.Generator(device="cpu").manual_seed(seed)
+        latents_torch = torch.randn(
+            (1, int(
+                dit.config["in_channels"]), plan.latent_frames, plan.stage1_latent_height, plan.stage1_latent_width),
+            generator=latents_seed,
+            dtype=torch.float32,
+        )
+        noise = mx.array(latents_torch.numpy()).astype(mx.float16)
+        encoder_hidden_states = mx.array(prompt_embeds.numpy()).astype(mx.float16)
+        freqs_cis = _make_wan_rotary_embeddings(
+            dit.config,
+            latent_frames=plan.latent_frames,
+            latent_height=plan.stage1_latent_height,
+            latent_width=plan.stage1_latent_width,
+        )
+
+        started = time.perf_counter()
+        mx.reset_peak_memory()
+        # sample_wan22_dmd's own re-noise seed defaults to 0 in the CLI script
+        # (--renoise-seed), independent of --seed; matched here rather than
+        # exposed as a second knob nobody overrides in the validated recipe.
+        latents = sample_wan22_dmd(
+            dit,
+            encoder_hidden_states,
+            noise,
+            freqs_cis,
+            dmd_denoising_steps=list(dmd_denoising_steps),
+            flow_shift=flow_shift,
+            warp_denoising_step=True,
+            seed=0,
+        )
+        timings["denoise_s"] = time.perf_counter() - started
+        timings["denoise_peak_gib"] = _peak_memory_gib()
+
+        latents_np = np.array(latents.astype(mx.float32))
+        # Free the DiT before decode, matching the CLI script's phase-memory
+        # policy -- the 5B DiT and the decoder are not held resident together.
+        del dit, latents, encoder_hidden_states, freqs_cis, noise
+        cleanup_mlx()
+
+        started = time.perf_counter()
+        output_path = Path(output_path)
+        decode_latents_to_video(
+            latents_np,
+            output_path,
+            fps=fps,
+            backend="taehv",
+            z_dim=latents_np.shape[1],
+            taehv_checkpoint=None,
+            torch_device="auto",
+        )
+        timings["decode_s"] = time.perf_counter() - started
+        cleanup_torch_mps()
+
+        return GenerationResult(video_path=str(output_path),
+                                timings=timings,
+                                peak_memory_gib={
+                                    k: v
+                                    for k, v in timings.items() if k.endswith("_gib")
+                                })