Repository navigation
feat(image): mflux seam + Qwen-Image-2.1 serving — /v1/images/generations, Viggle turbo fast lane (#91, #101) - #109
hsliuustc0106 wants to merge 5 commits into
Conversation
…ions, Viggle turbo fast lane (#91, #101) The diffusion seam from #91 lands here with its first mflux backend (that seam had not landed yet); #91's Z-Image-Turbo reference-path acceptance stays open for #99. - diffusion/ seam: config (family registry + mflux loader, license surfaced) + service (validation, single-user lock, prompt → PNG bytes). Zero native model code; mflux pinned >=0.21,<0.22 under the new [image] extra. - POST /v1/images/generations: OpenAI images shape, b64_json; size/n plus steps/guidance/seed/negative_prompt; multiples-of-16 enforced (no silent rounding); 'url' response_format rejected. - CLI: serve <image-repo> (layout-marker sniffing), --image-model, --image-lora, --image-scheduler, one-shot 'vllm-omni-mlx image'. - Verdicts recorded on #101: license non-commercial (Qwen Research License, all paths incl. Viggle/Pruna); fast variant = Viggle turbo v0.3 LoRA via the built-in viggle_turbo scheduler (6 steps enforced); Pruna not selectable (merged transformers, no mflux loader path). - Eval (M4 24 GB, solo, 4-bit, fixed seeds): base 40-step 512²→1024² 190→852 s tti (peak MLX 13.6→23.3 GiB); turbo 6-step 31 s @512², 134 s @1024² (6.3× vs base; peak 24.6 GiB). scripts/bench_image.py reproduces. - Tests: 13 CI-safe unit tests (validation/seed/lock/HTTP shapes) + weight-gated e2e (byte-identical same-seed repro, output stats gates, IMAGE_E2E_FULL=1-gated 1024²/40-step acceptance cell). Serve-path e2e: boot + client round-trip + 400/401 shapes verified.
|
Serve-path e2e (weight-gated, on the M4 box):
One process note for reviewers running the suite: the 1024²/40-step acceptance cell is gated behind Also, per the maintainer's 1328² concern: agreed and recorded in the tables — 1328² was not measured because the 1024² fast lane already peaks at 24.6 GiB MLX-active here; any 1328² path needs >24 GB. Extrapolating per-token puts base-40-step at roughly 29–30 GiB, so it would not have completed on this box. |
…h; on-the-fly quantize covers the DiT only Measured on the M4 (512², 8 steps, same seed): dense + --quantize 8 → 22.4 GiB resident, ~18 s/step, 24.6 GiB peak (visibly cleaner output but not servable); dense + --quantize 4 → 19.1 GiB, ~8 s/step. Root cause: mflux's on-the-fly quantize predicate covers only the transformer — the 8B text encoder stays bf16 (14.1 GiB). Follow-up lever: hybrid 4-bit DiT+TE + bf16 VAE repo.
|
Quantization A/B added (docs/image.md updated, commit on the branch):
So: the served default is the quantization answer — 4-bit everywhere, and the pre-quantized repo is the only path that both fits (the on-the-fly predicate leaves the 8B text encoder bf16 at 14.1 GiB) and is fastest. The 8-bit quality win is real but the actionable lever is a hybrid 4-bit DiT+TE + bf16 VAE repo (+0.9 GiB) plus an upstream mflux quantize-predicate fix — filed as follow-up, not in this PR. |
…E decode scales with pixels, prompt encode ~0.3s (cached)
…t4 affine DiT; nvfp4 matches int4 speed and near-int8 flat gradients (mflux monkeypatch experiment)
|
fp4 quantization research (docs/image.md updated): where does the faint crosshatch come from, and would mxfp4/nvfp4 fix it? Root cause ablation (same seed/prompt, DiT-only quantization, TE/VAE bf16): the artifact is the int4 affine DiT, not the 4-bit VAE — int4-DiT + bf16-VAE keeps the crosshatch, int8-DiT removes it. Format ablation (MLX 0.32
Conclusion: nvfp4 is the same size/speed class as the served int4 with visibly better flat gradients — the right target format for the next community quant repo. Blocker is upstream: mflux 0.21 passes no |
| if lora_paths: | ||
| config = replace(config, lora_paths=tuple(lora_paths)) | ||
| if scheduler: | ||
| config = config.with_overrides(scheduler=scheduler) |
There was a problem hiding this comment.
[P1] The Viggle turbo fast lane rejects every default request — both documented invocations are broken. _load_image_service overrides only scheduler; ImageConfig.steps keeps its default 40 (config.py:35), and the service gate rejects any steps != 6 under viggle_turbo. So the documented serve command (README + docs/image.md, ending --image-scheduler viggle_turbo) builds a server that 400s every request unless the client explicitly sends steps: 6 — and there is no --image-steps serve flag to fix it server-side. The one-shot path fails the same way: image --scheduler viggle_turbo without --steps 6 hits the same gate → ImageError → exit 1 (the --steps default is None, which with_overrides drops). scripts/bench_image.py is the only caller that sets steps=VIGGLE_TURBO_STEPS, which is why the bench worked while the shipped CLI paths don't. Fix: when scheduler == viggle_turbo, default steps to VIGGLE_TURBO_STEPS unless explicitly overridden (serve and one-shot paths), or add --image-steps.
|
|
||
| from .config import VIGGLE_TURBO_SCHEDULER, VIGGLE_TURBO_STEPS, ImageConfig | ||
|
|
||
| MIN_SIDE, MAX_SIDE = 256, 2048 |
There was a problem hiding this comment.
[P2] MAX_SIDE = 2048 admits sizes the PR itself documents as unservable. The docs/image.md tables show 1024² turbo peaking 24.6 GiB on the 24 GB host and state any 1328² path wants more than 24 GB — yet validation accepts sides up to 2048 (4× the pixels of 1328², certain swap/OOM), and even the default 1024x1024 sits at 23.3 GiB peak on exactly that class of box. Oversized requests are documented, not guarded. Fix: cap MAX_SIDE at the measured envelope (e.g. 1328) or make it a config knob tied to the box class.
| cls.model = load_image_model(cls.config) | ||
| cls.service = ImageService(cls.model, cls.config) | ||
|
|
||
| def test_same_seed_is_byte_identical(self): |
There was a problem hiding this comment.
[P2] No worker-thread coverage with the real engine. The repo's serving bar requires it — main-thread tests structurally miss thread-binding bugs. test_image_service.py exercises the HTTP→asyncio.to_thread path only with FakeMfluxModel, and this e2e drives service.generate on the main thread. mflux builds mx.compile closures lazily on first generate and mutates shared state (prompt_cache, _text_caches) — serialized by the service lock, but never exercised on a serving worker thread with the real engine. Fix: add a gated serve smoke here (TestClient over create_app(image_service=…) or asyncio.to_thread(service.generate, …)).
| if not isinstance(n, int) or not 1 <= n <= 4: | ||
| raise ImageError("n must be an integer within [1, 4]") | ||
| if seed is not None: | ||
| seed = self._int(seed, "seed") |
There was a problem hiding this comment.
[P3] Negative seed yields a 500, not a 400. _int checks only isinstance(value, int); mx.random.key(-1) raises TypeError (verified against installed mlx 0.32.3), which escapes as the route's generic 500 instead of a validation error. Fix: reject seed < 0 (or bound to [0, 2**32)) alongside the int check.
| seed_i = seed if seed is not None else random.randint(0, 2**31 - 1) | ||
| seed_i += i | ||
| started = time.perf_counter() | ||
| generated = self._model.generate_image( |
There was a problem hiding this comment.
[P3] mflux's prompt-embedding cache grows unbounded in a long-running server. mflux 0.21 writes every unique prompt's embeddings into model.prompt_cache and never evicts (verified in the installed 0.21.0). The per-run finally clears transformer._text_caches but not prompt_cache, so each unique prompt permanently adds resident memory on a box the docs say peaks at 23.3–24.6 of 24 GiB. Fix: cap or periodically clear prompt_cache from the seam (docs/image.md already treats the prompt cache as a known lever).
| help="also serve image generation on /v1/images/generations (mflux via the [image] extra); the only model when <model> is omitted", | ||
| ) | ||
| parser.add_argument( | ||
| "--image-lora", |
There was a problem hiding this comment.
[P3] --image-lora / --image-scheduler are silently dropped when no image model is served. They are consumed only inside _load_image_service, so serve --tts-model X --image-scheduler viggle_turbo starts a server with the flags ignored. Fix: error or warn in _serve when these flags are set without an image model being loaded.
| encode per new prompt (free on repeats — mflux caches the text-prefix KV | ||
| in-process), <1% noise init + PNG encode. One-time per process: 7–8 s load, | ||
| plus first-generation kernel compile per resolution. | ||
| - **Prompt encoding is per-request** (~2 s at load-warm state): the text |
There was a problem hiding this comment.
[P3] This bullet contradicts the first bullet (and the code). mflux 0.21 already caches exact-prompt embeddings (model.prompt_cache), so repeats are free — bullet 1 correctly says ~0.3 s and free on repeats, while this bullet claims per-request ~2 s and proposes the cache as a follow-up. Delete or rewrite to match.
|
|
||
| - **Solo-only residency.** Weights are ~8.9 GiB resident and generation peaks | ||
| ~14–17 GiB MLX-active; on a 16 GB Mac serve this model alone and keep chat | ||
| out of the same process. The 24 GB class runs it comfortably. |
There was a problem hiding this comment.
[P3] "runs it comfortably" contradicts this page's own tables. The tables report 23.3 GiB peak (base 1024²) and 24.6 GiB (turbo 1024²), and the 1328² note says the base path at ≥1024² and any path at 1328² want more than 24 GB; the design-note figures (~14–17 GiB) match only the ≤768² cells. Reconcile the design note with the tables.
|
|
||
| # default cells: "SIZE:runs" — three runs up to 1024²-side, two for the | ||
| # multi-minute big cells (min/median/max still meaningful at n=2/3) | ||
| BASE_CELLS = "512x512:3,768x768:3,1024x1024:3,1280x720:2,1328x1328:2" |
There was a problem hiding this comment.
[P3] The documented repro does not reproduce the published tables. The default preset includes 1328x1328:2, which the docs explicitly say was not measured and wants more than 24 GB (the default preset would attempt unservable cells on the measurement machine), and the table run counts (768²=1, 1280×720=1, turbo cells=1) don't match the preset defaults (3/2/3). Fix: align preset defaults with the published cells or record the --cells/run overrides actually used.
hsliuustc0106
left a comment
hsliuustc0106
left a comment
There was a problem hiding this comment.
Delegated review @ ee177496 — verdict: request-changes (9 findings posted inline; the P1 blocks).
The P1: the Viggle turbo fast lane — the headline feature — rejects every default request. Both documented invocations (serve --image-scheduler viggle_turbo, one-shot image --scheduler viggle_turbo) leave steps at the config default 40, and the service gate requires exactly 6 under that scheduler; only bench_image.py sets it. Plus a P2 size cap (2048 admits documented-unservable requests) and a P2 missing worker-thread coverage for the real engine.
Verified against the installed stack (mflux 0.21.0 / mlx 0.32.3, read-only):
- every
QwenImage21/generate_imagekwarg used by the service is valid in 0.21.0;viggle_turbois a registered scheduler with fixed 6-step sigma nodes (its own loader raises if steps ≠ 6, mirroring our 400 gate); - license print at load present (3 sites); CI gating is structural (
local_snapshotgate precedes the lazy mflux import; no ungated mflux/torch imports in CI-collected tests); - text-prefix KV caching and the exact-prompt
prompt_cacheexist as claimed (grounds the unbounded-growth P3); mx.random.key(-1)raises TypeError (grounds the seed P3).
Not executed (weight-gated): all eval/battery/e2e runs — evidence reviewed only.
Questions (could not verify offline):
- Thread affinity — generation runs on
asyncio.to_threadpool threads against a model loaded on the main thread, with lazily-builtmx.compileclosures. The lock serializes access, but was any serve-path generation (HTTP round-trip, not CLI/bench) actually run? _looks_like_imageroutes any repo withtransformer/config.jsonand no_class_nameto the image path — known text/vlm repos that would be misrouted?turbo_lora_path()'s bare-repo-id fallback (Viggle/Qwen-Image-2.1-viggle-turbowithout:filename) — does mflux resolve that as a single adapter or a full-repo download?
What
Closes #101 — Qwen-Image-2.1 served on the OpenAI images API, and with it the mflux
diffusion seam from #91 (that seam had not landed yet, so this PR lays it down together with
the first mflux backend; #91's Z-Image-Turbo reference-path acceptance remains open for #99).
vllm_omni_mlx/diffusion/— the seam (Image generation engine: /v1/images/generations on MLX — mflux seam, [image] extra, time-to-image doctrine #91 design):config.py(model registry + loader;license surfaced per family) and
service.py(validation + single-user lock, prompt → PNGbytes). Zero native model code — the architecture rule holds; model math is mflux's,
pinned
>=0.21.0,<0.22under a new[image]extra.POST /v1/images/generationson the Starlette server: OpenAI images shape (b64_json),OpenAI
size/nplus diffusion extrassteps,guidance,seed,negative_prompt.sizesides must be multiples of 16 — rejected with a 400 rather than silently rounded.serve <image-repo>(checkpoint sniffing via diffusers/mflux layout markers),serve --image-model <repo>(solo image serving),--image-lora/--image-scheduler(the fast lane), and a one-shot
vllm-omni-mlx image --prompt … --out out.png.scripts/bench_image.py— the Image generation engine: /v1/images/generations on MLX — mflux seam, [image] extra, time-to-image doctrine #91 metrics doctrine for images: time-to-image, per-stepcadence (mflux in-loop callback), MLX peak/active/cache + process RSS, per-resolution
cells, cool-down sleeps; JSON + markdown output. Repro for every number below.
Verdicts (recorded first, per the issue's acceptance)
License — non-commercial, all paths.
Qwen/Qwen-Image-2.1ships the Qwen ResearchLicense (2026-09-20): use, reproduction, and derivatives for non-commercial
research/evaluation only; commercial use needs a separate Qwen license. The fast variants
inherit it (Viggle turbo LICENSE and Pruna LICENSE are the same agreement). Served with the
license printed at load and flagged in README/docs. The Apache-2.0 tier (#99/#100) is the
commercial-safe path on the same seam.
Fast variant — Viggle turbo wins. mflux carries first-class support: a DMD-distilled
LoRA over the base DiT + a built-in
viggle_turboscheduler sampling the LoRA's 6 trainedsigma nodes (6 steps, no CFG — enforced in our service validation). A/B at 512²/seed 1000:
v0.3 LoRA shows less ghosting than v0.2.1 → v0.3 is the default fast lane. Pruna ships
merged dense transformers (
p_qwen_image_2.1_{5,8}step_v0.1.safetensors) with no mfluxloader path — recorded as follow-up, not selected. Base-40-step remains the quality
reference and the serving default.
Evaluation (M4 24 GB, solo, 4-bit mflux-q4, guidance-free, fixed seeds, cool-down sleeps)
Weights resident at load: 8.93 GiB MLX-active (4-bit DiT + 4-bit TE + VAE), load 7–8 s
from page cache.
ru_maxrsshigh-water stays ~8.9–9.5 GiB.Base path — 40 steps, guidance-free (time-to-image; min/median/max of N runs)
Fast lane — Viggle turbo v0.3 LoRA, 6 steps
Reading: latency scales with latent tokens ((H/16)·(W/16)) — ~4× pixels ≈ 4× per-step
(compute-bound at 4-bit). The turbo lane is 6.3× faster end-to-end at 1024² (134 s vs
852 s) with visible distillation ghosting; the base path is the quality reference. The
issue's "~10–30 s per image" expectation for the official base holds only on
Max/Ultra-class bandwidth — on a base M4 the 40-step path is minutes, which is exactly why
the fast lane matters as the servable default. Weights resident: 8.93 GiB (base), 12.2 GiB
(with the baked LoRA). 1328² not measured: the 1024² fast lane already peaks at 24.6 GiB
MLX-active on this 24 GB machine — ≥1024² base and any 1328² path want >24 GB.
Gates
tests/test_image_e2e.py, 3/3 green locally): same-seedbyte-identical output; non-degenerate PNG stats (stddev band, channel means, size
floor); the issue's acceptance cell — a valid 1024² image from the default 40-step path.
tests/test_image_service.py, 13 tests): size/steps/guidance/n/seedvalidation, seed incrementing for
n>1, theviggle_turbosteps=6 gate, lock behavior,HTTP surface (OpenAI shape, 400 mapping,
response_format: urlrejected, model listing).checkpoints are absent, unchanged).
vllm-omni-mlx serve+ client round-trip (see comment below).Footprint note
[image]is the first extra that drags torch (arm64 build, a hard mflux dependency)plus matplotlib/opencv: 61 resolved packages, ~1.3 GB installed, ~750 MB over
[tts].README's footprint table gains the row with an explicit "keep it out of installs that
don't serve images".
Follow-ups (not in this PR)
Reference editing +
/v1/images/edits(mflux prefix-KV — second edit-native model after#100) · Z-Image-Turbo (#99) and FLUX.2-klein (#100) as Apache-2.0 entries on the same seam ·
prompt-embedding cache (the TTS prefix-cache analog; prompt encode costs ~2 s/request) ·
step-cache (
step_cache_ratio) as a latency lever · Pruna merged-transformer loader.