feat: diffusion engine + Z-Image-Turbo on /v1/images/generations (#91, #99) by hsliuustc0106 · Pull Request #108 · InferMatrixLab/vllm-omni-mlx

hsliuustc0106 · 2026-10-07T01:37:33Z

Lands the diffusion engine (#91) with its reference model Z-Image-Turbo (#99): vllm_omni_mlx/diffusion/ (family registry → ResolvedModel weight routing → mflux seam), OpenAI-compatible POST /v1/images/generations, serve --image-model, the [image] extra (mflux pinned >=0.21.0,<0.22), and tests/diffusion/ mirroring src (#106's convention).

Stacked on #106 (tests-mirror layout + -t . discovery) — merge that first; rebase is trivial if it lands via squash.

How it serves

vllm-omni-mlx serve --image-model z-image-turbo        # needs [image]
curl -s localhost:8000/v1/images/generations -d '{"prompt":"a red apple","size":"512x512","steps":9,"seed":42}'
# → {"data":[{"b64_json":...}], "seeds":[42], "generation_time_ms":[...]}

Family registry: z-image-turbo alias → pre-quantized 4-bit mirror (filipstrand/Z-Image-Turbo-mflux-4bit, 5.5 GiB — the canonical fp16 repo is ~12.3 GiB and doesn't fit the dev disk); Tongyi-MAI/Z-Image-Turbo → canonical repo with on-load 4-bit. Unknown refs → 400 listing what's supported.
OpenAI surface (prompt/n/size/b64_json) + steps/guidance/seed; guidance is rejected 400 on this family (guidance-distilled — mflux forces it to 0; we fail early instead of ignoring). n generates sequentially with seeds seed+i.
Response extensions: seeds + generation_time_ms (the doctrine's metrics, machine-readable).
One bug found by the serve smoke (the TTS model: VoxCPM2 (tier 1) — presets + cloning, 48 kHz, 4-bit 2.3 GB #71 lesson working as intended): load on the uvicorn main thread + generate on a worker thread trips MLX's per-thread stream context (There is no Stream(cpu, 0) in current thread). Fix: MFluxService owns a single dedicated worker thread for load and generate — which also serializes requests (batch-1) for free. Main-thread unit tests never saw it; the serve smoke did.

Measurements (M4 16 GB, 4-bit, offline cache, 2026-10-07)

workload	time-to-image	peak memory
model load (mmap)	1.8 s	—
512² × 9 steps, solo	64.4 s	6.42 GiB
512² × 4 steps, solo (serve API)	30.3 s	~6.4 GiB
1024² × 9 steps, solo	387.8 s	8.79 GiB
512² × 9 steps, mixed w/ 0.5B chat	95.4 s — chat replies OK	6.68 GiB
determinism	same seed → byte-identical PNG	—

Per-step ≈ 6.7–7.5 s at 512² — the profile conclusion: compute-bound on the base M4 (6B × ~1.5k-token context ≈ 18 TFLOP/step ÷ ~2.8 effective TFLOPS ≈ 6.5 s — matches), not a bug and not the weight-stream bound #99 guessed. So Z-Image-Turbo is the quality tier; klein-4B (#100, 4 steps, 1.7 GB/step) is the latency hope. Everything fits solo under the 10.7 GiB working set; mixed with 0.5B chat verified.

Tests

tests/diffusion/test_config.py — resolution + parsing, CI-safe (19 cases incl. endpoint suite w/ fake service; no mflux import needed)
tests/diffusion/test_images_endpoint.py — route: auth, 400s, b64/seeds shape, model listing, absent-without-service
tests/diffusion/test_mflux_zimage.py — weight-gated e2e: PNG gate (magic + size floor), dims/steps/seed honored, same-seed byte-identity, guidance 400; donor release in tearDownClass

Docs: README Diffusion rows + [image] footprint row (honest torch note: +870 MB, 583 of it torch — the one extra that pulls it); docs/adding-a-model.md gains the diffusion pattern section. Audition artifact: /tmp/zimage_serve.png (512², "a tiny wooden sailing boat").

Follow-ups (not this PR): img2img (image/strength — mflux takes a path, needs temp-file plumbing), klein-4B on the same seam (#100), per-step timings via mflux callbacks → SSE preview (the time-to-first-preview lever).

…99) - vllm_omni_mlx/diffusion/: config (family registry + ResolvedModel weight routing), service protocol (ImageResult), mflux seam (MFluxService) - POST /v1/images/generations: OpenAI-compatible (prompt/n/size) + steps/ guidance/seed; b64_json + seeds/generation_time_ms extensions - serve --image-model flag; [image] extra (mflux>=0.21,<0.22); torch note - tests/diffusion/ mirrors src: CI-safe config+endpoint (fake service), weight-gated Z-Image-Turbo e2e (512x512/9 steps 64.4s, peak 6.42 GiB, same-seed byte-identical PNGs)

hsliuustc0106 · 2026-10-07T22:41:12Z

+    @classmethod
+    def load(cls, resolved: ResolvedModel) -> "MFluxService":
+        from mflux.models.common.config.model_config import ModelConfig
+        from mflux.models.z_image import ZImage


[P2] Missing [image] extra gives a raw traceback instead of the repo's install hint. With pip install -e . (CI's own install, and the README default), serve --image-model z-image-turbo raises ModuleNotFoundError, which is not in _serve's except (RuntimeError, ValueError, OSError) → ugly traceback. This PR's own tree establishes the pattern three times (backends.py:237 for [omni], tts/config.py:52 for [tts], main's #104 for [asr]; #109 also converts it). Fix: wrap these imports in try/except ImportError: raise RuntimeError("image generation needs mflux; install it with: pip install 'vllm-omni-mlx[image]'") from exc.

hsliuustc0106 · 2026-10-07T22:41:13Z

+        height=height,
+        steps=_opt_int(payload.get("steps"), "steps"),
+        guidance=_opt_float(payload.get("guidance"), "guidance"),
+        seed=_opt_int(payload.get("seed"), "seed"),


[P3] Unbounded seed yields an opaque 500, not a 400. _opt_int accepts any int and generate does no bounds check; mx.random.key(-1) and key(2**64) both raise TypeError in mlx 0.32.3 (verified; 2**31/2**63 accepted), and mflux calls it per generation → escapes as the route's generic 500 for a client-input problem. seed + i for n>1 can also overflow the accepted range. Fix: reject seed < 0 or seed >= 2**32 in parse_image_generation or generate with a 400.

hsliuustc0106 · 2026-10-07T22:41:15Z

            return _openai_error(exc)
        return JSONResponse({"object": "list", "voices": tts_service.voices})

+    async def images_generations(request: Request) -> Response:


[P3] response_format is silently ignored. An OpenAI client requesting response_format: "url" silently receives b64_json. The sibling /v1/audio/speech route validates the field (server.py:332), and #109's images route does too (response_format must be 'b64_json'…). Fix: accept only None/"b64_json" with an explicit 400.

hsliuustc0106 · 2026-10-07T22:41:16Z

+
+
+@unittest.skipUnless(_snapshot_cached(), f"{DEFAULT_MODEL} not in the local HF cache")
+class ZImageTurboE2ETest(unittest.TestCase):


[P3] Weight-gated class skips the ReleaseAfterClass taxonomy. Every weight-gated class in the repo inherits tests/_teardown.ReleaseAfterClass (16 TTS classes in this tree; tests/test_asr.py:286 on main). The hand-rolled tearDownClass is functionally equivalent today but diverges from the documented mechanism and won't pick up future ReleaseAfterClass additions. Fix: class ZImageTurboE2ETest(ReleaseAfterClass, unittest.TestCase): and drop the custom tearDownClass.

hsliuustc0106 · 2026-10-07T22:41:18Z

+            model_path=resolved.weights_ref,
+        ).result()
+        service.bits = getattr(service._model, "bits", resolved.quantize)
+        print(


[P3] Load diagnostics print to stdout, not stderr. All sibling load/serve diagnostics go to stderr (__main__.py:165, :177, :234, main's _load_asr); stdout pollutes any JSON-piped consumer of the server process. Fix: print(..., file=sys.stderr, flush=True).

hsliuustc0106 · 2026-10-07T22:41:19Z

+    )
+
+
+@dataclass(frozen=True)


[P3] DiffusionConfig is exported dead code. Its only references are the definition and the __init__.py re-export; nothing in server.py, __main__.py, or tests consumes it, and its width/height = 1024 defaults duplicate schemas.ImageGenerationRequest's. Fix: drop it until a consumer exists (or wire it into the CLI).

hsliuustc0106

Delegated review @ d49375b — verdict: request-changes (6 findings posted inline + 2 branch-level blockers below). The engine code itself is in good shape; the blockers are branch topology and seam coordination.

Blocker 1 — stale base branch, retarget + rebase required. The base chore/tests-mirror-src is fully superseded by main: merge-base is 7ddd019, and the base's distinguishing commit (28fb328) is content-identical under tests/+CONTRIBUTING.md to main's 0e8f4e3 (#106). Main has since gained #104 (asr), #106, #107, #89. Conflicts on rebase are guaranteed in: server.py (create_app signature vs #104's asr_service — the merged version must carry both asr_service and image_service; /v1/models listing), __main__.py:116 (the no-model guard line now includes asr_model), pyproject.toml ([project.optional-dependencies] and packages=[...] — main inserted the asr lines in the same spots), and README.md (#89 rewrote the install block this appends to). Good news: this PR's own tests already use the post-#106 tests/<family>/ layout (contrast #109, which adds flat test files).

Blocker 2 — add/add seam collision with #109. Both PRs create vllm_omni_mlx/diffusion/{__init__,config,service}.py with different contents, byte-identical pyproject.toml lines (image = ["mflux>=0.21.0,<0.22"], vllm_omni_mlx.diffusion package), the same --image-model flag, route, handler name, and create_app(..., image_service=None) — every one conflicts; they cannot both merge without reconciliation. Of the two skeletons, this one is the better shared base: the Family/FAMILIES/resolve_model registry + DiffusionService Protocol is multi-family by design (docs/adding-a-model.md here codifies "add a Family…"), it pins load+generate to a single executor thread (the #71 VoxCPM2 stream-affinity lesson), and it parses requests in the schema layer like the other endpoints. #109 brings the richer per-model surface (guidance/negative_prompt/scheduler/lora, license print at load, one-shot CLI, bench, examples, docs/image.md) and its ImageConfig hard-rejects any family but qwen-image-2.1. Recommended path: adopt this PR's skeleton, fold #109's Qwen family into the registry, port its missing request fields (negative_prompt, lora knobs) into the protocol.

Verified against installed mflux 0.21.0 (read-only): ZImage.__init__ / generate_image kwargs all valid; guidance is forced to 0.0 by mflux when the family doesn't support it (our 400 is stricter, fine); ModelConfig.z_image_turbo() exists; WeightApplier.apply_and_quantize honors stored quant levels and sets model.bits; GeneratedImage.image is PIL; mflux's own default steps for z-image-turbo is 9 (matches family.steps = 9). mx.random.key bounds probed (grounds the seed finding). CI safety structural (no ungated mflux/torch imports; e2e skip gate is a pure glob). No MLX-trap patterns in the diffusion code.

Questions: (1) the default weights are a third-party mirror (filipstrand/Z-Image-Turbo-mflux-4bit), not the canonical Apache-2.0 Tongyi-MAI/Z-Image-Turbo — state the mirror's license at load or in the README, per #109's precedent? (2) peak_memory_gib is computed per image but not surfaced in the HTTP response — intentional? (3) steps cap 50 + n≤4 at 2048² is minutes of GPU per request on a 24 GB host — is that the whole latency story?

Not executed (weight-gated): tests, loads, generations — evidence reviewed only.

This was referenced Oct 7, 2026

Image model: Z-Image-Turbo 6B (tier 1, engine reference path candidate) — 592k/mo, 9-step, Apache-2.0, mflux-native #99

Open

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

Open

hsliuustc0106 commented Oct 7, 2026

View reviewed changes

hsliuustc0106 mentioned this pull request Oct 10, 2026

docs: README catch-up — one model matrix, ASR row served, [asr] install + news #114

Merged

Navigation Menu

Repository navigation

feat: diffusion engine + Z-Image-Turbo on /v1/images/generations (#91, #99) - #108

feat: diffusion engine + Z-Image-Turbo on /v1/images/generations (#91, #99)#108
hsliuustc0106 wants to merge 1 commit into
chore/tests-mirror-srcfrom
feat/diffusion-engine-zimage

hsliuustc0106 commented Oct 7, 2026

Uh oh!

hsliuustc0106 Oct 7, 2026

Uh oh!

hsliuustc0106 Oct 7, 2026

Uh oh!

hsliuustc0106 Oct 7, 2026

Uh oh!

hsliuustc0106 Oct 7, 2026

Uh oh!

hsliuustc0106 Oct 7, 2026

Uh oh!

hsliuustc0106 Oct 7, 2026

Uh oh!

hsliuustc0106 left a comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant



		@unittest.skipUnless(_snapshot_cached(), f"{DEFAULT_MODEL} not in the local HF cache")
		class ZImageTurboE2ETest(unittest.TestCase):

		)


		@dataclass(frozen=True)

Conversation

hsliuustc0106 commented Oct 7, 2026

How it serves

Measurements (M4 16 GB, 4-bit, offline cache, 2026-10-07)

Tests

Uh oh!

hsliuustc0106 Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 Oct 7, 2026

Choose a reason for hiding this comment

Uh oh!

hsliuustc0106 left a comment

Choose a reason for hiding this comment

Uh oh!

Reviewers

Uh oh!

Assignees

Labels

Projects

Milestone

Development

Uh oh!

1 participant