Architecture

MLX-Gen is an independent package forked from mflux. It keeps the MLX-native model runtime from mflux while exposing a cleaner mlxgen command surface for new users and applications. The video paths include Wan2.2 TI2V-5B text-to-video, TI2V-5B first-frame image-to-video, Wan2.2 A14B text-to-video, Wan2.2 A14B image-to-video and video-to-video, Wan2.1-VACE-1.3B conditioning, and an experimental Bernini-R 1.3B role-aware reference renderer whose current visual validation fails. SeedVR2 image and video restoration use mlxgen upscale.

System Overview

flowchart TD
    U[User or application] -->|mlxgen generate / upscale / capabilities / download / prepare| R[mlxgen CLI router]
    U -->|"load_generation_model(...) / generate_outputs(...)"| P[Python runtime helpers]
    R --> TI[Task inference and capability planner]
    P --> TI
    TI -->|generation plan: task, mode, capability, handler| B[Backend command and runtime selection]
    B --> M[Model runtime variants<br/>FLUX.2 Klein, Qwen, Z-Image, ERNIE, FIBO, Bonsai, Wan, SeedVR2]
    M --> W[Weight loading<br/>single or factored pinned sources / local packages]
    M --> CB[Progress callbacks and runtime memory telemetry]
    M --> O[Saved images and videos with metadata]

The router consumes routing options such as --model, primary images/videos, semantic --reference-image values, --mask-path, and strength/padding flags, resolves one generation plan through the capability planner, and forwards a normalized invocation to the selected backend command. The same planner powers the Python helpers, so embedded applications make identical routing decisions before loading any weights. Reference images remain a distinct role: they are never silently reinterpreted as first-frame image-to-video inputs. This durable boundary is recorded in ADR 0007.

Package Shape

• PyPI distribution: mlx-gen
• Public CLI root: mlxgen
• Compatibility CLI for older SeedVR2 scripts: mflux-upscale-seedvr2
• New application import identity: mlxgen
• Current runtime internals: primarily mflux.*

The project keeps some mflux vocabulary and compatibility entry points while the fork evolves. New docs and integrations should use mlxgen commands where the workflow is available and treat mflux.* internals as inherited implementation detail unless a specific model class or dedicated command currently requires them.

Command Boundary

The public command surface separates setup from inference:

• mlxgen download is an explicit cache population command.
• mlxgen prepare is an explicit local model-folder creation command.
• mlxgen generate is the inference command and does not start downloads by default.
• mlxgen upscale is the inference command for SeedVR2 image and video restoration.

This boundary is important for embedded workflow systems such as AbstractVision: a generation request should not unexpectedly start a large network transfer in the middle of a larger job.

Model File Lifecycle

Source model files usually come from Hugging Face. They can be used in two ways:

1. Cache the source files with mlxgen download and run by alias or repository id.
2. Create a reusable local MLX-Gen model package with mlxgen prepare --model ... --path ... --quantize ....

flowchart LR
    HF[Hugging Face source repositories] -->|mlxgen download| C[Local Hugging Face cache]
    HF -->|mlxgen prepare --quantize| PKG[Local MLX-Gen model package<br/>optionally q4/q8 + generated model card]
    C --> G[mlxgen generate / mlxgen upscale]
    PKG --> G
    G --> OUT[Images and videos with embedded or sidecar metadata]

MLX-Gen model packages use the MLX/mflux saved-weight layout. They may contain MLX quantization tensors and generated Hugging Face model cards. They are intended for MLX-Gen and compatible mflux code, not direct Diffusers or Transformers loading.

Video support follows the same setup/runtime boundary. Wan2.2 loads local source files and writes MP4 output. Text-to-video starts from random video latents. TI2V-5B image-to-video VAE-encodes the first frame, masks first-frame timesteps, keeps the condition active during denoising, and reinserts the condition before decode. A14B uses Diffusers-compatible two-transformer boundary routing and, for the separate I2V model, concatenated image-condition latents.

Bernini is a dedicated single-transformer renderer variant rather than a VACE or first-frame-I2V branch. It independently VAE-encodes a source video and each ordered reference, patches each into a heterogeneous token segment with a source-ID rotary phase, appends the noisy target as source zero, and returns only the target segment. R2V uses chained APG across empty/reference/text branches; RV2V uses four sequential empty/video/video+reference/text branches; V2V uses source-conditioned APG. Sequential branches, condition-by-condition encoding, denoiser release, and streamed VAE decode make the path practical on bounded unified-memory hosts.

Bernini also exercises factored model ownership. mlxgen download resolves pinned tokenizer, UMT5, VAE, and scheduler files from a Wan2.1 base repository and the renderer transformer from a separate ByteDance repository. Preflight validates both revisions, component configs, required files, and aggregate free space; provenance for every component is retained in output metadata.

Runtime Failure Contract

Runtime model construction and generation use files that are already available locally. Missing required files raise DownloadRequiredError, which is also a FileNotFoundError for compatibility with existing callers.

The error includes actionable command fields such as download_command and, when applicable, prepare_command. CLI entry points print the human-readable error without a traceback for common missing-artifact cases.

Quantization Policy

Quantization is model-specific. Qwen and ERNIE q4 paths use mixed q4/q8 policies because fully q4 checkpoints can lose coherent generative behavior for those model families. SeedVR2 3B and 7B use q4/q8 MLX-Gen packages for the transformer linears and VAE attention linears that support MLX quantization. Bonsai Image uses Prism's pre-packed ternary 2-bit transformer path instead of MLX-Gen's q4/q8 prepare flow; it follows the same quality principle of keeping sensitive paths at higher precision, but ships as a pre-packed artifact. Bernini is BF16-only: its generic Wan q4 path failed transformer/video gates, while nominal q8 quantized no renderer linears and would misstate the execution. Other model families keep their existing predicates unless model behavior requires a dedicated policy.

See Quantization for the current rules.

Python Integration Boundary

The current Python API still exposes many model classes through inherited mflux modules. MLX-Gen's near-term integration contract is:

• prepare files before constructing models;
• fail early when required artifacts are missing;
• keep model instances as stateful runtime objects;
• publish lightweight progress events through mflux.callbacks.ProgressEvent without exposing latents or model tensors;
• expose clearer public orchestration APIs over time without breaking existing compatibility paths unnecessarily.

See Python Integration and API And CLI for current usage.