diff --git a/.github/workflows/ci-macos-mlx.yml b/.github/workflows/ci-macos-mlx.yml index 62054c8393..39e5db64e2 100644 --- a/.github/workflows/ci-macos-mlx.yml +++ b/.github/workflows/ci-macos-mlx.yml @@ -94,6 +94,7 @@ jobs: fastvideo/tests/mlx/test_mlx_refine.py \ fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \ fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \ + fastvideo/tests/mlx/test_mlx_wan_pipeline.py \ fastvideo/tests/mlx/test_wan22_sample.py \ fastvideo/tests/mlx/test_windowed_attention.py \ fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_download_unavailable_has_specific_error \ @@ -158,6 +159,7 @@ jobs: fastvideo/tests/mlx/test_mlx_refine.py \ fastvideo/tests/mlx/test_mlx_prompt_to_video_decode.py \ fastvideo/tests/mlx/test_mlx_wan22_prompt_cache_fingerprint.py \ + fastvideo/tests/mlx/test_mlx_wan_pipeline.py \ fastvideo/tests/mlx/test_wan22_sample.py \ fastvideo/tests/mlx/test_windowed_attention.py \ fastvideo/tests/mlx/test_mlx_rife_interpolation.py::test_rife_download_unavailable_has_specific_error \ diff --git a/examples/serving/mlx_wan21_14b.yaml b/examples/serving/mlx_wan21_14b.yaml new file mode 100644 index 0000000000..959ddbca91 --- /dev/null +++ b/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 diff --git a/examples/serving/mlx_wan21_1_3b.yaml b/examples/serving/mlx_wan21_1_3b.yaml new file mode 100644 index 0000000000..dc09bcd28e --- /dev/null +++ b/examples/serving/mlx_wan21_1_3b.yaml @@ -0,0 +1,23 @@ +# Native MLX FastMetal 1.3B server. Run from the repository root after +# downloading FastVideo/FastMetal-1.3B-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_1_3b.yaml +runtime: mlx +generator: + model_path: FastVideo/FastMetal-1.3B-QAD + model_root: ./FastMetal-1.3B-QAD + mlx_checkpoint: ./FastMetal-1.3B-QAD +server: + host: 127.0.0.1 + port: 8000 + served_model_name: fastwan21-1.3b-mlx + output_dir: outputs/mlx_wan21_1_3b +default_request: + sampling: + height: 480 + width: 832 + num_frames: 81 + fps: 16 + num_inference_steps: 3 + guidance_scale: 1.0 + seed: 1024 diff --git a/examples/serving/mlx_wan22_5b.yaml b/examples/serving/mlx_wan22_5b.yaml new file mode 100644 index 0000000000..82934403ab --- /dev/null +++ b/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 diff --git a/fastvideo/entrypoints/openai/api_server.py b/fastvideo/entrypoints/openai/api_server.py index 81d669e5b2..6fa909d081 100644 --- a/fastvideo/entrypoints/openai/api_server.py +++ b/fastvideo/entrypoints/openai/api_server.py @@ -62,6 +62,8 @@ async def lifespan(app: FastAPI) -> AsyncIterator[None]: default_request: GenerationRequest | None = getattr(app.state, "default_request", None) logger.info("Initializing %s generation runtime for %s ...", app.state.runtime, args.model_path) + # A non-CUDA runtime (e.g. MLX) supplies its own generator_factory instead + # of the default VideoGenerator/Executor path -- see MLXWanGenerator. factory = app.state.generator_factory generator = factory() if factory is not None else VideoGenerator.from_fastvideo_args(args) serving_engine = OpenAIServingEngine(generator, app.state.video_request_validator) @@ -98,7 +100,13 @@ def create_app( video_request_validator: Callable[[VideoGenerationRequest], None] | None = None, runtime: str = "cuda", ) -> FastAPI: - """Build the FastAPI application with all routers mounted""" + """Build the FastAPI application with all routers mounted. + + ``generator_factory``/``video_request_validator``/``runtime`` let a + non-CUDA backend (MLX) plug in its own generator and request rules + without a separate app-building path. Defaults reproduce the original + CUDA-only behavior exactly. + """ app = FastAPI( title="FastVideo OpenAI-Compatible API", @@ -158,6 +166,8 @@ async def openai_validation_error(_request: Request, exc: RequestValidationError app.include_router(common_router) app.include_router(video_router) + # The MLX runtime only wires video-with-audio generation (see + # fastvideo/mlx_runtime/); image generation has no MLX backend yet. if runtime != "mlx": app.include_router(image_router) app.include_router(playground_router) diff --git a/fastvideo/entrypoints/openai/mlx_wan_server.py b/fastvideo/entrypoints/openai/mlx_wan_server.py new file mode 100644 index 0000000000..1119683cc3 --- /dev/null +++ b/fastvideo/entrypoints/openai/mlx_wan_server.py @@ -0,0 +1,269 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Serve native FastMetal (Wan2.1) MLX through the shared video-job API and playground.""" + +from __future__ import annotations + +import argparse +from concurrent.futures import ThreadPoolExecutor +from functools import partial +from pathlib import Path +import platform +import shutil +import time +from types import SimpleNamespace +from typing import Any, Literal + +from pydantic import BaseModel, ConfigDict, Field +import uvicorn +import yaml + +from fastvideo.api.compat import explicit_request_updates, normalize_generation_request +from fastvideo.api.schema import GenerationRequest +from fastvideo.entrypoints.openai.api_server import create_app +from fastvideo.entrypoints.openai.protocol import VideoGenerationRequest + +# 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 + +# Wan's VAE compresses this many input frames into one latent frame, so a legal +# frame count is 1 modulo the stride; plan_refine_resolutions enforces it. +_VAE_TEMPORAL_COMPRESSION = 4 + +# The adapter's own fps fallback, used when a caller validates a request outside +# a configured server (create_mlx_wan_app binds the served fps instead). +_FALLBACK_FPS = 24 + +# 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", "FastVideo/FastMetal-14B-QAD", "FastVideo/FastMetal-5B-QAD"] + model_root: str + mlx_checkpoint: str + + +class MLXWanServerConfig(BaseModel): + """Host/port/output shape for an MLX serve YAML's ``server:`` block. + + Generic across MLX-served models -- if a second MLX server config needs + the same shape, promote this (and its H3 counterpart) to one shared + module instead of a third copy. + """ + model_config = ConfigDict(extra="forbid") + host: str = "127.0.0.1" + port: int = Field(default=8000, ge=1, le=65535) + output_dir: str = "outputs/mlx_wan" + served_model_name: str = Field(default="fastwan", min_length=1) + + +class MLXWanServeConfig(BaseModel): + """Top-level ``mlx_wan_*.yaml`` shape read by --config.""" + model_config = ConfigDict(extra="forbid") + runtime: Literal["mlx"] + generator: MLXWanGeneratorConfig + server: MLXWanServerConfig = Field(default_factory=MLXWanServerConfig) + default_request: dict[str, Any] + + +def _aligned_num_frames(num_frames: int) -> int: + """Round up to the next frame count Wan accepts (1 modulo the VAE temporal stride).""" + remainder = (num_frames - 1) % _VAE_TEMPORAL_COMPRESSION + if remainder == 0: + return num_frames + return num_frames + (_VAE_TEMPORAL_COMPRESSION - remainder) + + +def _align_seconds_to_frame_grid(request: VideoGenerationRequest, *, default_fps: int) -> None: + """Resolve an explicit ``seconds`` into a Wan-legal ``num_frames`` before admission. + + ``build_generation_request`` turns ``seconds`` into ``seconds * fps``, which lands on + 0 modulo the VAE temporal stride at every fps, while ``plan_refine_resolutions`` + requires 1. Left alone the job is admitted and only fails inside generation, so + resolve it synchronously here; the adapter prefers an explicit ``num_frames`` over + ``seconds``, and assigning one records it in ``model_fields_set``. + + Mirrors the adapter's explicit-field precedence, including the nested + ``video_params`` spelling, so a request using either form is aligned. + """ + body_set = request.model_fields_set + nested_set = request.video_params.model_fields_set if request.video_params is not None else set() + if "seconds" not in body_set or request.seconds is None: + return + frames_explicit = ("num_frames" in body_set + and request.num_frames is not None) or ("video_params" in body_set and "num_frames" in nested_set + and request.video_params.num_frames is not None) + if frames_explicit: + return + fps = default_fps + if "fps" in body_set and request.fps is not None: + fps = request.fps + elif "video_params" in body_set and "fps" in nested_set and request.video_params.fps is not None: + fps = request.video_params.fps + request.num_frames = _aligned_num_frames(int(request.seconds) * int(fps)) + + +def validate_wan_video_request(request: VideoGenerationRequest, *, default_fps: int = _FALLBACK_FPS) -> None: + """Reject unsupported inputs before fetching media or creating a job. + + Also normalizes the two request shapes the shared adapter would otherwise + reject after the job is already admitted: an explicit ``task`` and a + ``seconds`` duration that does not land on Wan's frame grid. + """ + allowed = { + "model", + "prompt", + "seed", + "size", + "width", + "height", + "fps", + "num_frames", + "seconds", + "video_params", + "task", + "guidance_scale", + "num_inference_steps", + } + unsupported = request.model_fields_set - allowed + if unsupported: + raise ValueError("Wan MLX serving does not support: " + ", ".join(sorted(unsupported))) + if request.task not in (None, "t2v"): + raise ValueError("Wan MLX serving supports task=t2v only.") + if request.task is not None: + request.task = None + request.model_fields_set.discard("task") + if request.guidance_scale not in (None, 1.0): + raise ValueError("FastMetal MLX is DMD-distilled and requires guidance_scale=1.") + if request.num_inference_steps not in (None, _DMD_STEP_COUNT): + raise ValueError(f"Wan MLX serving uses a fixed {_DMD_STEP_COUNT}-step DMD ladder; " + f"num_inference_steps must be {_DMD_STEP_COUNT}.") + if request.seed is not None and not 0 <= request.seed <= 2**32 - 1: + raise ValueError("Wan MLX seed must be between 0 and 4294967295.") + _align_seconds_to_frame_grid(request, default_fps=default_fps) + + +class MLXWanGenerator: + """Keep one FastMetal pipeline on one MLX thread across requests.""" + + def __init__(self, config: MLXWanGeneratorConfig) -> None: + self._worker = ThreadPoolExecutor(max_workers=1, thread_name_prefix="wan-mlx") + try: + self._pipeline = self._worker.submit(self._load, config).result() + except BaseException: + self._worker.shutdown(wait=True) + raise + + @staticmethod + def _load(config: MLXWanGeneratorConfig): + """Load the pipeline; must run on the MLX worker thread.""" + if platform.system() != "Darwin" or platform.machine() != "arm64": + 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.") + import fastvideo.mlx_runtime.wan_pipeline as wan_pipeline_module + + 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(), + ) + + def generate(self, request: GenerationRequest) -> dict[str, Any]: + """Run one generation on the MLX worker thread; block until it finishes.""" + return self._worker.submit(self._generate, request).result() + + def _generate(self, request: GenerationRequest) -> dict[str, Any]: + """The actual pipeline call; must run on the MLX worker thread.""" + started = time.perf_counter() + result = self._pipeline.generate( + request.prompt, + output_path=request.output.output_path, + width=request.sampling.width, + height=request.sampling.height, + num_frames=request.sampling.num_frames, + seed=request.sampling.seed, + fps=request.sampling.fps, + ) + return {"video_path": str(result.video_path), "generation_time": time.perf_counter() - started} + + def shutdown(self) -> None: + """Release the pipeline and stop the MLX worker thread.""" + + def release(): + self._pipeline = None + from fastvideo.mlx_runtime.memory import cleanup_mlx + + cleanup_mlx() + + try: + self._worker.submit(release).result() + finally: + self._worker.shutdown(wait=True) + + +def load_config(path: str) -> MLXWanServeConfig: + """Parse a Wan MLX serve YAML into its typed config.""" + with open(path, encoding="utf-8") as source: + return MLXWanServeConfig.model_validate(yaml.safe_load(source)) + + +def create_mlx_wan_app(config: MLXWanServeConfig): + """Build the FastAPI app for a validated Wan MLX serve config.""" + request = normalize_generation_request(config.default_request) + explicit = explicit_request_updates(request) + supported = {"width", "height", "num_frames", "fps", "seed", "num_inference_steps", "guidance_scale"} + if set(explicit) - supported: + raise ValueError("Wan MLX default_request contains unsupported fields: " + + ", ".join(sorted(set(explicit) - supported))) + required = {"width", "height", "num_frames", "fps"} + if required - set(explicit): + 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=config.generator.model_path, + lora_path=None, + lora_nickname="default", + lora_strength=1.0, + override_pipeline_cls_name=None) + from fastvideo.entrypoints.openai.request_adapter import build_generation_request + + build_generation_request("config-check", + VideoGenerationRequest(prompt="validate config"), + args, + served_model_name=config.server.served_model_name, + output_dir=config.server.output_dir, + default_request=request) + return create_app( + args, + config.server.output_dir, + request, + config.server.served_model_name, + generator_factory=lambda: MLXWanGenerator(config.generator), + video_request_validator=partial(validate_wan_video_request, default_fps=request.sampling.fps), + runtime="mlx", + ) + + +def main() -> None: + """CLI entrypoint: python -m fastvideo.entrypoints.openai.mlx_wan_server --config ...""" + parser = argparse.ArgumentParser(description=__doc__) + parser.add_argument("--config", + required=True, + help="Wan MLX serving YAML; paths are relative to the working directory") + args = parser.parse_args() + config = load_config(args.config) + uvicorn.run(create_mlx_wan_app(config), host=config.server.host, port=config.server.port) + + +if __name__ == "__main__": + main() diff --git a/fastvideo/entrypoints/openai/serving_engine.py b/fastvideo/entrypoints/openai/serving_engine.py index e01b1c99d0..8d619b6495 100644 --- a/fastvideo/entrypoints/openai/serving_engine.py +++ b/fastvideo/entrypoints/openai/serving_engine.py @@ -14,6 +14,11 @@ class ServingGenerator(Protocol): + """The minimal shape OpenAIServingEngine needs from a generator. + + VideoGenerator (CUDA/mp/ray) satisfies this structurally already; an MLX + generator (e.g. MLXWanGenerator) implements just these two methods. + """ def generate(self, request: GenerationRequest) -> Any: ... @@ -47,6 +52,7 @@ def generator(self) -> ServingGenerator: return self._generator def validate_video_request(self, request: VideoGenerationRequest) -> None: + """Runtime-specific request checks (e.g. MLX's supported-field allowlist).""" if self._video_request_validator is not None: self._video_request_validator(request) @@ -143,4 +149,4 @@ async def shutdown(self) -> None: await asyncio.to_thread(self._generator.shutdown) -__all__ = ["OpenAIServingEngine"] +__all__ = ["OpenAIServingEngine", "ServingGenerator"] diff --git a/fastvideo/entrypoints/openai/state.py b/fastvideo/entrypoints/openai/state.py index 66fd39a456..c990bb6c46 100644 --- a/fastvideo/entrypoints/openai/state.py +++ b/fastvideo/entrypoints/openai/state.py @@ -25,7 +25,8 @@ def get_generator() -> ServingGenerator: - """Return the global VideoGenerator instance (set during startup).""" + """Return the global generator (set during startup). + """ assert _generator is not None, "Server not initialized — generator is None" return _generator diff --git a/fastvideo/entrypoints/openai/video_api.py b/fastvideo/entrypoints/openai/video_api.py index 5c1c7a2350..cdc83e9a8a 100644 --- a/fastvideo/entrypoints/openai/video_api.py +++ b/fastvideo/entrypoints/openai/video_api.py @@ -331,6 +331,8 @@ async def _parse_video_request(raw_request: Request) -> VideoGenerationRequest: async def _adapt_request(request_id: str, request: VideoGenerationRequest) -> GenerationRequest: try: + # Runtime-specific checks (e.g. MLX's supported-field allowlist) run + # before the CUDA-oriented model/LoRA validation below. get_serving_engine().validate_video_request(request) validate_model_and_lora(request, get_server_args(), get_served_model_name()) await prepare_reference_media(request_id, request, get_output_dir()) diff --git a/fastvideo/mlx_runtime/wan_pipeline.py b/fastvideo/mlx_runtime/wan_pipeline.py new file mode 100644 index 0000000000..16c3e5e8f6 --- /dev/null +++ b/fastvideo/mlx_runtime/wan_pipeline.py @@ -0,0 +1,462 @@ +# SPDX-License-Identifier: Apache-2.0 +"""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 + +import numpy as np + +from fastvideo.mlx_runtime.checkpoint_compat import ( + UnsupportedMLXCheckpointError, + raise_if_unsupported_mlx_checkpoint, +) +from fastvideo.mlx_runtime.memory import cleanup_mlx, cleanup_torch_mps +from fastvideo.mlx_runtime.refine import plan_refine_resolutions +from fastvideo.utils import init_logger + +logger = init_logger(__name__) + +# 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) + + +@dataclass +class GenerationResult: + """Everything a caller needs after one generate() call.""" + video_path: str + timings: dict[str, float] = field(default_factory=dict) + peak_memory_gib: dict[str, float] = field(default_factory=dict) + + +def _peak_memory_gib() -> float: + """Read MLX's peak-memory counter in GiB.""" + import mlx.core as mx + + return mx.get_peak_memory() / 2**30 + + +def _resolve_wan_torch_device(device_arg: str): + """Pick the torch device for UMT5 text encoding and TAEHV decode.""" + import torch + + if device_arg == "auto": + return torch.device("mps" if torch.backends.mps.is_available() else "cpu") + return torch.device(device_arg) + + +def _resolve_wan_torch_dtype(dtype_arg: str): + """Map a recipe dtype name to its torch dtype (mirrors the reference scripts).""" + import torch + + dtypes = {"fp16": torch.float16, "bf16": torch.bfloat16, "fp32": torch.float32} + if dtype_arg not in dtypes: + raise ValueError(f"Unsupported text-encoder dtype {dtype_arg!r}; expected one of {sorted(dtypes)}.") + return dtypes[dtype_arg] + + +def _encode_wan_prompt(*, + model_root: Path, + prompt: str, + max_sequence_length: int, + device_arg: str = "auto", + dtype_arg: str = "bf16"): + """Encode a prompt with UMT5, padded/truncated to max_sequence_length. + + The dtype/device pair is a per-family recipe invariant, not a performance knob: + each pipeline passes the values its validated reference script uses. The + defaults are Wan2.1's (bf16 on MPS, matching mlx_wan_prompt_to_video.py); + Wan2.2-TI2V overrides them -- see MLXWan22Pipeline.generate. + """ + import torch + from transformers import AutoTokenizer, UMT5EncoderModel + + dtype = _resolve_wan_torch_dtype(dtype_arg) + device = _resolve_wan_torch_device(device_arg) + tokenizer = AutoTokenizer.from_pretrained(model_root / "tokenizer", local_files_only=True) + text_encoder = UMT5EncoderModel.from_pretrained( + model_root / "text_encoder", + torch_dtype=dtype, + low_cpu_mem_usage=True, + 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", + ) + input_ids = text_inputs.input_ids.to(device) + attention_mask = text_inputs.attention_mask.to(device) + valid_lengths = attention_mask.gt(0).sum(dim=1).long() + + with torch.no_grad(): + hidden_states = text_encoder(input_ids, attention_mask).last_hidden_state + hidden_states = hidden_states.to(dtype=dtype) + trimmed = [row[:length] for row, length in zip(hidden_states, valid_lengths, strict=False)] + padded = torch.stack( + [torch.cat([row, row.new_zeros(max_sequence_length - row.size(0), row.size(1))]) for row in trimmed], + dim=0, + ) + if padded.dtype == torch.bfloat16: + padded = padded.float() + padded = padded.cpu().contiguous() + del text_encoder, tokenizer, text_inputs, input_ids, attention_mask, valid_lengths + cleanup_torch_mps() + return padded + + +def _make_wan_rotary_embeddings(config: dict[str, Any], *, latent_frames: int, latent_height: int, latent_width: int): + """Build the RoPE cos/sin tables the DiT's attention layers expect.""" + 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"]) + 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, + num_heads * head_dim, + 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 _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 + if not isinstance(manifest, dict): + return None + config = manifest.get("config", manifest) + if not isinstance(config, dict): + return None + channels = config.get("in_channels") + if channels is None: + return None + try: + return int(channels) + except (TypeError, ValueError): + return None + + +class MLXWanPipeline: + """Text-to-video generation through the native MLX runtime (Wan2.1/FastMetal).""" + + 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 == _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, + prompt: str, + *, + output_path: str | Path, + height: int = 480, + width: int = 832, + num_frames: int = 81, + seed: int = 0, + dmd_denoising_steps: tuple[int, ...] = _DEFAULT_DMD_STEPS, + flow_shift: float = 8.0, + fps: int = 16, + max_sequence_length: int = 512, + ) -> GenerationResult: + import mlx.core as mx + import torch + + from fastvideo.mlx_runtime.checkpoint import load_mlx_dit_checkpoint + from fastvideo.mlx_runtime.sampling import MLXDMDSchedule, dmd_step + from fastvideo.mlx_runtime.wan_vae import decode_latents_to_video + from fastvideo.models.schedulers.scheduling_flow_match_euler_discrete import FlowMatchEulerDiscreteScheduler + + 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=_WAN21_SPATIAL_COMPRESSION, + vae_temporal_compression=_WAN21_TEMPORAL_COMPRESSION, + enabled=False, + ) + + started = time.perf_counter() + mx.clear_cache() + mx.reset_peak_memory() + dit = load_mlx_dit_checkpoint(self.mlx_checkpoint, compile=True) + timings["load_s"] = time.perf_counter() - started + timings["load_peak_gib"] = _peak_memory_gib() + + scheduler = FlowMatchEulerDiscreteScheduler(shift=flow_shift) + timesteps = torch.tensor(list(dmd_denoising_steps), dtype=torch.long) + dmd_schedule = MLXDMDSchedule.from_torch_scheduler(scheduler) + + 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, + ) + latents = 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() + for step_index, timestep in enumerate(timesteps): + timestep_mx = mx.array([float(timestep.item())]).astype(mx.float32) + noise_pred = dit(latents.astype(mx.float16), encoder_hidden_states, timestep_mx, freqs_cis) + + is_last_step = step_index == len(timesteps) - 1 + next_timestep = None if is_last_step else float(timesteps[step_index + 1].item()) + renoise = None if is_last_step else 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=noise_pred.astype(mx.float32), + schedule=dmd_schedule, + timestep=float(timestep.item()), + next_timestep=next_timestep, + noise=renoise, + ).astype(mx.float16) + mx.eval(latents) + logger.info("Wan MLX denoise step %d/%d complete", step_index + 1, len(timesteps)) + 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 -- this is a resident server, but MLX's + # unified memory still holds one heavyweight phase at a time, matching + # the CLI script's proven memory behavior on this hardware class. + del dit, latents, encoder_hidden_states, freqs_cis + 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") + }) + + +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, + device_arg="cpu", + dtype_arg="fp16") + 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") + }) diff --git a/fastvideo/tests/entrypoints/test_mlx_wan_server.py b/fastvideo/tests/entrypoints/test_mlx_wan_server.py new file mode 100644 index 0000000000..0a904acec9 --- /dev/null +++ b/fastvideo/tests/entrypoints/test_mlx_wan_server.py @@ -0,0 +1,357 @@ +# SPDX-License-Identifier: Apache-2.0 +"""Wan MLX serving config, request validation, and generator dispatch. + +Mirrors fastvideo/tests/entrypoints/test_openai_video_client.py's MLX H3 +coverage pattern. MLXWanGenerator only imports the real `mlx` package inside +_load, which runs on a worker thread and is stubbed out here -- nothing in +this file needs Apple Silicon. +""" +from __future__ import annotations + +from pathlib import Path +from types import SimpleNamespace +from unittest.mock import patch + +import pytest +from pydantic import ValidationError + +from fastvideo.entrypoints.openai.mlx_wan_server import ( + MLXWanGenerator, + create_mlx_wan_app, + load_config, + validate_wan_video_request, +) +from fastvideo.entrypoints.openai.protocol import VideoGenerationRequest + +ROOT = Path(__file__).resolve().parents[3] +CONFIG_PATH = ROOT / "examples/serving/mlx_wan21_1_3b.yaml" +CONFIG_PATH_14B = ROOT / "examples/serving/mlx_wan21_14b.yaml" +CONFIG_PATH_5B = ROOT / "examples/serving/mlx_wan22_5b.yaml" + + +def test_load_config_parses_the_checked_in_yaml() -> None: + config = load_config(str(CONFIG_PATH)) + assert config.runtime == "mlx" + assert config.generator.model_path == "FastVideo/FastMetal-1.3B-QAD" + assert config.generator.model_root == "./FastMetal-1.3B-QAD" + assert config.generator.mlx_checkpoint == "./FastMetal-1.3B-QAD" + assert config.server.served_model_name == "fastwan21-1.3b-mlx" + + +def test_load_config_parses_the_14b_yaml() -> None: + config = load_config(str(CONFIG_PATH_14B)) + assert config.generator.model_path == "FastVideo/FastMetal-14B-QAD" + assert config.generator.model_root == "./FastMetal-14B-QAD" + assert config.server.served_model_name == "fastwan21-14b-mlx" + + +def test_create_app_uses_the_14b_config_model_path_not_the_1_3b_default() -> None: + """Regression guard: the served model must come from the config, not a + module-level constant left over from when only 1.3B was supported.""" + config = load_config(str(CONFIG_PATH_14B)) + app = create_mlx_wan_app(config) + assert app.state.fastvideo_args.model_path == "FastVideo/FastMetal-14B-QAD" + + +def test_load_config_parses_the_5b_yaml() -> None: + config = load_config(str(CONFIG_PATH_5B)) + assert config.generator.model_path == "FastVideo/FastMetal-5B-QAD" + assert config.server.served_model_name == "fastwan22-5b-mlx" + assert config.default_request["sampling"]["height"] == 704 + assert config.default_request["sampling"]["width"] == 1280 + + +def test_generator_config_rejects_unknown_fields() -> None: + from fastvideo.entrypoints.openai.mlx_wan_server import MLXWanGeneratorConfig + with pytest.raises(ValidationError): + MLXWanGeneratorConfig(model_path="FastVideo/FastMetal-1.3B-QAD", model_root="x", mlx_checkpoint="y", + extra_field="not allowed") + + +def test_generator_config_requires_a_model_path() -> None: + from fastvideo.entrypoints.openai.mlx_wan_server import MLXWanGeneratorConfig + with pytest.raises(ValidationError): + MLXWanGeneratorConfig(model_root="x", mlx_checkpoint="y") + + +@pytest.mark.parametrize("model_path", ["FastVideo/FastMetal-1.3B-QAD", "FastVideo/FastMetal-14B-QAD"]) +def test_generator_config_accepts_the_two_wan21_sizes(model_path) -> None: + """1.3B and 14B share MLXWanPipeline's Wan2.1 architecture.""" + from fastvideo.entrypoints.openai.mlx_wan_server import MLXWanGeneratorConfig + config = MLXWanGeneratorConfig(model_path=model_path, model_root="x", mlx_checkpoint="y") + assert config.model_path == model_path + + +def test_generator_config_accepts_the_wan22_5b_model() -> None: + """FastMetal 5B is Wan2.2-TI2V; MLXWanGenerator routes it to MLXWan22Pipeline + (see test_generator_loads_the_pipeline_class_matching_model_path).""" + from fastvideo.entrypoints.openai.mlx_wan_server import MLXWanGeneratorConfig + config = MLXWanGeneratorConfig(model_path="FastVideo/FastMetal-5B-QAD", model_root="x", mlx_checkpoint="y") + assert config.model_path == "FastVideo/FastMetal-5B-QAD" + + +def test_generator_config_rejects_an_unregistered_model() -> None: + from fastvideo.entrypoints.openai.mlx_wan_server import MLXWanGeneratorConfig + with pytest.raises(ValidationError): + MLXWanGeneratorConfig(model_path="FastVideo/SomeOtherModel", model_root="x", mlx_checkpoint="y") + + +@pytest.mark.parametrize("field,value", [ + ("negative_prompt", "bad"), + ("task", "t2va"), + ("guidance_scale", 2.0), + ("num_inference_steps", 10), + ("seed", -1), +]) +def test_validate_rejects_unsupported_request_fields(field, value) -> None: + with pytest.raises(ValueError): + validate_wan_video_request(VideoGenerationRequest(prompt="a cat", **{field: value})) + + +def test_validate_accepts_the_recipe_shape() -> None: + validate_wan_video_request(VideoGenerationRequest(prompt="a cat", num_inference_steps=3, guidance_scale=1.0)) + + +def test_validate_accepts_explicit_task_t2v() -> None: + validate_wan_video_request(VideoGenerationRequest(prompt="a cat", task="t2v")) + + +def test_validate_normalizes_task_t2v_to_none() -> None: + """build_generation_request rejects every non-None task outside MiniMax-H3, so an + accepted t2v must not survive into the adapter -- t2v is all this runtime does.""" + request = VideoGenerationRequest(prompt="a cat", task="t2v") + validate_wan_video_request(request) + assert request.task is None + assert "task" not in request.model_fields_set + + +@pytest.mark.parametrize("seconds,fps,expected", [ + (3, 16, 49), + (2, 16, 33), + (1, 24, 25), + (3, 24, 73), +]) +def test_validate_aligns_seconds_to_the_wan_frame_grid(seconds, fps, expected) -> None: + """seconds * fps lands on 0 modulo 4 at every fps, but Wan requires 1 -- without + alignment the job is admitted and only fails inside plan_refine_resolutions.""" + request = VideoGenerationRequest(prompt="a cat", seconds=seconds, fps=fps) + validate_wan_video_request(request) + assert request.num_frames == expected + assert (request.num_frames - 1) % 4 == 0 + + +def test_validate_aligns_seconds_against_the_served_fps() -> None: + """A request that omits fps must align against the server's configured fps.""" + request = VideoGenerationRequest(prompt="a cat", seconds=3) + validate_wan_video_request(request, default_fps=16) + assert request.num_frames == 49 + + +def test_validate_aligns_seconds_from_the_nested_video_params_spelling() -> None: + """The adapter honours video_params.fps, so alignment must read it too.""" + request = VideoGenerationRequest(prompt="a cat", seconds=3, video_params={"fps": 16}) + validate_wan_video_request(request) + assert request.num_frames == 49 + + +def test_validate_leaves_an_explicit_num_frames_alone() -> None: + """num_frames already wins over seconds in the adapter; don't second-guess it.""" + request = VideoGenerationRequest(prompt="a cat", seconds=3, fps=16, num_frames=81) + validate_wan_video_request(request) + assert request.num_frames == 81 + + +def test_validate_leaves_a_request_without_seconds_untouched() -> None: + request = VideoGenerationRequest(prompt="a cat", fps=16) + validate_wan_video_request(request) + assert request.num_frames is None + assert "num_frames" not in request.model_fields_set + + +def test_create_app_binds_the_served_fps_to_the_request_validator() -> None: + """Regression guard: a seconds-only request must align against the config's fps + (16 for 1.3B), not the adapter's generic 24 fallback.""" + config = load_config(str(CONFIG_PATH)) + app = create_mlx_wan_app(config) + request = VideoGenerationRequest(prompt="a cat", seconds=3) + app.state.video_request_validator(request) + assert request.num_frames == 49 + + +@pytest.mark.parametrize("seed", [0, 2**32 - 1]) +def test_validate_accepts_seed_boundary_values(seed) -> None: + validate_wan_video_request(VideoGenerationRequest(prompt="a cat", seed=seed)) + + +def test_validate_rejects_seed_above_the_uint32_range() -> None: + with pytest.raises(ValueError): + validate_wan_video_request(VideoGenerationRequest(prompt="a cat", seed=2**32)) + + +def test_create_app_builds_without_loading_the_mlx_pipeline() -> None: + """generator_factory is lazy -- building the app must not import mlx.core.""" + config = load_config(str(CONFIG_PATH)) + app = create_mlx_wan_app(config) + assert app.state.served_model_name == "fastwan21-1.3b-mlx" + + +def test_create_app_rejects_unsupported_default_request_fields(tmp_path) -> None: + bad_config_path = tmp_path / "bad.yaml" + bad_config_path.write_text(""" +runtime: mlx +generator: + model_path: FastVideo/FastMetal-1.3B-QAD + model_root: ./FastMetal-1.3B-QAD + mlx_checkpoint: ./FastMetal-1.3B-QAD +server: + served_model_name: bad +default_request: + negative_prompt: not-supported + sampling: + height: 480 + width: 832 + num_frames: 81 + fps: 16 +""") + with pytest.raises(ValueError, match="unsupported fields"): + create_mlx_wan_app(load_config(str(bad_config_path))) + + +def test_create_app_rejects_a_default_request_missing_required_sampling_fields(tmp_path) -> None: + incomplete_config_path = tmp_path / "incomplete.yaml" + incomplete_config_path.write_text(""" +runtime: mlx +generator: + model_path: FastVideo/FastMetal-1.3B-QAD + model_root: ./FastMetal-1.3B-QAD + mlx_checkpoint: ./FastMetal-1.3B-QAD +server: + served_model_name: incomplete +default_request: + sampling: + height: 480 + width: 832 +""") + with pytest.raises(ValueError, match="must set"): + create_mlx_wan_app(load_config(str(incomplete_config_path))) + + +def test_generator_dispatches_request_to_pipeline_on_worker_thread() -> None: + """MLXWanGenerator must call MLXWanPipeline.generate with the request's + sampling fields, entirely off the calling thread.""" + calls: dict[str, object] = {} + + class _StubPipeline: + + def __init__(self, **kwargs): + calls["init_kwargs"] = kwargs + + def generate(self, prompt, **kwargs): + calls["prompt"] = prompt + calls["generate_kwargs"] = kwargs + return SimpleNamespace(video_path="outputs/stub.mp4") + + generator_config = SimpleNamespace(model_path="FastVideo/FastMetal-1.3B-QAD", model_root="./FastMetal-1.3B-QAD", + mlx_checkpoint="./FastMetal-1.3B-QAD") + + with patch("platform.system", return_value="Darwin"), \ + patch("platform.machine", return_value="arm64"), \ + patch("shutil.which", return_value="/usr/bin/ffmpeg"), \ + patch("fastvideo.mlx_runtime.wan_pipeline.MLXWanPipeline", _StubPipeline), \ + patch("fastvideo.mlx_runtime.memory.cleanup_mlx", lambda: calls.setdefault("cleaned_up", True)): + generator = MLXWanGenerator(generator_config) + try: + request = SimpleNamespace( + prompt="a red panda", + output=SimpleNamespace(output_path="outputs/out.mp4"), + sampling=SimpleNamespace(width=832, height=480, num_frames=81, seed=1024, fps=16), + ) + result = generator.generate(request) + finally: + generator.shutdown() + + assert result["video_path"] == "outputs/stub.mp4" + assert calls["prompt"] == "a red panda" + assert calls["generate_kwargs"] == {"output_path": "outputs/out.mp4", "width": 832, "height": 480, + "num_frames": 81, "seed": 1024, "fps": 16} + assert calls["cleaned_up"] is True + + +@pytest.mark.parametrize("model_path,expected_class_name", [ + ("FastVideo/FastMetal-1.3B-QAD", "MLXWanPipeline"), + ("FastVideo/FastMetal-14B-QAD", "MLXWanPipeline"), + ("FastVideo/FastMetal-5B-QAD", "MLXWan22Pipeline"), +]) +def test_generator_loads_the_pipeline_class_matching_model_path(model_path, expected_class_name) -> None: + """1.3B/14B must load MLXWanPipeline (Wan2.1); 5B must load MLXWan22Pipeline + (Wan2.2-TI2V) -- these are different architectures, not interchangeable.""" + loaded: dict[str, object] = {} + + class _StubWan21Pipeline: + + def __init__(self, **kwargs): + loaded["class_name"] = "MLXWanPipeline" + + class _StubWan22Pipeline: + + def __init__(self, **kwargs): + loaded["class_name"] = "MLXWan22Pipeline" + + generator_config = SimpleNamespace(model_path=model_path, model_root="x", mlx_checkpoint="y") + + with patch("platform.system", return_value="Darwin"), \ + patch("platform.machine", return_value="arm64"), \ + patch("shutil.which", return_value="/usr/bin/ffmpeg"), \ + patch("fastvideo.mlx_runtime.wan_pipeline.MLXWanPipeline", _StubWan21Pipeline), \ + patch("fastvideo.mlx_runtime.wan_pipeline.MLXWan22Pipeline", _StubWan22Pipeline), \ + patch("fastvideo.mlx_runtime.memory.cleanup_mlx", lambda: None): + generator = MLXWanGenerator(generator_config) + generator.shutdown() + + assert loaded["class_name"] == expected_class_name + + +def test_generator_rejects_non_apple_silicon() -> None: + with patch("platform.system", return_value="Linux"): + with pytest.raises(RuntimeError, match="Apple Silicon"): + MLXWanGenerator(SimpleNamespace(model_root="x", mlx_checkpoint="y")) + + +def test_generator_rejects_intel_mac() -> None: + """Darwin alone isn't enough -- x86_64 Macs have no Metal MLX support.""" + with patch("platform.system", return_value="Darwin"), \ + patch("platform.machine", return_value="x86_64"): + with pytest.raises(RuntimeError, match="Apple Silicon"): + MLXWanGenerator(SimpleNamespace(model_root="x", mlx_checkpoint="y")) + + +def test_generator_rejects_missing_ffmpeg() -> None: + with patch("platform.system", return_value="Darwin"), \ + patch("platform.machine", return_value="arm64"), \ + patch("shutil.which", return_value=None): + with pytest.raises(RuntimeError, match="ffmpeg"): + MLXWanGenerator(SimpleNamespace(model_root="x", mlx_checkpoint="y")) + + +def test_generator_shuts_down_its_worker_thread_when_load_fails() -> None: + """__init__ must not leak a running thread pool if the pipeline fails to load.""" + import threading + + class _BrokenPipeline: + + def __init__(self, **_kwargs): + raise RuntimeError("checkpoint is corrupt") + + threads_before = {thread.name for thread in threading.enumerate()} + + with patch("platform.system", return_value="Darwin"), \ + patch("platform.machine", return_value="arm64"), \ + patch("shutil.which", return_value="/usr/bin/ffmpeg"), \ + patch("fastvideo.mlx_runtime.wan_pipeline.MLXWanPipeline", _BrokenPipeline): + with pytest.raises(RuntimeError, match="checkpoint is corrupt"): + MLXWanGenerator(SimpleNamespace(model_path="FastVideo/FastMetal-1.3B-QAD", model_root="x", + mlx_checkpoint="y")) + + # The worker thread pool must not outlive the failed __init__ call. + threads_after = {thread.name for thread in threading.enumerate()} + assert not any("wan-mlx" in name for name in threads_after - threads_before) diff --git a/fastvideo/tests/mlx/test_mlx_wan_pipeline.py b/fastvideo/tests/mlx/test_mlx_wan_pipeline.py new file mode 100644 index 0000000000..7f16559bcf --- /dev/null +++ b/fastvideo/tests/mlx/test_mlx_wan_pipeline.py @@ -0,0 +1,189 @@ +# SPDX-License-Identifier: Apache-2.0 +"""MLXWanPipeline.__init__ validation. + +The constructor only does filesystem/path checks (no real weights loaded, no +mlx.core import) -- these run for real here, no stubbing, since nothing in +__init__ needs Apple Silicon. Generation itself (.generate()) does need +Metal and is out of scope for this file; see mlx_wan_server tests for the +serving-layer coverage that stubs it out. +""" +from __future__ import annotations + +import json +from pathlib import Path + +import pytest + +from fastvideo.mlx_runtime.wan_pipeline import ( + MLXWan22Pipeline, + MLXWanPipeline, + _resolve_wan_torch_dtype, +) + + +def _make_model_root(tmp_path: Path) -> Path: + """A model root with just enough structure to pass the constructor's checks.""" + model_root = tmp_path / "model_root" + (model_root / "tokenizer").mkdir(parents=True) + (model_root / "text_encoder").mkdir(parents=True) + return model_root + + +def _make_packed_checkpoint(tmp_path: Path, name: str = "FastMetal-1.3B-QAD-mlx", in_channels: int | None = None) -> Path: + """A directory shaped like a real packed MLX DiT checkpoint.""" + checkpoint = tmp_path / name + checkpoint.mkdir(parents=True) + manifest = {"config": {"in_channels": in_channels}} if in_channels is not None else {} + (checkpoint / "mlx_dit.json").write_text(json.dumps(manifest)) + (checkpoint / "mlx_dit.safetensors").write_bytes(b"") + return checkpoint + + +def test_init_accepts_a_valid_model_root_and_checkpoint(tmp_path) -> None: + model_root = _make_model_root(tmp_path) + checkpoint = _make_packed_checkpoint(tmp_path) + pipeline = MLXWanPipeline(model_root=model_root, mlx_checkpoint=checkpoint) + assert pipeline.model_root == model_root + assert pipeline.mlx_checkpoint == checkpoint + + +def test_init_rejects_missing_tokenizer_dir(tmp_path) -> None: + model_root = tmp_path / "model_root" + (model_root / "text_encoder").mkdir(parents=True) + checkpoint = _make_packed_checkpoint(tmp_path) + with pytest.raises(FileNotFoundError, match="tokenizer"): + MLXWanPipeline(model_root=model_root, mlx_checkpoint=checkpoint) + + +def test_init_rejects_missing_text_encoder_dir(tmp_path) -> None: + model_root = tmp_path / "model_root" + (model_root / "tokenizer").mkdir(parents=True) + checkpoint = _make_packed_checkpoint(tmp_path) + with pytest.raises(FileNotFoundError, match="text_encoder"): + MLXWanPipeline(model_root=model_root, mlx_checkpoint=checkpoint) + + +def test_init_rejects_nvidia_fastwan_qad_checkpoint_name(tmp_path) -> None: + """FastWan-QAD is the NVIDIA NVFP4/FP8 release; loading it through MLX + silently requantizes the wrong weights (see checkpoint_compat.py).""" + model_root = _make_model_root(tmp_path) + nvidia_checkpoint = tmp_path / "FastWan-QAD-1.3B" + nvidia_checkpoint.mkdir() + with pytest.raises(ValueError, match="FastWan-QAD"): + MLXWanPipeline(model_root=model_root, mlx_checkpoint=nvidia_checkpoint) + + +def test_init_allows_the_legacy_int8_nvidia_named_directory(tmp_path) -> None: + """FastWan-QAD-INT8-* predates the mlx_dit.json packing convention but is + a real Apple checkpoint; the name-based NVIDIA check must not reject it.""" + model_root = _make_model_root(tmp_path) + checkpoint = _make_packed_checkpoint(tmp_path, name="FastWan-QAD-INT8-1.3B") + MLXWanPipeline(model_root=model_root, mlx_checkpoint=checkpoint) + + +def test_init_does_not_validate_checkpoint_contents(tmp_path) -> None: + """A directory that is neither a packed MLX checkpoint nor NVIDIA-flagged + passes __init__ unexamined -- real content validation happens inside + generate() when the weights are actually loaded (needs Metal).""" + model_root = _make_model_root(tmp_path) + empty_checkpoint = tmp_path / "not_a_real_checkpoint" + empty_checkpoint.mkdir() + MLXWanPipeline(model_root=model_root, mlx_checkpoint=empty_checkpoint) + + +def test_init_rejects_a_wan22_checkpoint(tmp_path) -> None: + """Pointing the Wan2.1 pipeline at a 48-channel FastMetal-5B-QAD checkpoint + would silently produce garbled output (wrong VAE compression assumed) -- + this must be caught here, not discovered downstream.""" + model_root = _make_model_root(tmp_path) + wan22_checkpoint = _make_packed_checkpoint(tmp_path, name="FastMetal-5B-QAD", in_channels=48) + with pytest.raises(ValueError, match="Wan2.2-TI2V"): + MLXWanPipeline(model_root=model_root, mlx_checkpoint=wan22_checkpoint) + + +def test_init_accepts_a_checkpoint_with_declared_16_channels(tmp_path) -> None: + model_root = _make_model_root(tmp_path) + checkpoint = _make_packed_checkpoint(tmp_path, in_channels=16) + MLXWanPipeline(model_root=model_root, mlx_checkpoint=checkpoint) + + +def _write_manifest(tmp_path: Path, body: str) -> Path: + """A packed-checkpoint directory whose mlx_dit.json holds arbitrary JSON.""" + checkpoint = tmp_path / "FastMetal-1.3B-QAD-mlx" + checkpoint.mkdir(parents=True) + (checkpoint / "mlx_dit.json").write_text(body) + (checkpoint / "mlx_dit.safetensors").write_bytes(b"") + return checkpoint + + +@pytest.mark.parametrize("body", ["[]", '"nope"', "42", "null"]) +def test_init_tolerates_a_manifest_that_is_not_an_object(tmp_path, body) -> None: + """The channel probe is documented as best-effort: a readable but wrongly shaped + manifest must fall through to the weight load, not raise out of __init__.""" + model_root = _make_model_root(tmp_path) + MLXWanPipeline(model_root=model_root, mlx_checkpoint=_write_manifest(tmp_path, body)) + + +def test_init_tolerates_a_manifest_whose_config_is_not_an_object(tmp_path) -> None: + model_root = _make_model_root(tmp_path) + checkpoint = _write_manifest(tmp_path, json.dumps({"config": ["in_channels"]})) + MLXWanPipeline(model_root=model_root, mlx_checkpoint=checkpoint) + + +@pytest.mark.parametrize("channels", ["sixteen", [16], {"value": 16}]) +def test_init_tolerates_a_non_numeric_in_channels(tmp_path, channels) -> None: + model_root = _make_model_root(tmp_path) + checkpoint = _write_manifest(tmp_path, json.dumps({"config": {"in_channels": channels}})) + MLXWanPipeline(model_root=model_root, mlx_checkpoint=checkpoint) + + +def test_init_still_reads_a_numeric_string_in_channels(tmp_path) -> None: + """Tolerating junk must not stop the probe recognising a real Wan2.2 checkpoint.""" + model_root = _make_model_root(tmp_path) + checkpoint = _write_manifest(tmp_path, json.dumps({"config": {"in_channels": "48"}})) + with pytest.raises(ValueError, match="Wan2.2-TI2V"): + MLXWanPipeline(model_root=model_root, mlx_checkpoint=checkpoint) + + +def test_resolve_torch_dtype_maps_the_recipe_names() -> None: + """Wan2.1 and Wan2.2-TI2V pass different names; both must resolve exactly.""" + torch = pytest.importorskip("torch") + assert _resolve_wan_torch_dtype("bf16") is torch.bfloat16 + assert _resolve_wan_torch_dtype("fp16") is torch.float16 + assert _resolve_wan_torch_dtype("fp32") is torch.float32 + + +def test_resolve_torch_dtype_rejects_an_unknown_name() -> None: + pytest.importorskip("torch") + with pytest.raises(ValueError, match="Unsupported text-encoder dtype"): + _resolve_wan_torch_dtype("float8") + + +class TestMLXWan22Pipeline: + """Mirrors the MLXWanPipeline coverage above for the Wan2.2-TI2V (5B) pipeline.""" + + def test_init_accepts_a_valid_model_root_and_checkpoint(self, tmp_path) -> None: + model_root = _make_model_root(tmp_path) + checkpoint = _make_packed_checkpoint(tmp_path, name="FastMetal-5B-QAD", in_channels=48) + pipeline = MLXWan22Pipeline(model_root=model_root, mlx_checkpoint=checkpoint) + assert pipeline.model_root == model_root + + def test_init_rejects_missing_tokenizer_dir(self, tmp_path) -> None: + model_root = tmp_path / "model_root" + (model_root / "text_encoder").mkdir(parents=True) + checkpoint = _make_packed_checkpoint(tmp_path, name="FastMetal-5B-QAD", in_channels=48) + with pytest.raises(FileNotFoundError, match="tokenizer"): + MLXWan22Pipeline(model_root=model_root, mlx_checkpoint=checkpoint) + + def test_init_rejects_a_wan21_checkpoint(self, tmp_path) -> None: + """The reverse mistake: pointing the 5B pipeline at a 1.3B/14B checkpoint.""" + model_root = _make_model_root(tmp_path) + wan21_checkpoint = _make_packed_checkpoint(tmp_path, in_channels=16) + with pytest.raises(ValueError, match="MLXWanPipeline for 1.3B/14B"): + MLXWan22Pipeline(model_root=model_root, mlx_checkpoint=wan21_checkpoint) + + def test_init_does_not_validate_checkpoint_contents(self, tmp_path) -> None: + model_root = _make_model_root(tmp_path) + empty_checkpoint = tmp_path / "not_a_real_checkpoint" + empty_checkpoint.mkdir() + MLXWan22Pipeline(model_root=model_root, mlx_checkpoint=empty_checkpoint)