From d49375bfe7c0e6c4a291ba5ea52d32724083098b Mon Sep 17 00:00:00 2001 From: hsliu_ustc Date: Wed, 7 Oct 2026 09:16:35 +0800 Subject: [PATCH] feat: diffusion engine + Z-Image-Turbo on /v1/images/generations (#91/#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) --- .zcodeignore | 38 ++++++++ README.md | 15 ++- docs/adding-a-model.md | 27 ++++++ pyproject.toml | 5 +- tests/diffusion/__init__.py | 0 tests/diffusion/test_config.py | 79 +++++++++++++++ tests/diffusion/test_images_endpoint.py | 97 +++++++++++++++++++ tests/diffusion/test_mflux_zimage.py | 77 +++++++++++++++ vllm_omni_mlx/__main__.py | 43 ++++++++- vllm_omni_mlx/diffusion/__init__.py | 18 ++++ vllm_omni_mlx/diffusion/config.py | 99 +++++++++++++++++++ vllm_omni_mlx/diffusion/mflux_service.py | 116 +++++++++++++++++++++++ vllm_omni_mlx/diffusion/service.py | 59 ++++++++++++ vllm_omni_mlx/schemas.py | 53 +++++++++++ vllm_omni_mlx/server.py | 55 ++++++++++- 15 files changed, 769 insertions(+), 12 deletions(-) create mode 100644 .zcodeignore create mode 100644 tests/diffusion/__init__.py create mode 100644 tests/diffusion/test_config.py create mode 100644 tests/diffusion/test_images_endpoint.py create mode 100644 tests/diffusion/test_mflux_zimage.py create mode 100644 vllm_omni_mlx/diffusion/__init__.py create mode 100644 vllm_omni_mlx/diffusion/config.py create mode 100644 vllm_omni_mlx/diffusion/mflux_service.py create mode 100644 vllm_omni_mlx/diffusion/service.py diff --git a/.zcodeignore b/.zcodeignore new file mode 100644 index 0000000..7131520 --- /dev/null +++ b/.zcodeignore @@ -0,0 +1,38 @@ +.venv/ +__pycache__/ +*.pyc +*.egg-info/ +dist/ +build/ +.DS_Store +.env + +# ===== ↑ 以上同步自 .gitignore（「从 .gitignore 同步」只重写以上部分）===== +.git/ +.hg/ +.svn/ +node_modules/ +bower_components/ +jspm_packages/ +site-packages/ +venv/ +coverage/ +htmlcov/ +lcov-report/ +cmakefiles/ +cmake-build-*/ +bazel-*/ +pods/ +deriveddata/ +storybook-static/ +playwright-report/ +test-results/ +allure-results/ +allure-report/ +cdk.out/ +*.dist-info/ +eggs/ +pip-wheel-metadata/ +wheels/ +# ----- ↑ 以上为 ZCode 默认排除规则（自定义规则请写在本行下方，不会被同步/恢复改动）----- +# 自定义规则写在下方（本行提示可删除） diff --git a/README.md b/README.md index 807fa9c..5538971 100644 --- a/README.md +++ b/README.md @@ -55,14 +55,14 @@ vllm-omni-mlx is flexible and easy to use with: | **TTS** — text → speech | Qwen3-TTS (CustomVoice · Base · VoiceDesign, 0.6B/1.7B) | ✅ verified end-to-end, buffered + streaming — [speech guide](docs/speech.md) | | **Omni** — any-to-any chat | Qwen3-Omni 30B-A3B | 🚧 chat works via mlx-vlm; speech-out chat in progress | | **ASR** — speech → text | Whisper, Qwen3-ASR, Voxtral, … | 🚧 planned — `/v1/audio/transcriptions` (#68) | -| **Diffusion** — text/image → image | — | 🚧 planned — roadmap (#2) | +| **Diffusion** — text/image → image | Z-Image-Turbo (6B, 9-step distilled, 4-bit) | ✅ `/v1/images/generations` (mflux seam) — SD tier next (#95) | | Modality | Examples | Status | | --- | --- | --- | | **TTS** — text → speech | Qwen3-TTS (CustomVoice · Base · VoiceDesign, 0.6B/1.7B); VoxCPM2 (zero-shot · cloned · described voice, 30+ languages, 48 kHz) | ✅ verified end-to-end, buffered + streaming — [speech guide](docs/speech.md) | | **Omni** — any-to-any chat | Qwen3-Omni 30B-A3B | 🚧 chat works via mlx-vlm; speech-out chat in progress | | **ASR** — speech → text | Whisper, Qwen3-ASR, Voxtral, … | 🚧 planned — `/v1/audio/transcriptions` (#68) | -| **Diffusion** — text/image → image | — | 🚧 planned — roadmap (#2) | +| **Diffusion** — text/image → image | Z-Image-Turbo (6B, 9-step distilled, 4-bit) | ✅ `/v1/images/generations` (mflux seam) — SD tier next (#95) | Text-only LLMs and image-in/text-out VLMs load through their engines but are not this server's target categories. @@ -76,6 +76,7 @@ python -m venv .venv && source .venv/bin/activate pip install -e . # server core (MLX engine stack) pip install -e '.[omni]' # + vision/audio models (mlx-vlm) pip install -e '.[tts]' # + speech synthesis (mlx-audio) +pip install -e '.[image]' # + image generation, Z-Image-Turbo (mflux) ``` ### Dependency footprint @@ -85,6 +86,7 @@ pip install -e '.[tts]' # + speech synthesis (mlx-audio) | core | `mlx-lm`, `starlette`, `uvicorn` | 38 | ~440 MB | | + `[tts]` | + `mlx-audio` | 44 | ~560 MB | | + `[omni]` | + `mlx-vlm` | 59 | ~750 MB | +| + `[image]` | + `mflux` | ~55 | ~870 MB | The core install pulls in the MLX stack (`mlx` + `mlx-metal` kernels, `transformers`, `tokenizers`, `huggingface_hub`) plus starlette/uvicorn and almost nothing else — @@ -92,8 +94,10 @@ The core install pulls in the MLX stack (`mlx` + `mlx-metal` kernels, `transform `mlx-audio` (miniaudio, sounddevice — still no torch). The `[omni]` extra adds ~315 MB through `mlx-vlm` (opencv, pillow, scipy — and mlx-audio, so `[omni]` implies `[tts]`; that path does drag in fastapi/pydantic, contained to the -optional extra). Measured on macOS arm64 / Python 3.13 with mlx-lm 0.32, -mlx-vlm 0.7, and mlx-audio 0.5.7. +optional extra). The `[image]` extra adds ~870 MB through `mflux` — the one +extra that does pull `torch` (583 MB of it, for mflux's utility stack; MLX +still runs every model). Measured on macOS arm64 / Python 3.13 with mlx-lm 0.32, +mlx-vlm 0.7, mlx-audio 0.5.7, and mflux 0.21. ### Run @@ -104,6 +108,9 @@ vllm-omni-mlx serve mlx-community/Qwen3-Omni-30B-A3B-Instruct-4bit \ # speech-only server vllm-omni-mlx serve mlx-community/Qwen3-TTS-12Hz-1.7B-CustomVoice-4bit --omni + +# image-only server (Z-Image-Turbo, 4-bit ~5.5 GiB of weights) +vllm-omni-mlx serve --image-model z-image-turbo # needs [image] ``` Qwen3-Omni serves text-out chat today (speech-out chat is in progress); the 30B-A3B diff --git a/docs/adding-a-model.md b/docs/adding-a-model.md index af873e8..55f7273 100644 --- a/docs/adding-a-model.md +++ b/docs/adding-a-model.md @@ -162,3 +162,30 @@ Apple-Silicon budget, not the ceiling. - [ ] Numbers in the PR: TTFA p50/p95, RTF distribution, peak memory — solo *and* with a chat model resident - [ ] HNR gate calibrated per voice, human audition done + +## Image models (`diffusion/`, #91) + +The seam is adapter-shaped, not loop-shaped — no streaming, one blocking +call per request: + +1. **Family registry** (`diffusion/config.py`): add a `Family` (aliases, + canonical repo, default steps, `supports_guidance`) and route refs in + `resolve_model()` → `ResolvedModel(family, weights_ref, quantize)`. + Weight routing matters: pre-quantized mirror repos honor their stored + level (`quantize=None`); canonical fp16 repos quantize on load + (`quantize=4`). Keep `config.py` importable without the `[image]` + extra (CI installs none). +2. **Service** (`diffusion/mflux_service.py`): `MFluxService` maps the + request onto the mflux backend call, validates family capabilities + (guidance on a guidance-distilled model → `ValueError`, the route + maps it to a 400), runs load+generate on one dedicated worker thread + (MLX stream affinity — serve smoke is mandatory, see #71), and + returns `ImageResult` (PNG bytes + seed + timings + peak + memory) — the metrics the doctrine reports. +3. **Tests** (`tests/diffusion/`, mirroring src): CI-safe resolution + + endpoint tests with a fake service; one weight-gated file per model + (per-file subprocess, offline, donor release in `tearDownClass`). + Gates: valid PNG above a size floor, dimensions/steps honored, + **same seed → byte-identical PNG** (characterize, don't assume, if + fusion flips appear), family-capability 400s. +4. **Docs**: README models-table row + the `[image]` footprint row. diff --git a/pyproject.toml b/pyproject.toml index cb6612c..1570ee2 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -37,10 +37,13 @@ Changelog = "https://github.com/ThinkFlowLab/vllm-omni-mlx/releases" [project.optional-dependencies] omni = ["mlx-vlm>=0.7"] tts = ["mlx-audio>=0.5.7,<0.6"] +# minor-capped like mlx-audio: the mflux seam touches mflux.models.* internals +# (loader classes), so a minor bump is an adaptation PR, not a free upgrade +image = ["mflux>=0.21.0,<0.22"] [project.scripts] # serve/tts subcommands, upstream vllm-omni's `vllm serve --omni` shape vllm-omni-mlx = "vllm_omni_mlx.__main__:main" [tool.setuptools] -packages = ["vllm_omni_mlx", "vllm_omni_mlx.tts"] +packages = ["vllm_omni_mlx", "vllm_omni_mlx.tts", "vllm_omni_mlx.diffusion"] diff --git a/tests/diffusion/__init__.py b/tests/diffusion/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tests/diffusion/test_config.py b/tests/diffusion/test_config.py new file mode 100644 index 0000000..3b0542e --- /dev/null +++ b/tests/diffusion/test_config.py @@ -0,0 +1,79 @@ +"""Family registry + request parsing for /v1/images/generations (#91). + +CI-safe: no mflux import (config.py keeps backend imports lazy) and no +weights — resolution and defaults only. +""" + +import unittest + +from vllm_omni_mlx.diffusion.config import DEFAULT_MODEL, FAMILIES, resolve_model +from vllm_omni_mlx.schemas import ApiError, parse_image_generation + + +class ResolveModelTest(unittest.TestCase): + def test_alias_routes_to_prequant_default(self): + for alias in ("z-image-turbo", "ZImage-Turbo"): + resolved = resolve_model(alias) + self.assertEqual(resolved.family.name, "z-image-turbo") + self.assertEqual(resolved.weights_ref, DEFAULT_MODEL) + self.assertIsNone(resolved.quantize) # stored 4-bit level honored + + def test_canonical_repo_quantizes_on_load(self): + resolved = resolve_model("Tongyi-MAI/Z-Image-Turbo") + self.assertEqual(resolved.weights_ref, "Tongyi-MAI/Z-Image-Turbo") + self.assertEqual(resolved.quantize, 4) + + def test_prequant_mirror(self): + resolved = resolve_model(DEFAULT_MODEL) + self.assertEqual(resolved.weights_ref, DEFAULT_MODEL) + self.assertIsNone(resolved.quantize) + + def test_unknown_ref_lists_supported(self): + with self.assertRaises(ValueError) as ctx: + resolve_model("stabilityai/stable-diffusion-xl-base-1.0") + self.assertIn("z-image-turbo", str(ctx.exception)) + + def test_every_family_has_serving_defaults(self): + for family in FAMILIES: + self.assertGreaterEqual(family.steps, 1, family.name) + + +class ParseImageGenerationTest(unittest.TestCase): + def test_minimal_request(self): + req = parse_image_generation({"prompt": "a cat"}) + self.assertEqual((req.width, req.height, req.n, req.steps), (1024, 1024, 1, None)) + + def test_size_and_diffusion_params(self): + req = parse_image_generation( + {"prompt": "a cat", "n": 2, "size": "768X1344", "steps": 9, "guidance": 0.0, "seed": 7} + ) + self.assertEqual((req.width, req.height, req.n, req.steps, req.guidance, req.seed), (768, 1344, 2, 9, 0.0, 7)) + + def test_prompt_must_be_nonempty_string(self): + for bad in (None, 5, "", " "): + with self.assertRaises(ApiError): + parse_image_generation({"prompt": bad}) + + def test_size_format(self): + for bad in ("1024", "ax1024", "1024x", 1024, "1024.5x1024"): + with self.assertRaises(ApiError): + parse_image_generation({"prompt": "a cat", "size": bad}) + + def test_dimension_bounds_and_grid(self): + for bad in ("240x1024", "2112x1024", "1024x1000"): # below/above bounds, not %16 + with self.assertRaises(ApiError): + parse_image_generation({"prompt": "a cat", "size": bad}) + + def test_n_range(self): + for bad in (0, 5, "2", True): + with self.assertRaises(ApiError): + parse_image_generation({"prompt": "a cat", "n": bad}) + + def test_scalar_types(self): + for payload in ({"prompt": "a cat", "steps": "9"}, {"prompt": "a cat", "guidance": "0"}, {"prompt": "a cat", "seed": 1.5}): + with self.assertRaises(ApiError): + parse_image_generation(payload) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/diffusion/test_images_endpoint.py b/tests/diffusion/test_images_endpoint.py new file mode 100644 index 0000000..1798f4c --- /dev/null +++ b/tests/diffusion/test_images_endpoint.py @@ -0,0 +1,97 @@ +"""POST /v1/images/generations (#91): validation, auth, model listing, +and the b64/seeds response shape with a fake service. The real +Z-Image-Turbo round-trip is weight-gated in test_mflux_zimage.py.""" + +import base64 +import unittest + +from starlette.testclient import TestClient + +from vllm_omni_mlx.diffusion.service import ImageResult +from vllm_omni_mlx.server import create_app + + +class FakeImageService: + name = "fake-image-model" + model_type = "image" + + def __init__(self): + self.calls = [] + + def generate(self, prompt, n=1, width=1024, height=1024, steps=None, guidance=None, seed=None): + if not prompt or not prompt.strip(): + raise ValueError("prompt must be a non-empty string") + if not 1 <= n <= 4: + raise ValueError(f"n must be between 1 and 4, got {n}") + if guidance is not None: + raise ValueError("guidance is not supported on z-image-turbo (guidance-distilled)") + self.calls.append((prompt, n, width, height, steps, guidance, seed)) + return [ + ImageResult( + png=b"\x89PNG-fake-" + bytes([i]), + seed=(seed if seed is not None else 1234) + i, + width=width, + height=height, + steps=steps or 9, + generation_time=0.5, + peak_memory_gib=1.0, + ) + for i in range(n) + ] + + +class ImagesEndpointTest(unittest.TestCase): + def setUp(self): + self.service = FakeImageService() + self.client = TestClient(create_app(image_service=self.service)) + + def test_round_trip(self): + response = self.client.post( + "/v1/images/generations", json={"prompt": "a cat", "size": "512x512", "steps": 9, "seed": 42} + ) + self.assertEqual(response.status_code, 200) + body = response.json() + self.assertEqual(body["model"], "fake-image-model") + self.assertEqual(len(body["data"]), 1) + self.assertEqual(base64.b64decode(body["data"][0]["b64_json"]), b"\x89PNG-fake-\x00") + self.assertEqual(body["seeds"], [42]) + self.assertEqual(body["generation_time_ms"], [500]) + self.assertEqual(self.service.calls, [("a cat", 1, 512, 512, 9, None, 42)]) + + def test_n_sequential_seeds(self): + body = self.client.post("/v1/images/generations", json={"prompt": "a cat", "n": 3, "seed": 10}).json() + self.assertEqual(body["seeds"], [10, 11, 12]) + self.assertEqual(len(body["data"]), 3) + + def test_validation_errors_are_400(self): + for payload in ({"prompt": ""}, {"prompt": "a cat", "size": "1000x1024"}, {"prompt": "a cat", "n": 9}): + response = self.client.post("/v1/images/generations", json=payload) + self.assertEqual(response.status_code, 400, payload) + self.assertEqual(response.json()["error"]["type"], "invalid_request_error") + + def test_family_capability_error_is_400(self): + response = self.client.post("/v1/images/generations", json={"prompt": "a cat", "guidance": 3.5}) + self.assertEqual(response.status_code, 400) + self.assertIn("guidance-distilled", response.json()["error"]["message"]) + + def test_auth(self): + client = TestClient(create_app(image_service=self.service, api_key="secret")) + self.assertEqual(client.post("/v1/images/generations", json={"prompt": "a cat"}).status_code, 401) + self.assertEqual( + client.post( + "/v1/images/generations", json={"prompt": "a cat"}, headers={"Authorization": "Bearer secret"} + ).status_code, + 200, + ) + + def test_route_absent_without_service(self): + client = TestClient(create_app()) + self.assertEqual(client.post("/v1/images/generations", json={"prompt": "a cat"}).status_code, 404) + + def test_model_listing(self): + body = self.client.get("/v1/models").json() + self.assertIn("fake-image-model", [entry["id"] for entry in body["data"]]) + + +if __name__ == "__main__": + unittest.main() diff --git a/tests/diffusion/test_mflux_zimage.py b/tests/diffusion/test_mflux_zimage.py new file mode 100644 index 0000000..3202302 --- /dev/null +++ b/tests/diffusion/test_mflux_zimage.py @@ -0,0 +1,77 @@ +"""Real Z-Image-Turbo round-trip through the mflux seam (#99). + +Weight-gated: runs against the local HF snapshot when cached and skips +elsewhere (CI has no weights, so a green CI run does not exercise this +file). One file = one process (CONTRIBUTING) — the 4-bit model peaks +around 6–7 GiB resident and holds it for the class. +""" + +import os + +# offline keeps the load fast and hang-free when the network is flaky +os.environ.setdefault("HF_HUB_OFFLINE", "1") + +import unittest + +from vllm_omni_mlx.diffusion.config import DEFAULT_MODEL, load_image_service + + +def _snapshot_cached() -> bool: + # direct check rather than snapshot_download(local_files_only=True): + # a patterns-filtered download counts as "incomplete" to the hub API + # even when every weight file is present + import glob + import os + + pattern = os.path.expanduser( + f"~/.cache/huggingface/hub/models--{DEFAULT_MODEL.replace('/', '--')}/snapshots/*/transformer/*.safetensors" + ) + return bool(glob.glob(pattern)) + + +@unittest.skipUnless(_snapshot_cached(), f"{DEFAULT_MODEL} not in the local HF cache") +class ZImageTurboE2ETest(unittest.TestCase): + """512² × 9 steps keeps the run quick while exercising the full path: + pre-quant load → text encode → 9 denoise steps → VAE decode → PNG.""" + + @classmethod + def setUpClass(cls): + import mlx.core as mx + + cls.mx = mx + cls.service = load_image_service(DEFAULT_MODEL) + cls.prompt = "a crisp red apple on a wooden table, studio light" + + @classmethod + def tearDownClass(cls): + # donor release: one heavy thing at a time on a 16 GB machine + cls.service = None + cls.mx.clear_cache() + + def test_generate_png_shape_and_gate(self): + (result,) = self.service.generate(self.prompt, width=512, height=512, steps=9, seed=20261007) + self.assertTrue(result.png.startswith(b"\x89PNG"), "output must be a PNG") + self.assertGreater(len(result.png), 50_000, "a 512² PNG should not be tiny (non-degenerate gate)") + self.assertEqual((result.width, result.height, result.steps), (512, 512, 9)) + self.assertEqual(result.seed, 20261007) + self.assertGreater(result.generation_time, 0.1) + print( + f"\nz-image-turbo 512x512/9 steps: {result.generation_time:.2f}s, " + f"peak {result.peak_memory_gib:.2f} GiB, png {len(result.png) / 1e6:.2f} MB", + flush=True, + ) + + def test_same_seed_reproduces_bytes(self): + first = self.service.generate(self.prompt, width=512, height=512, steps=9, seed=7)[0].png + second = self.service.generate(self.prompt, width=512, height=512, steps=9, seed=7)[0].png + self.assertEqual( + first, second, "same-seed calls must reproduce byte-identical PNGs (fp16 fusion-flip precedent: characterize if this trips)" + ) + + def test_family_capability_validation(self): + with self.assertRaises(ValueError): + self.service.generate(self.prompt, width=512, height=512, guidance=3.5) + + +if __name__ == "__main__": + unittest.main() diff --git a/vllm_omni_mlx/__main__.py b/vllm_omni_mlx/__main__.py index 6ab9ef3..a35e409 100644 --- a/vllm_omni_mlx/__main__.py +++ b/vllm_omni_mlx/__main__.py @@ -27,6 +27,11 @@ def build_serve_parser() -> argparse.ArgumentParser: default=None, help="also serve speech synthesis on /v1/audio/* (Qwen3-TTS via the [tts] extra); the only model when is omitted", ) + parser.add_argument( + "--image-model", + default=None, + help="also serve image generation on /v1/images/generations (Z-Image-Turbo via the [image] extra); the only model when is omitted", + ) parser.add_argument("--host", default="127.0.0.1", help="bind address (default: 127.0.0.1)") parser.add_argument("--port", type=int, default=8000, help="bind port (default: 8000)") parser.add_argument( @@ -108,8 +113,8 @@ def main(argv: list[str] | None = None) -> int: def _serve(args) -> int: - if not args.model and not args.tts_model: - build_serve_parser().error("a model is required (or --tts-model to serve speech alone)") + if not args.model and not args.tts_model and not args.image_model: + build_serve_parser().error("a model is required (or --tts-model / --image-model to serve one alone)") if args.omni and not args.model: build_serve_parser().error("--omni applies to the served model") if args.omni and args.tts_model: @@ -152,12 +157,42 @@ def _serve(args) -> int: print(f"error: failed to load TTS model '{args.tts_model}': {exc}", file=sys.stderr) return 1 - names = [n for n in (getattr(backend, "name", None), getattr(tts_service, "name", None)) if n] + image_service = None + if args.image_model: + try: + image_service = _load_image(args.image_model) + except (RuntimeError, ValueError, OSError) as exc: + print(f"error: failed to load image model '{args.image_model}': {exc}", file=sys.stderr) + return 1 + + names = [ + n + for n in ( + getattr(backend, "name", None), + getattr(tts_service, "name", None), + getattr(image_service, "name", None), + ) + if n + ] print(f"serving {', '.join(names)} on http://{args.host}:{args.port}", file=sys.stderr) - uvicorn.run(create_app(backend, api_key=args.api_key, tts_service=tts_service), host=args.host, port=args.port, log_level=args.log_level) + uvicorn.run( + create_app(backend, api_key=args.api_key, tts_service=tts_service, image_service=image_service), + host=args.host, + port=args.port, + log_level=args.log_level, + ) return 0 +def _load_image(model_ref: str): + """Load an image-generation model for serving (#91/#99). Z-Image-Turbo + via the mflux seam: pre-quantized mirror repos honor their stored + 4-bit level; the canonical fp16 repo quantizes on load.""" + from .diffusion import load_image_service + + return load_image_service(model_ref) + + def _load_tts(model_ref: str): """Load a TTS checkpoint for serving — shared by the --omni and --tts-model paths. Qwen3-TTS: load, reject variants this build can't diff --git a/vllm_omni_mlx/diffusion/__init__.py b/vllm_omni_mlx/diffusion/__init__.py new file mode 100644 index 0000000..96c0495 --- /dev/null +++ b/vllm_omni_mlx/diffusion/__init__.py @@ -0,0 +1,18 @@ +"""Image generation seam (diffusion, #91/#99): config + family registry, +service protocol, and the mflux adapter. See docs/architecture.md for the +layer map and the #91 design record for the API/metrics doctrine.""" + +from .config import DEFAULT_MODEL, DiffusionConfig, Family, FAMILIES, ResolvedModel, load_image_service, resolve_model +from .service import DiffusionService, ImageResult + +__all__ = [ + "DEFAULT_MODEL", + "DiffusionConfig", + "DiffusionService", + "FAMILIES", + "Family", + "ImageResult", + "ResolvedModel", + "load_image_service", + "resolve_model", +] diff --git a/vllm_omni_mlx/diffusion/config.py b/vllm_omni_mlx/diffusion/config.py new file mode 100644 index 0000000..b07a546 --- /dev/null +++ b/vllm_omni_mlx/diffusion/config.py @@ -0,0 +1,99 @@ +"""Typed configuration and family registry for image generation (#91, #99). + +The diffusion package is a seam over library backends (mflux first, the +vendored mlx-examples SD family later) — no native model code lives here +(same rule as tts/ over mlx-audio, see docs/architecture.md). This module +resolves a served model ref to a family and the weights source that goes +with it. It imports no backend at module scope so CI (which installs the +package without extras) can exercise resolution and request validation. + +Z-Image-Turbo notes (mflux 0.21): guidance-distilled — guidance is forced +to 0.0 and a negative prompt is never encoded (mflux's own CLI ignores +both with a warning; we reject guidance with a 400 instead, the same +early-clear-signal choice as instruct-on-0.6B). The default source is the +pre-quantized 4-bit mirror (~5.5 GiB on disk vs ~12.3 GiB fp16 for the +canonical repo, whose stored quantization level is honored on load). +""" + +from __future__ import annotations + +from dataclasses import dataclass +from typing import Optional + +# Pre-quantized 4-bit mirror of Tongyi-MAI/Z-Image-Turbo in mflux's save +# format — the default source (the canonical fp16 repo does not fit the +# dev machine's disk budget). +DEFAULT_MODEL = "filipstrand/Z-Image-Turbo-mflux-4bit" + + +@dataclass(frozen=True) +class Family: + """A model family behind the mflux seam: what a request may set and + which refs identify it. Weight routing lives in ResolvedModel.""" + + name: str + aliases: tuple[str, ...] + canonical_repo: str + steps: int + supports_guidance: bool = False + + +@dataclass(frozen=True) +class ResolvedModel: + family: Family + weights_ref: str # repo id or local path handed to mflux as model_path + quantize: Optional[int] = None # None → honor the checkpoint's stored level + + +_FAMILIES: tuple[Family, ...] = ( + Family( + name="z-image-turbo", + aliases=("z-image-turbo", "zimage-turbo"), + canonical_repo="Tongyi-MAI/Z-Image-Turbo", + steps=9, + ), +) + +FAMILIES = _FAMILIES # re-exported for tests/introspection + + +def resolve_model(model_ref: str) -> ResolvedModel: + ref = (model_ref or "").strip() + low = ref.lower() + for family in _FAMILIES: + if low in family.aliases: + # the alias is the family's disk/memory-friendly default: the + # pre-quantized mirror with its stored 4-bit level + return ResolvedModel(family, weights_ref=DEFAULT_MODEL) + if low == family.canonical_repo.lower(): + # canonical fp16 repo: quantize 4-bit on load (memory + parity + # with the mirror; the fp16 download is also ~12 GiB of disk) + return ResolvedModel(family, weights_ref=ref, quantize=4) + if low == DEFAULT_MODEL.lower(): + return ResolvedModel(_FAMILIES[0], weights_ref=DEFAULT_MODEL) + known = ", ".join(sorted({alias for f in _FAMILIES for alias in f.aliases})) + raise ValueError( + f"unknown image model '{model_ref}'; supported: {known}, " + f"canonical repos ({', '.join(f.canonical_repo for f in _FAMILIES)}), " + f"or the pre-quantized mirror ({DEFAULT_MODEL})" + ) + + +@dataclass(frozen=True) +class DiffusionConfig: + """Serving defaults for /v1/images/generations.""" + + model_ref: str = DEFAULT_MODEL + width: int = 1024 + height: int = 1024 + + +def load_image_service(model_ref: str = DEFAULT_MODEL) -> "MFluxService": # noqa: F821 + """Resolve ``model_ref`` and load it via the mflux seam. + + Imports mflux lazily: the [image] extra is optional, and config/test + imports must not require it (CI installs no extras). + """ + from .mflux_service import MFluxService + + return MFluxService.load(resolve_model(model_ref)) diff --git a/vllm_omni_mlx/diffusion/mflux_service.py b/vllm_omni_mlx/diffusion/mflux_service.py new file mode 100644 index 0000000..8a1f87e --- /dev/null +++ b/vllm_omni_mlx/diffusion/mflux_service.py @@ -0,0 +1,116 @@ +"""mflux seam: the adapter over the pip-installed mflux library (#91/#99). + +Everything model-specific stays inside mflux — this module maps a Family +onto mflux's loader (ZImage + ModelConfig + model_path for pre-quantized +mirrors), validates family capabilities, and wraps the blocking +generate_image call with the timings/metrics the image doctrine wants. +mflux is pinned like mlx-audio (see CONTRIBUTING.md): the classes this +module touches live under mflux.models.*, so a minor bump is an +adaptation PR, not a free upgrade. +""" + +from __future__ import annotations + +import secrets +import time +from concurrent.futures import ThreadPoolExecutor + +import mlx.core as mx + +from .config import Family, ResolvedModel +from .service import ImageResult, png_bytes + + +class MFluxService: + """All MLX work — load included — runs on one dedicated worker thread: + splitting load (uvicorn main thread) from generate (`asyncio.to_thread`) + trips MLX's per-thread stream context ("There is no Stream(cpu, 0) in + current thread" — same class as the VoxCPM2 serve-smoke bug, #71). The + single-thread executor doubles as the batch-1 serializer.""" + + def __init__(self, family: Family, resolved: ResolvedModel): + self.family = family + self.weights_ref = resolved.weights_ref + self._executor = ThreadPoolExecutor(max_workers=1, thread_name_prefix="mflux") + self._model = None + self.bits = None + + @classmethod + def load(cls, resolved: ResolvedModel) -> "MFluxService": + from mflux.models.common.config.model_config import ModelConfig + from mflux.models.z_image import ZImage + + service = cls(resolved.family, resolved) + started = time.perf_counter() + # model_path routes weights (pre-quantized mirrors keep their stored + # level when quantize=None); quantize applies on-load for fp16 repos. + service._model = service._executor.submit( + ZImage, + model_config=ModelConfig.z_image_turbo(), + quantize=resolved.quantize, + model_path=resolved.weights_ref, + ).result() + service.bits = getattr(service._model, "bits", resolved.quantize) + print( + f"image model '{resolved.family.name}' loaded from '{resolved.weights_ref}' in " + f"{time.perf_counter() - started:.1f}s ({service.bits}-bit)", + flush=True, + ) + return service + + @property + def name(self) -> str: + return self.weights_ref + + @property + def model_type(self) -> str: + return "image" + + def generate( + self, + prompt: str, + n: int = 1, + width: int = 1024, + height: int = 1024, + steps: int | None = None, + guidance: float | None = None, + seed: int | None = None, + ) -> list[ImageResult]: + if not prompt or not prompt.strip(): + raise ValueError("prompt must be a non-empty string") + if not 1 <= n <= 4: + raise ValueError(f"n must be between 1 and 4, got {n}") + if guidance is not None and not self.family.supports_guidance: + raise ValueError(f"guidance is not supported on {self.family.name} (guidance-distilled)") + steps = self.family.steps if steps is None else steps + if not 1 <= steps <= 50: + raise ValueError(f"steps must be between 1 and 50, got {steps}") + base_seed = seed if seed is not None else secrets.randbelow(2**32) + return self._executor.submit(self._generate_many, prompt, n, width, height, steps, guidance, base_seed).result() + + def _generate_many( + self, prompt: str, n: int, width: int, height: int, steps: int, guidance, base_seed: int + ) -> list[ImageResult]: + return [self._generate_one(prompt, width, height, steps, guidance, base_seed + i) for i in range(n)] + + def _generate_one(self, prompt: str, width: int, height: int, steps: int, guidance, seed: int) -> ImageResult: + mx.reset_peak_memory() + started = time.perf_counter() + image = self._model.generate_image( + seed=seed, + prompt=prompt, + num_inference_steps=steps, + height=height, + width=width, + guidance=guidance, + ) + elapsed = time.perf_counter() - started + return ImageResult( + png=png_bytes(image.image), + seed=seed, + width=width, + height=height, + steps=steps, + generation_time=elapsed, + peak_memory_gib=mx.get_peak_memory() / 2**30, + ) diff --git a/vllm_omni_mlx/diffusion/service.py b/vllm_omni_mlx/diffusion/service.py new file mode 100644 index 0000000..c077495 --- /dev/null +++ b/vllm_omni_mlx/diffusion/service.py @@ -0,0 +1,59 @@ +"""DiffusionService protocol + result type (#91). + +One long blocking call per request, not a token stream — so image +generation is a sibling service to tts/ rather than a Backend (whose +protocol is chat-shaped, see backends.py). The server injects whichever +service load_image_service() picked and runs generate() on a worker +thread; a per-service lock enforces one generation at a time (batch-1 +doctrine, same as the chat backend lock). +""" + +from __future__ import annotations + +import io +from dataclasses import dataclass +from typing import Protocol, runtime_checkable + + +@dataclass(frozen=True) +class ImageResult: + png: bytes + seed: int + width: int + height: int + steps: int + generation_time: float # seconds, measured around the model call + peak_memory_gib: float + + +@runtime_checkable +class DiffusionService(Protocol): + """What the /v1/images/generations route needs from a backend seam.""" + + @property + def name(self) -> str: ... + + @property + def model_type(self) -> str: ... + + def generate( + self, + prompt: str, + n: int = 1, + width: int = 1024, + height: int = 1024, + steps: int | None = None, + guidance: float | None = None, + seed: int | None = None, + ) -> list[ImageResult]: + """Generate ``n`` images sequentially; seeds are ``seed + i`` when a + seed is given (deterministic per index), fresh randoms otherwise. + Raises ValueError for family-unsupported parameters (mapped to a + 400 by the route, the audio_speech pattern).""" + ... + + +def png_bytes(image) -> bytes: + buffer = io.BytesIO() + image.save(buffer, format="PNG") + return buffer.getvalue() diff --git a/vllm_omni_mlx/schemas.py b/vllm_omni_mlx/schemas.py index 4cb40e1..e285c6d 100644 --- a/vllm_omni_mlx/schemas.py +++ b/vllm_omni_mlx/schemas.py @@ -291,3 +291,56 @@ def _opt_int(value, name: str) -> Optional[int]: if not isinstance(value, int) or isinstance(value, bool): raise ApiError(400, f"'{name}' must be an integer") return value + + +# -------------------------------------------------------------------------- +# /v1/images/generations (#91): OpenAI-compatible request, plus the +# diffusion params the category needs (steps/guidance/seed). Family +# capability checks (e.g. guidance on guidance-distilled models) live in +# the service, not here. +# -------------------------------------------------------------------------- + + +@dataclass +class ImageGenerationRequest: + prompt: str + n: int = 1 + width: int = 1024 + height: int = 1024 + steps: Optional[int] = None + guidance: Optional[float] = None + seed: Optional[int] = None + + +def parse_image_generation(payload: dict) -> ImageGenerationRequest: + prompt = payload.get("prompt") + if not isinstance(prompt, str) or not prompt.strip(): + raise ApiError(400, "prompt must be a non-empty string") + + n = payload.get("n", 1) + if not isinstance(n, int) or isinstance(n, bool) or not 1 <= n <= 4: + raise ApiError(400, "'n' must be an integer between 1 and 4") + + width, height = 1024, 1024 + if payload.get("size") is not None: + size = payload["size"] + if not isinstance(size, str) or "x" not in size.lower(): + raise ApiError(400, "'size' must be a 'x' string, e.g. '1024x1024'") + parts = size.lower().split("x", 1) + if not (parts[0].isdigit() and parts[1].isdigit()): + raise ApiError(400, f"'size' must be a 'x' string, got '{size}'") + width, height = int(parts[0]), int(parts[1]) + if not (256 <= width <= 2048 and 256 <= height <= 2048): + raise ApiError(400, f"image dimensions must be within 256–2048, got {width}x{height}") + if width % 16 or height % 16: + raise ApiError(400, f"image dimensions must be multiples of 16 (latent grid), got {width}x{height}") + + return ImageGenerationRequest( + prompt=prompt, + n=n, + width=width, + height=height, + steps=_opt_int(payload.get("steps"), "steps"), + guidance=_opt_float(payload.get("guidance"), "guidance"), + seed=_opt_int(payload.get("seed"), "seed"), + ) diff --git a/vllm_omni_mlx/server.py b/vllm_omni_mlx/server.py index 97c994e..65f0e3b 100644 --- a/vllm_omni_mlx/server.py +++ b/vllm_omni_mlx/server.py @@ -3,6 +3,7 @@ from __future__ import annotations import asyncio +import base64 import json import threading import time @@ -15,7 +16,13 @@ from starlette.routing import Route from .backends import Backend, Chunk -from .schemas import ApiError, UnifiedRequest, normalize_anthropic, normalize_openai +from .schemas import ( + ApiError, + UnifiedRequest, + normalize_anthropic, + normalize_openai, + parse_image_generation, +) _STARTED_AT = int(time.time()) @@ -221,7 +228,7 @@ async def stream() -> AsyncIterator[bytes]: # app factory # -------------------------------------------------------------------------- -def create_app(backend: Backend | None = None, api_key: str | None = None, tts_service=None) -> Starlette: +def create_app(backend: Backend | None = None, api_key: str | None = None, tts_service=None, image_service=None) -> Starlette: async def health(request: Request) -> Response: return JSONResponse({"status": "ok"}) @@ -233,7 +240,11 @@ async def models(request: Request) -> Response: "created": _STARTED_AT, "owned_by": "vllm-omni-mlx", } - for name in ([backend.name] if backend is not None else []) + ([tts_service.name] if tts_service is not None else []) + for name in ( + ([backend.name] if backend is not None else []) + + ([tts_service.name] if tts_service is not None else []) + + ([image_service.name] if image_service is not None else []) + ) ] return JSONResponse({"object": "list", "data": entries}) @@ -371,6 +382,40 @@ async def audio_voices(request: Request) -> Response: return _openai_error(exc) return JSONResponse({"object": "list", "voices": tts_service.voices}) + async def images_generations(request: Request) -> Response: + try: + _check_auth(request, api_key) + req = parse_image_generation(await _json_body(request)) + # blocking, seconds-to-tens-of-seconds: off the event loop, one + # request at a time inside the service (per-service lock) + results = await asyncio.to_thread( + image_service.generate, + req.prompt, + n=req.n, + width=req.width, + height=req.height, + steps=req.steps, + guidance=req.guidance, + seed=req.seed, + ) + return JSONResponse( + { + "created": int(time.time()), + "model": image_service.name, + "data": [{"b64_json": base64.b64encode(r.png).decode("ascii")} for r in results], + # vllm-omni-mlx extensions: reproducibility + the image + # doctrine's metrics (per image) + "seeds": [r.seed for r in results], + "generation_time_ms": [round(r.generation_time * 1000) for r in results], + } + ) + except ApiError as exc: + return _openai_error(exc) + except ValueError as exc: + return _openai_error(ApiError(400, str(exc), err_type="invalid_request_error")) + except Exception as exc: + return _openai_error(ApiError(500, f"image generation failed: {exc}", err_type="server_error")) + routes = [ Route("/health", health), Route("/v1/models", models), @@ -385,4 +430,8 @@ async def audio_voices(request: Request) -> Response: Route("/v1/audio/speech", audio_speech, methods=["POST"]), Route("/v1/audio/voices", audio_voices), ] + if image_service is not None: + routes += [ + Route("/v1/images/generations", images_generations, methods=["POST"]), + ] return Starlette(routes=routes)