Image generation engine: /v1/images/generations on MLX — mflux seam, [image] extra, time-to-image doctrine

Parent: #95 · Sequencing: first — this seam unblocks every image-model issue. Maintainer steer (2026-10-06): small + hot models first → after the hot-list refresh, the engine's reference path is Z-Image-Turbo (#99) via mflux (592k dl/mo, 6B, 9-step, Apache-2.0, pre-quantized MLX repo — the best small+hot+recent+MLX intersection). Orthogonal to the ASR (#68) and TTS-breadth (#69) tracks: new package, no shared code with tts/.

Motivation

The README's model matrix has four categories; three have landed code or open issues (TTS ✅, ASR #68, Omni via mlx-vlm). Diffusion (text/image → image) is the only category with nothing behind it. Meanwhile:

Upstream vllm-omni now ships a full diffusion subsystem — vllm_omni/diffusion/ (engine, per-step scheduling, step-wise continuous batching, quantization, LoRA), user-guide docs, and recipes for Qwen-Image (+Edit), Stable-Diffusion-3.5 ("OpenAI-compatible images API"), SDXL, Stable-Audio-Open, and video (Wan2.2, LTX, Cosmos).
Hub demand is real: SDXL-base 4.37M dl/mo, FLUX.1-schnell 637k, Qwen-Image 256k (HF 30-day, pulled 2026-10-06).
A dedicated, pip-installable MLX library exists for the quality tier: mflux (PyPI, v0.14.2, active) — FLUX port with 3/4/5/6/8-bit quantization, img2img, LoRA/fill.

This issue defines the serving seam so per-model work stays adapter-only — same doctrine as tts/ over mlx-audio, chat over mlx-lm.

Design decisions (make them here)

Library seam. Candidates (verification results, 2026-10-06):
- mlx-diffusers — SD 1.5/SDXL/SD3.5 breadth; not on PyPI (verified via PyPI + pip index) → git-only dependency, disfavored under pin discipline.
- Apple mlx-examples stable_diffusion — verified: SDXL is the default path, plus SD 1.x/2.x and sdxl-turbo (their README demo runs it on an 8 GB M1 mini); built-in --quantize (4/8-bit, on the fly); image2image.py; reads original diffusers-format repos directly, so no pre-converted mlx-community weights are required. Not pip-published → vendoring (precedented by mlx-audio internals).
- mflux (PyPI, v0.21.0, active — verified) — FLUX-only (+ newer FLUX models).
  Proposal (post hot-list refresh): mflux seam first — it turned out to be a 14-model multi-family library, not FLUX-only (Z-Image/Turbo, FLUX.2 klein 4B/9B + KV-cache edit variant, Qwen-Image 2.1 with reference-editing, Krea 2, FLUX.1, ERNIE, Lens, Ming, Boogu, FIBO, Ideogram 4, SeedVR2, Depth Pro), pip-installable and actively maintained → the engine + Image model: Z-Image-Turbo 6B (tier 1, engine reference path candidate) — 592k/mo, 9-step, Apache-2.0, mflux-native #99 Z-Image-Turbo land with zero vendoring, and Image model: FLUX.2-klein-4B (tier 1) — 394k/mo, 4-step, Apache-2.0, smallest of the category, edit-native #100/Image model: Qwen-Image-2.1 + fast variants (tier 1) — the currently-trending image family (3M+/mo wave), upstream flagship parity #101 follow on the same dependency. The SD seam (vendored mlx-examples, since mlx-diffusers is off PyPI) is the second PR for the evergreen SD tier (Image model: SDXL-base 4-bit (tier 1) — most-used image model on the Hub, CLIP-only memory tier #93/Image model: Stable Diffusion 1.5 4-bit (tier 1) — 2nd-hottest image model on the Hub, ~1 GB, small+hot pick #96/Image fast lane: low-step generation — SDXL-Turbo (943k/mo) vs SDXL-Lightning vs LCM-LoRA #97). Both dependencies live under the same [image] extra with independent pins.
Pin discipline. Cap the image dependency the way mlx-audio is capped (>=0.5.7,<0.6, see CONTRIBUTING.md): propose mflux>=<validated>,<next-minor>.
API surface. OpenAI-compatible POST /v1/images/generations on the Starlette server, mirroring upstream's images API: model, prompt, n, size, response_format (b64_json first), plus diffusion params steps, guidance, seed, negative_prompt. img2img (image + strength) rides the SDXL issue; an /v1/images/edits shape can follow.
Package layout. vllm_omni_mlx/diffusion/ seam (config + loader + service) mirroring tts/; zero native model code (architecture rule, docs/architecture.md). [image] optional extra in pyproject (like [tts]).
Residency model. Image models are big. Start solo-per-category (like TTS); record mixed-residency numbers wherever they fit (SDXL should; FLUX likely not — see its issue).

Metrics doctrine — the image analog of the TTS five (#39-style)

Time-to-image p50/p95 @ 1024×1024 — the headline (TTFA analog).
Per-step latency + stability — the step loop is our decode loop (same doctrine as the CFM/predictor work in feat(tts): compiled per-frame decode — fused predictor block, shapeless talker step, sync-free loop (#65) #67/VoxCPM2 batch-1 latency: vendored compiled per-patch loop (CFM 77% of wall) #79).
Time-to-first-preview — optional product lever: progressive JPEG every N steps. Stretch, not launch-blocking.
Output gate — deterministic-seed reproducibility + non-degenerate output stats in weight-gated tests; a CLIP-score floor only if calibrated (HNR-floor precedent: calibrate before hard-asserting).
Peak memory, solo + mixed — capacity constraint (16 GB M4: ~10.7 GB Metal working set; chat/TTS coexistence budget).

Work

Dependency decision + pin (SD-family library first; mflux pinned when Image model: FLUX.1-schnell 4-bit (tier 1) — mflux seam, 4-step fast lane, txt2img + img2img #92 lands); [image] extra in pyproject
diffusion/ seam: config, loader, service
POST /v1/images/generations (schema above), b64_json responses, OpenAI-compatible error shapes
Server wiring: model-id routing / CLI flag; README Getting Started + Diffusion-row update
Weight-gated e2e scaffold (HF_HUB_OFFLINE=1, donor-release rule, per-file subprocess; CI-safe stubs) + deterministic-seed reproducibility test
Baseline latency/memory numbers on Z-Image-Turbo 4-bit (first row of the doctrine above — delivered with Image model: Z-Image-Turbo 6B (tier 1, engine reference path candidate) — 592k/mo, 9-step, Apache-2.0, mflux-native #99)

Acceptance

/v1/images/generations returns a valid 1024² image from a weight-gated e2e run on Z-Image-Turbo 4-bit; same-seed calls reproduce byte-identical output (or deviations are characterized — fp16 fusion-flip precedent, #67); baseline time-to-image + peak-memory table in the PR; README Diffusion row points at real usage.

Verification results (2026-10-06), replacing the open questions in the library-seam decision:

mlx-examples stable_diffusion (via gh api, source of truth): SDXL is the default model family; SD 1.x/2.x supported; sdxl-turbo explicitly supported (README's sample image is sdxl-turbo on an 8 GB M1 Mac mini); txt2image.py --quantize for 4/8-bit on the fly; separate image2image.py. It reads original diffusers-format repos directly — stabilityai/stable-diffusion-xl-base-1.0 and stable-diffusion-v1-5/stable-diffusion-v1-5 both verified reachable — so the whole SD1.5/SDXL finetune ecosystem is servable with no pre-converted mlx-community weights needed. (My earlier guessed mlx-community 4-bit repo IDs 404 — irrelevant under this path; no LoRA script in the current example, so LoRA stays a stretch everywhere.)
mlx-diffusers: not on PyPI (PyPI JSON + pip index versions both checked) → git-only dep, disfavored under our pin discipline. The SD-family call effectively collapses to vendoring mlx-examples — mlx-audio-internals precedent.
mflux: PyPI v0.21.0 (search summaries said 0.14.x — stale; it's actively moving). Pinned when Image model: FLUX.1-schnell 4-bit (tier 1) — mflux seam, 4-step fast lane, txt2img + img2img #92 lands.

Net: every small+hot tier-1 model (#93 SDXL, #96 SD 1.5, #97 turbo/Lightning/LCM) rides one verified seam with no external conversion dependency.

Design record for the engine PR (from the arch/folder discussion):

Core decision: diffusion/ is a sibling service to tts/, not a Backend. The Backend protocol is chat-shaped (token Chunk iterators); image gen is one long blocking call. It follows the tts_service injection pattern exactly: constructed once, passed to create_app(..., image_service=...), own lock discipline.

POST /v1/images/generations
  → server.py (Route, schema ImageGenerationRequest in schemas.py, auth)
  → diffusion/config.py: load_image_service(model_ref)
      family registry: repo-id/org-prefix → {mflux | sd} (sniff like backends._peek_config;
      unknown id → 400 ApiError — instruct-400-on-0b6 precedent)
  → DiffusionService.generate(req) in a worker thread via the _bridge pattern
      (cancel at step boundary; one generation at a time — batch-1 doctrine)
      ├─ mflux_service.py  (pip mflux: Z-Image-Turbo, klein, Qwen-2.1, FLUX.1) — PR1 with #99
      └─ sd_service.py → vendored sd/ (mlx-examples stable_diffusion, LICENSE + SHA pin;
                          mlx-audio-internals precedent) — PR2 with #93
  → ImageResult {png, seed, timings{te_encode, steps[], vae_decode}, peak_mem} → b64_json

Folder design:

vllm_omni_mlx/diffusion/
├── __init__.py
├── config.py        # DiffusionConfig (quant bits, per-family steps/guidance/size presets)
│                    #   + FAMILY registry + load_image_service() factory
├── service.py       # DiffusionService protocol: generate() → ImageResult, name/model_type,
│                    #   per-model lock, donor-release, timings
├── mflux_service.py # PR1
├── sd_service.py    # PR2
└── sd/              # vendored mlx-examples stable_diffusion (LICENSE + pinned SHA)

tests/test_diffusion_config.py        # CI-safe: dispatch, schemas, defaults
tests/test_image_<model>.py           # weight-gated e2e, one file per model
scripts/image_probe.py                # latency_probe analog: time-to-image p50/p95, per-step, peak mem
docs/images.md                        # post-landing (speech.md precedent)

Details:

pyproject: add vllm_omni_mlx.diffusion to the explicit [tool.setuptools] packages list (currently only root + tts — silently breaks the wheel if missed) + [image] extra (mflux pinned at PR1; vendored SD adds no dep).
Per-family defaults in DiffusionConfig: Z-Image-Turbo 9 steps/guidance-off; klein-4B 4; Qwen-2.1 40 guidance-free; SDXL 30/g5 (+1024², 1344×768, 768×1344 presets); SD1.5 512²-native.
The mflux step callback is the product seam — timings, the future SSE preview (time-to-first-preview), and cancellation all hang off it.
Seed plumbed through both seams (deterministic-reproducibility gate).

Test-layout convention (maintainer preference): tests mirror src with second-level folders — landing repo-wide as #106 (tests/tts/ etc., unittest discover -s tests -t .). The diffusion work follows from day one: tests/diffusion/ (test_config, test_service CI-safe at root of it; one weight-gated file per model: test_mflux_zimage.py, test_image_sdxl.py, …) — replaces the flat tests/test_diffusion_config.py sketch in the design comment above.

Engine + reference path landed via #108 (stacked on #106): diffusion/ seam, /v1/images/generations, --image-model, [image] extra (mflux pinned), tests/diffusion/ mirroring src, docs pattern section, baseline doctrine table on #99. Notable engineering find: MLX stream-affinity bug on the serve path (load on main thread + generate on worker → 'no Stream(cpu, 0)') — fixed by a dedicated single worker thread owning load+generate; caught by the mandatory serve smoke, missed by main-thread unit tests (#71 lesson holding). Open follow-ups on this issue's list: SD seam PR2 (vendored mlx-examples, #93), SSE preview via mflux step callbacks, /v1/images/edits shape.

chore: tests mirror src layout — tests/tts/ second level (#106)

Navigation Menu

Repository navigation

Image generation engine: /v1/images/generations on MLX — mflux seam, [image] extra, time-to-image doctrine #91

Motivation

Design decisions (make them here)

Metrics doctrine — the image analog of the TTS five (#39-style)

Work

Acceptance

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Image generation engine: /v1/images/generations on MLX — mflux seam, [image] extra, time-to-image doctrine #91

Description

Motivation

Design decisions (make them here)

Metrics doctrine — the image analog of the TTS five (#39-style)

Work

Acceptance

Activity

hsliuustc0106 commented on Oct 6, 2026

hsliuustc0106 commented on Oct 7, 2026

hsliuustc0106 commented on Oct 7, 2026

hsliuustc0106 commented on Oct 7, 2026

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions