Roadmap: MLX backend strategy + HiDream-O1-Image-Dev

[localkin]

Context

Tracking the next architectural bet for ollamadiffuser. v2.0.14 added two new registry entries (FLUX.1-Kontext-dev, Chroma1-HD) that ride on existing strategies. The next batch needs a new strategy class — specifically for Apple Silicon (MLX) native inference.

Why MLX

ollamadiffuser currently routes everything through PyTorch (diffusers library). On Apple Silicon, PyTorch + MPS is slower than Apple's MLX framework for many ops (often 2-3× on the same hardware).

Concretely, the mlx-community org publishes MLX-quantized ports of major diffusion models, and mflux (0.17.5, April 2026) is the de-facto MLX inference library — supports FLUX.1, FLUX.2 (4B + 9B), Z-Image, FIBO, SeedVR2, Qwen-Image, Depth Pro.

Status of competition: ComfyUI, A1111, InvokeAI, diffusers-studio — none have a native MLX backend. This is a real differentiation opportunity for the "Ollama-for-X" niche.

Concrete plan

Phase 1: MLXStrategy base class

New file: ollamadiffuser/core/inference/strategies/mlx_strategy.py

Mirror the shape of FluxStrategy / GenericPipelineStrategy but route through mlx-community model packages. Registry entries opt in via model_type: "mlx" + parameters.mlx_backend: "mflux" | "mlx-community".

Phase 2: HiDream-O1-Image-Dev as first consumer

HiDream-ai/HiDream-O1-Image-Dev (MIT, May 8 2026):

• 8B unified transformer on Qwen3-VL backbone
• No VAE — predicts raw 32×32 RGB patches directly
• Already has an MLX port: mlx-community/HiDream-O1-Image-Dev-mlx-bf16
• 16 GB peak at 1024×1024 — fits the M4 16GB
• ~67s/image at 1024² on Apple Silicon

This is the natural first consumer of MLXStrategy because (a) it has zero diffusers integration, so the alternative is a custom HF transformers wrapper; (b) the MLX path is already published.

Phase 3: mflux integration (FLUX family)

After HiDream-O1 proves the pattern, extend MLXStrategy to support mflux:

• FLUX.1-schnell / dev (already in registry as PyTorch — would get an mlx variant)
• FLUX.2-klein-4B / 9B (gives M4 16GB users a usable FLUX.2 path)
• Z-Image-Turbo (already in registry — would get MLX acceleration)

Hardware constraints

Maintainer can only test on:

• Mac Pro M1 32GB (~24 GB UMA)
• Mac Mini M4 16GB (~12 GB UMA)

So this work is doubly good for the project: solves a real differentiation gap AND aligns with what the maintainer can actually develop on.

Effort estimate

• Phase 1 (base class): ~2-3 days
• Phase 2 (HiDream-O1): ~1-2 days on top of Phase 1
• Phase 3 (mflux): ~2-3 days

Tracking

Reply to this issue with implementation progress, or open child issues per phase if discussion warrants.

cc / community: interested in helping? PRs welcome — the test fixture pattern from v2.0.13 (tests/unit/test_*_endpoint.py uses TestClient + monkeypatched loaded_models) is the template for adding a new strategy without needing actual GPU weights.

[localkin]

Phase 1 shipped in v2.0.15 (commit 02d3620). Issue stays open for Phase 2 + 3.

What's in

• New MLXStrategy in core/inference/strategies/mlx_strategy.py routes through mflux for Apple-Silicon-native FLUX inference.
• Dispatch wired: model_type: "mlx" in registry → MLXStrategy().
• Refuses to load on anything other than Darwin/arm64 with a clear error message — so registry entries are inert on Linux/Windows, no breakage.
• Overrides unload() (no torch cleanup needed) and load_lora_runtime() (mflux takes LoRAs at construction time, runtime LoRA hot-swap is deferred — see follow-ups below).
• 16 unit tests covering platform guard, config validation, mocked load+generate, registry-default propagation, and the runtime-LoRA error contract.

New registry entries

Name	mflux variant	quantize	disk	rec. VRAM	license
flux.1-schnell-mlx	flux1	8-bit	~14 GB	16 GB (M1 32GB)	Apache 2.0
flux.1-schnell-mlx-q4	flux1	4-bit	~8 GB	12 GB (M4 16GB)	Apache 2.0
flux.1-dev-mlx	flux1	8-bit	~14 GB	16 GB	FLUX Non-Commercial

Install

pip install 'ollamadiffuser[mlx]'   # mflux >= 0.17.0

mflux ships only the darwin/arm64 wheel, and the optional extra is pinned with sys_platform == 'darwin' and platform_machine == 'arm64' so pip refuses the install elsewhere — the extra is a no-op on non-Mac hosts.

Phase 2 — mflux family expansion (next)

SUPPORTED_MLX_VARIANTS is currently {"flux1"}. Adding to this set is a 5-10 line change per variant in _resolve_model_class():

• flux2 → mflux ships models/flux2/ already. Need to wire entries for flux2-klein-4b, flux2-klein-9b.
• z_image → mflux ships models/z_image/. Z-Image-Turbo would be the obvious entry.
• flux1-kontext → mflux ships models/flux/variants/kontext/. Pairs nicely with the existing flux.1-kontext-dev (PyTorch) added in v2.0.14.
• qwen-image and qwen-image-edit → mflux's ModelConfig.from_name already exposes these.

Phase 3 — HiDream-O1 (separate inference stack)

mlx-community/HiDream-O1-Image-Dev-mlx-bf16 is not on mflux — it ships its own loader. Will need either:

1. A separate HiDreamO1MLXStrategy, or
2. A more abstract MLXBackend enum that the existing MLXStrategy dispatches on (mlx_backend: "mflux" | "mlx-community").

I'd lean toward option 2 to keep one strategy class as the long-term home.

Open follow-up sub-tasks

If anyone wants to PR these in order of impact:

1. Wire flux1-kontext variant — pairs with the existing PyTorch flux.1-kontext-dev entry. ~30 min.
2. Wire flux2-klein-4b/9b MLX entries — the existing PyTorch flux.2-klein-4b gets a 2-3× faster sibling. ~1 hour.
3. Wire z_image-turbo MLX entry — already in registry as PyTorch; the MLX path makes it usable on M4 16GB at speed. ~1 hour.
4. Runtime LoRA via reconstruction — mflux takes lora_paths/lora_scales in the constructor. A clean implementation rebuilds the Flux1 instance on load_lora_runtime(). ~2 hours.
5. HiDream-O1 sub-strategy — biggest item. ~3-4 days.

Re-pin if you'd like one of these picked up next.

[localkin]

Phase 2 shipped in v2.0.16 (commit 60e2db1). Phase 3 (HiDream-O1) remains open.

What's in

MLXStrategy now routes all five mflux model families. Seven new registry entries on top of Phase 1's three (10 MLX entries total).

Entry	Variant	Quant	License	Fits M4 16GB
flux.1-kontext-dev-mlx	flux1-kontext	Q8	Non-Commercial	❌
flux.2-klein-4b-mlx	flux2	Q8	Apache 2.0	✅
flux.2-klein-9b-mlx	flux2	Q8	Apache 2.0	❌
z-image-turbo-mlx	z_image	Q8	Apache 2.0	✅
qwen-image-mlx	qwen-image	Q8	Apache 2.0	❌
qwen-image-edit-mlx	qwen-image	Q8	Apache 2.0	❌

Key refactors

1. Unified load path. _resolve_model_class → _resolve_model_and_config returning (ModelClass, mflux_model_config). All five families now load via cls(quantize=Q, model_config=mc) — only Flux1 has the from_name() factory, the other four require ModelConfig.<factory>() (e.g. ModelConfig.flux2_klein_4b()).
2. Per-variant kwarg filtering via inspect.signature() on the bound generate_image(). Necessary because Flux2Klein.generate_image() has no negative_prompt parameter — without filtering this crashes with TypeError.
3. flux1-kontext image-required guard — fails loud with ValueError rather than letting mflux crash mid-pipeline.

Tests

• 13 new variant-resolution tests (each family resolves to right class; each rejects bad model names).
• test_generate_filters_kwargs_for_variants_without_negative_prompt proves the filter actually drops unsupported kwargs.
• Existing 16 Phase-1 tests rewritten against the new resolution API.
• 107 tests pass (78 other + 29 MLX), 8 skipped.

Phase 3 — still open

mlx-community/HiDream-O1-Image-Dev-mlx-bf16 ships its own loader (not on mflux). Two paths:

1. Separate HiDreamO1MLXStrategy (cleanest separation)
2. mlx_backend: "mflux" | "mlx-community" enum on MLXStrategy (less code, more branching)

I'd lean toward option 2 to keep one strategy class as the long-term home. ~3-4 days estimated.

Open follow-ups for community PRs

1. Real-hardware MLX bench — run actual flux.1-schnell-mlx and flux.2-klein-4b-mlx end-to-end on M-series hardware to validate the PyTorch-vs-MLX speedup claim (currently uncalibrated in the registry's performance_notes).
2. Runtime LoRA for MLX backend — mflux takes lora_paths/lora_scales at construction time. Cleanest implementation rebuilds the model on load_lora_runtime(). ~2 hours.
3. mflux-controlnet / mflux-depth / mflux-fill variant wiring — mflux already supports these under models/flux/variants/{controlnet,depth,fill}/. Each adds maybe 30 lines to _resolve_model_and_config().
4. Phase 3: HiDream-O1 as above.

Re-pin if you'd like one of these picked up next.