AutomatosX/AX-DiffusionGemma-26B-A4B-IT-MLX-OptiQ-4bit

AX DiffusionGemma 26B-A4B IT MLX OptiQ 4-bit

Parameter count: approximately 25.25B logical parameters reported for the packaged language tensors (26B-class MoE, ~4B active per token). 4-bit names the low OptiQ tier, not a 4B model-size claim. The BF16 vision tower is packaged as a sidecar.

This is a revision-pinned AutomatosX package of the upstream OptiQ conversion mlx-community/diffusiongemma-26B-A4B-it-OptiQ-4bit at commit 30f3c7c7746bf41cfd1a290155cc3b777ab588b9.

AutomatosX did not train DiffusionGemma and did not re-run the OptiQ conversion. The weights, OptiQ allocation, tokenizer, processor, and generation settings are byte-identical to the pinned upstream revision. AutomatosX adds catalog packaging: model card, license copy, machine-readable provenance, and a verified logical total_parameters field on the shard index.

Why this OptiQ: DiffusionGemma is a block/masked-diffusion model. Upstream OptiQ measures per-layer KL on denoising-canvas logits along the actual denoising schedule (not autoregressive next-token KL). Attention, routers, and sensitive blocks stay at 8-bit; less sensitive expert projections take 4-bit. The SigLIP vision tower remains BF16 in optiq/optiq_vision.safetensors.

Package details

Property	Value
Format	MLX Safetensors
Quantization	OptiQ mixed 4/8-bit, group size 64
Measured allocation	52 components at 4-bit; 246 at 8-bit
Target / achieved BPW	4.66 / 4.6854
Vision tower	Bundled BF16 sidecar
Logical parameters (packaged)	25,250,898,462
Download size	Approximately 17.9 GB
Intended hardware	Apple Silicon
Runtime requirement	mlx-optiq ≥ 0.3.2

DiffusionGemma is not loadable by stock mlx-lm / plain mlx-vlm as a normal autoregressive model. Use OptiQ's vendored diffusion decoder.

Download

hf download AutomatosX/AX-DiffusionGemma-26B-A4B-IT-MLX-OptiQ-4bit \
  --local-dir ./AX-DiffusionGemma-26B-A4B-IT-MLX-OptiQ-4bit

Use with mlx-optiq

pip install -U "mlx-optiq>=0.3.2"

from optiq.vlm.diffusion_gemma import load, generate
from PIL import Image

model, tokenizer = load("AutomatosX/AX-DiffusionGemma-26B-A4B-IT-MLX-OptiQ-4bit")

# text
print(generate(model, tokenizer, "Write a haiku about Apple Silicon."))

# image + text
print(generate(
    model,
    tokenizer,
    "What is in this image?",
    images=[Image.open("photo.jpg")],
    sampler="confidence-threshold",  # recommended: much faster than entropy-bound
))

Try with AX Engine

AX Engine has an experimental native MLX path for DiffusionGemma:

ax-engine serve AutomatosX/AX-DiffusionGemma-26B-A4B-IT-MLX-OptiQ-4bit --download --port 31418

Treat DiffusionGemma support as experimental. First visible output is a committed diffusion block, so ordinary autoregressive TTFT comparisons do not apply.

Validation

• Upstream OptiQ metadata present (method=optiq)
• Weight shards and vision sidecar present and hashed in ax_provenance.json
• Local generation smoke via optiq.vlm.diffusion_gemma: passed
• AX Engine doctor: ready

Sample generation (truncated): "2:\nWait, it's not on the ladder, it's \"walking down\"\nSub-:-down:\nSub-walking-down:\nSub-walking-down:\nWait, the 2-walking down.\n2:\nWait, it's not on the ladder"

Quantization can change model quality. Evaluate on your own vision and text tasks. Prefer sampler="confidence-threshold" for substantially higher tok/s without quality loss versus the model default entropy-bound sampler.

Provenance

UPSTREAM_README.md preserves the mlx-community OptiQ card (Capability Score and build recipe). ax_provenance.json pins the source revision and SHA-256 digests for every Safetensors artifact.

License and responsible use

Apache License 2.0. See LICENSE, the google/diffusiongemma-26B-A4B-it model card, and UPSTREAM_README.md for usage limitations and safety guidance.