Qwen-Image-Flash-bf16 (MLX)

An MLX-ready bf16 snapshot of nvidia/Qwen-Image-Flash — NVIDIA's DMD2 four-step distillation of Qwen/Qwen-Image (20.43B-parameter MMDiT + Qwen2.5-VL text encoder + Qwen-Image VAE).

What this repo is, precisely

The weights are byte-identical to upstream. MLX and MLX-Swift consume the diffusers-layout safetensors directly — key sanitizing and the conv-layout transposes happen at load — so no dtype conversion or re-layout is involved, and none is claimed. This snapshot exists so an MLX consumer can materialize the whole pipeline from one repo, and it adds exactly one thing upstream does not ship:

• tokenizer/tokenizer.json — upstream ships slow-tokenizer files only (vocab.json + merges.txt), which the Swift swift-transformers stack cannot read. The fast tokenizer here was converted from those same files and verified to produce identical token ids (the encode template tokenizes to 34 prefix tokens either way, matching prompt_template_encode_start_idx = 34).

If you are using PyTorch/diffusers, use the upstream repo instead — it is the same weights, and NVIDIA is the source of truth.

Inference notes that are easy to get wrong

The distillation internalized CFG 4.0, and the packaged scheduler is static shift-3 (use_dynamic_shifting: false):

• num_inference_steps = 4, true_cfg_scale = 1.0 — applying CFG again double-counts guidance the student already absorbed, and doubles the transformer evaluations per step for nothing.
• The four-step trajectory is sigmas [1.0, 0.9, 0.75, 0.5, 0.0]. This is what the student was distilled on; substituting a dynamic-shift schedule degrades output.
• Tested at 1024 × 1024. Use width/height divisible by 16.

Use from Swift (MLXEngine)

import MLXQwenImageFlash
import MLXToolKit

let package = QwenImageFlashPackage(configuration: .init())   // materializes from this repo
try await package.load()
let response = try await package.run(T2IRequest(
    prompt: "A red fox in a snowy pine forest at golden hour, photorealistic",
    width: 1024, height: 1024, seed: 42)) as! T2IResponse
// response.image: PNG, 1024×1024 — 4 steps

Port: xocialize/qwen-image-edit-swift (MIT). The text-to-image path shares the parity-locked Qwen-Image DiT / VAE / Qwen2.5-VL core with the Qwen-Image-Edit packages — the Flash checkpoint is a pure weight delta on the same architecture.

Validated against PyTorch fp32 goldens (diffusers 0.37.1): prompt embeds cos 0.9999926 (bf16) · DiT step-0 cos 1.0000002 (fp32 CPU stream) / 0.99836 (bf16 GPU) · end-to-end 1024²/4-step render coherent in 83.3 s on an M5 Max, consistent with the diffusers bf16 reference render.

M5-class GPUs / mlx-swift ≤ 0.31.6: the 1024² text-to-image sequence (4096 image tokens) lands exactly on the upper edge of the NAX split-K GEMM window (mlx#3797) for this model's feed-forward down-projection, which returns NaN there. The Swift port works around it with an exact row-chunk; other MLX consumers on an unpatched runtime should verify a 1024² render before trusting output.

License

Governing terms: NVIDIA Open Model License Agreement. Additional information: Apache License 2.0 (LICENSE).

Licensed by NVIDIA Corporation under the NVIDIA Open Model License

Commercial use, derivative models, and redistribution are permitted; if you redistribute these weights you must pass on the Agreement and this notice (§3.1). See NOTICE.