docs/TECHNICAL.md

MiniMax-H3-Image-mlx

Native MLX/Metal inference for MiniMax-H3-Image on Apple silicon. The default uses the original BF16 transformer, all six distilled PDD steps, dense attention, FP32 timestep/input/output projections, and the released overlapping BF16 image VAE decoder. Q8, Q4, custom Metal normalization and fewer steps are explicit experiments.

The port reads the original safetensors directly. It does not need CUDA, SGLang, a remote API, or a second copy of the checkpoint. PyTorch is used only for the original CPU random-number stream and independent numerical tests; the model operators run in MLX on Metal.

The latest operator-tuning report applies the local tmp/mlx-skills workflow: native-order Q/K fusion, M5 FC2 tuning, physical-memory accounting, real activation sweeps and complete trajectory checks. Matched image stacks include the exact modes and links to earlier precision-loss examples.

On the tested Mac, paired warm latency improves 5.07% with dense BF16 storage and 3.57% with lossless packed storage, each against its own current control. Every pixel and saved denoising tensor matches on four cases. The largest sampled process footprint across candidate cases is 40.50 GiB dense / 32.56 GiB packed; see the report for timing ranges and memory scope.

The further optimization report covers compiled MLX, exact weight packing, asynchronous loading, PyTorch/MPS, custom Metal kernels, and a complete Core ML VAE tested on CPU, GPU, ANE and all devices. It includes real side-by-side image grids and failed experiments.

Install and generate

An Apple silicon Mac and a compatible MLX Metal installation are required. The tested machine is a 64 GB M5 Pro. BF16 inference uses a resident DiT core of about 36 GiB, plus activations, decoder and allocator overhead. Smaller memory budgets stream a fixed subset of blocks; Q8/Q4 reduce memory and change the image. Disk space for the original image model and full encoder is roughly 134 GB, independent of this small code repository.

uv venv --python 3.12
source .venv/bin/activate
uv pip install -e '.[audit]'
h3-mlx doctor

h3-mlx generate \
  --model /path/to/MiniMax-H3-Image \
  --text-encoder /path/to/MiniMax-H3/text_encoder \
  "An editorial photograph of a cobalt perfume bottle on pale stone" \
  --width 2048 --height 2048 --seed 731 --output outputs/bottle.png

requirements-macos.lock records the exact tested environment, including audit tools. Install it before the editable package when reproducing these measurements.

The neighboring checkpoint paths on the development machine are ../minimax-h3-image-opt/models/MiniMax-H3-Image and ../minimax-h3-image-opt/models/MiniMax-H3/text_encoder. Pass them explicitly; the package has no dependency on that repository at runtime.

If weights are not already present, download the official repositories:

hf download MiniMaxAI/MiniMax-H3-Image \
  --revision 6534ea1c979a5f8cb00c2bd1025d5ca4ee8d8901 \
  --local-dir models/MiniMax-H3-Image
hf download MiniMaxAI/MiniMax-H3 --include 'text_encoder/*' \
  --revision 0919fb684102ef115126612f3fbb84d1027e7317 \
  --local-dir models/MiniMax-H3

The model expects detailed English prompts. Prompt expansion is not performed silently. Requests smaller than 4 MP use the upstream larger generation canvas and are resized with LANCZOS. --min-generation-area 1024 disables that enlargement for quick low-resolution development trials; these are outside the model's usual 4–20 MP working range.

Every output PNG has a sibling JSON recording the prompt, seed, noise source, canvas, schedule, precision options, per-step times, model-load times, memory and software/hardware versions. --save-trace also saves conditioning, each velocity/Euler state and the final latent for numerical investigations.

Precision and performance controls

For the tested 64 GB M5 Pro, the recommended exact configuration is:

h3-mlx generate --model /path/to/MiniMax-H3-Image \
  --text-encoder /path/to/MiniMax-H3/text_encoder \
  --compile-blocks --eval-every 0 --prefetch-blocks 4 --wired-gb 42 \
  --compile-vae --vae-batch 4 \
  --kernel metal-qk-exact --tune-matmul \
  "A detailed image description" --output outputs/optimized.png

This keeps all six heads and the original BF16/FP32 arithmetic. The M5 tuning is scoped to MLX 0.32.4 and measured token ranges; other devices, versions and shapes retain native execution. See the report for complete image/trajectory checks and paired timing ranges. For lower memory use, add --weight-storage packed12 and use --wired-gb 32 on this machine. Packing is lossless, builds an approximately 27 GiB disk cache on first use, and trades a small amount of compute for smaller resident weights. Wired memory settings must fit the actual Mac; the runtime enforces Metal's limit.

Option	Effect
Default	BF16 core, original schedule, dense Metal SDPA, FP32 precision islands
--compile-blocks --eval-every 0	Compile DiT blocks and evaluate the complete warm step
--prefetch-blocks 4	Bound asynchronous loading/execution overlap to four blocks when weights need loading
--kernel metal-qk-exact	Native-order Q/K RMSNorm + partial RoPE fusion; keeps native full-width normalization
--tune-matmul	Metal 4 FC2 tile selected on real activations; M5 Pro/MLX 0.32.4 and 4096–8192 joint tokens, native fallback otherwise
--kernel metal-exact	Also fuse full-width RMSNorm/AdaLN; exact on the recorded suite, with a measured physical-memory tradeoff
--compile-vae --vae-batch 4	Faster decoder; identical pixels on the tested product and portrait
--vae-dtype float32	Original FP32 decoder arithmetic; slower, with a same-latent precision grid
--weight-storage packed12	Lossless BF16 storage with exact GPU reconstruction before native GEMMs
--wired-gb 42	Process-local residency on the tested 64 GB Mac; does not modify system settings
--no-refiner-cache	Bypass exact refined-conditioning caching for audits
--keep-refiners	Keep refiner weights for new conditioning instead of releasing them
--kernel metal --norm-threads 256	Experimental fused RMSNorm/AdaLN; preserves intermediate BF16 rounding but changes reduction order
--kernel metal-qk	Also fuse per-head Q/K normalization and partial 3D RoPE
--vae-batch 8	Decode several complete native tiles together; GEMM rounding can change
--mlp-chunk 1024	Lower activation memory; 0 uses a full-sequence MLP
--attention-chunk 512	Bound attention query working size while keeping all keys; benchmark before using
--dit-cache-gb 16	Pin up to 16 GiB of blocks and stream the rest without LRU thrashing
--quantize 8 / --quantize 4	Weight-only affine quantization of DiT attention/MLPs, group size 64; AdaLN and precision islands remain native
--group-size 32	Smaller quantization groups; separate quality/memory tradeoff
--encoder-quantize 8	Independently quantize the fifty language layers
--experimental-steps 0,2,4	Three original heads with larger Euler intervals; a preview approximation
--offload-after-denoise	Release DiT before decoding for tighter memory budgets
--rng mlx	Avoid the PyTorch dependency, using a different seed/noise stream
--vae-backend coreml	Experimental FP16 ML Package decoder; requires an exported package and matching tile batch

The JSON memory report separates MLX allocation counters from macOS parent physical footprint. Its process lifetime peak includes earlier requests; the experiment harness also samples each request separately. These counters overlap and must not be added. The text encoder now reads only requested vocabulary rows, preserving original BF16 payloads and avoiding a full vocabulary allocation on a conditioning-cache miss.

Core ML GPU and ALL had small image differences. The batch-eight ANE export had severe reconstruction artifacts and is a diagnostic option. See the device comparison grid and report before selecting it. The full denoiser remains MLX in every runtime configuration; only the VAE has a Core ML export.

For repeated quantized runs, preconvert attention/MLP matrices once. This avoids online conversion at startup; the original checkpoint remains the source of the exact modulation tables and FP32 layers:

python -m h3_mlx.convert --model /path/to/MiniMax-H3-Image --output models/core-q8 --bits 8
h3-mlx generate --model /path/to/MiniMax-H3-Image \
  --text-encoder /path/to/MiniMax-H3/text_encoder \
  --quantized-core models/core-q8 "A detailed image description" --output outputs/q8.png

Conversion writes one block per shard, verifies the source identity on load, and records SHA256 hashes. It requires additional disk space (roughly half the native core size for Q8). Quantization is a memory tradeoff; it does not automatically improve image-generation latency.

MLX's M5 defaults can reduce the precision of FP32 matrix operations. The package sets MLX_ENABLE_TF32=0 and MLX_CONV_WINOGRAD=0 before importing MLX, retaining full precision for the model's sensitive layers. Explicit environment overrides are respected and recorded. When embedding this package, import h3_mlx before using MLX, or set those variables before starting Python. See MLX's precision documentation.

Cross-backend byte identity with CUDA is not promised. Reduction order, attention arithmetic and GEMM algorithms can change a seeded diffusion trajectory. The tests distinguish ATen RMSNorm from Qwen's different BF16 rounding order and allow only narrowly bounded cross-device errors.

Reuse resident models

A fresh CLI process includes checkpoint loading. Use one Pipeline for multiple requests to retain the DiT and VAE in unified memory. Conditioning and native modulation tables are also cached on disk.

from h3_mlx.runtime import Pipeline

pipe = Pipeline("/path/to/MiniMax-H3-Image", "/path/to/MiniMax-H3/text_encoder")
try:
    for seed in (731, 732, 733):
        image, record = pipe.generate(
            "A detailed editorial photograph of a blue perfume bottle on pale stone",
            seed=seed, output=f"outputs/bottle-{seed}.png",
        )
finally:
    pipe.clear()

The package does not change system wired-memory settings. It caps block residency below Metal's recommended working set and keeps a fixed resident subset when the full core does not fit. Use a smaller block cache and MLP chunks for larger canvases or lower-memory Macs.

Image editing

The MLX implementation includes the Qwen3-VL vision tower, all three DeepStack features, multimodal RoPE, a single-frame reduction of the causal CNN VAE encoder, posterior sampling, and fixed noised reference anchors. References are processed in input order, up to sixteen:

h3-mlx generate \
  --model /path/to/MiniMax-H3-Image \
  --text-encoder /path/to/MiniMax-H3/text_encoder \
  "Change the bottle in <Picture 1> to emerald glass while preserving the label" \
  --image bottle.png --seed 42 --output outputs/edit.png

Qwen receives a 1024 px short edge. VAE references retain short edges in 1024–2048 px, with the native 16 px crop and even latent grid. Target noise, anchor noise and posterior samples consume the CPU generator in upstream order. This path requires more conditioning work than text-only generation.

Measurements and image comparisons

The experiment report describes measured latency, precision, failed performance hypotheses and remaining work. Raw operator trials are in operators.json. Measurements describe this Mac and the recorded prompts; they are not a general speed or image quality guarantee.

Generate a vertical stack, identical 1:1 crops and amplified RGB differences:

python -m h3_mlx.compare \
  --reference outputs/suite/product/native-1.png \
  --candidate outputs/suite/product/tuned-1.png \
  --candidate outputs/suite/product/q8-1.png \
  --label 'BF16 native' --label 'BF16 tuned' --label 'Q8' \
  --output comparisons/product

The comparison tool verifies prompt, seed, dimensions, generation canvas, noise generator and reference inputs from the inference records. It scores full-resolution RGB, includes the exact crop coordinates and hashes every source PNG. PSNR/SSIM quantify similarity to the reference trajectory; visual inspection remains necessary. The tool never fabricates or retouches evidence. It also writes grid.png and details-grid.png for direct side-by-side review.

Reproduce experiments

python -m pytest -q
python -m h3_mlx.benchmark --tokens 256 4096 --repeats 7
python experiments/run_suite.py \
  --model /path/to/MiniMax-H3-Image \
  --text-encoder /path/to/MiniMax-H3/text_encoder \
  --profiles native tuned q8 q4 preview --cases product portrait --repeat 3

The suite labels requests that include model loading separately from requests that reuse resident models. Conditioning-cache hits and all per-stage times are retained in the raw records. Do not run another GPU workload concurrently when measuring latency.

The central memory optimization computes the exact schedule's small AdaLN outputs once and releases roughly 26 GB of modulation-only DiT weights. It preserves every learned matrix and uses no curve fitting, pruning, interpolation or timestep skipping in the default path. The fifty remaining transformer blocks can then reside in this 64 GB Mac's unified memory.

Code is Apache-2.0; model weights retain their separate MiniMax licenses. NOTICE and source provenance identify the SGLang and Transformers implementations used for the port.