diff --git a/deploy/profiles/mac-apple-silicon.json b/deploy/profiles/mac-apple-silicon.json index bfeb4f8..c5e4235 100644 --- a/deploy/profiles/mac-apple-silicon.json +++ b/deploy/profiles/mac-apple-silicon.json @@ -14,7 +14,8 @@ "env": { "LR_BIN": "$REPO/zig-out/bin/localrouter", "DEVELOPER_DIR": "${DEVELOPER_DIR:-/Applications/Xcode.app/Contents/Developer}", - "STEPS": "bash $REPO/tools/install/steps.sh" + "STEPS": "bash $REPO/tools/install/steps.sh", + "IMAGE_ROOT": "$J/image-mac" }, "prereqs": [ { @@ -173,6 +174,39 @@ "fix": "network: run again (the download resumes). \"do not match SHA256SUMS\": delete $J/models/decide/knowline-4b-gen5.5 and run again" } ] + }, + { + "id": "qwen-image-2.1-turbo", + "key": "image", + "title": "Qwen-Image 2.1 Turbo text to image and edits of up to five images (MLX-Serve 4-bit, 8 steps; a cmd worker, loaded on demand)", + "licence": { + "name": "Qwen Research License: non-commercial research or evaluation only", + "url": "https://huggingface.co/Qwen/Qwen-Image-2.1/blob/main/LICENSE" + }, + "accept": [ + "the Qwen-Image 2.1 Turbo licence (link below)" + ], + "pins": "tools/image/qwen-mlx/pins-macos-q4.json (4-bit, installed), tools/image/qwen-mlx/pins-macos.json (the MLX-Serve binary; the 8-bit alternative), tools/image/qwen-mlx/requirements-macos.txt", + "memory_bytes": 10200547328, + "working_bytes": 6979321856, + "download_bytes": 10628459115, + "disk_bytes": 11500000000, + "warm": false, + "tool": "tools/image/qwen-mlx/config.example.json: cmd $IMAGE_ROOT/venv/bin/python -I $REPO/tools/image/qwen-mlx/qwen_mlx_worker.py, weights $IMAGE_ROOT/qwen-image-2.1-turbo-mlx-q4, options pack q4 and reference_working_bytes", + "skip": "leave image out of --models; picture requests then fail with 'no image model'", + "steps": [ + { + "id": "fetch", + "title": "the MLX-Serve 26.10.3 binary, the 4-bit Qwen-Image 2.1 Turbo pack (10.5 GB, every file against its pin) and the worker's Python", + "download_bytes": 10628459115, + "disk_bytes": 11500000000, + "needs_accept": true, + "run": "bash $REPO/tools/image/qwen-mlx/fetch.sh $IMAGE_ROOT --pack q4", + "done": "[ -x $IMAGE_ROOT/venv/bin/python ] && $IMAGE_ROOT/mlx-serve-macos-arm64/mlx-serve --version 2> /dev/null | grep -qx 'mlx-serve 26.10.3' && python3 -I $REPO/tools/package/verify_assets.py --quick $REPO/tools/image/qwen-mlx/pins-macos-q4.json $IMAGE_ROOT/qwen-image-2.1-turbo-mlx-q4", + "fix": "network: run again (finished files are kept, a partial one resumes). 'sha256 … pinned' or 'differs from': delete the named file and run again" + } + ], + "reference_working_bytes": 2147483648 } ] } diff --git a/docs/dev/IMAGE-EDIT.md b/docs/dev/IMAGE-EDIT.md index e290b6a..537c772 100644 --- a/docs/dev/IMAGE-EDIT.md +++ b/docs/dev/IMAGE-EDIT.md @@ -218,3 +218,104 @@ which differs between CUDA 13.0 (the box's nvcc, which builds the engine) and 13 four GELU ops differed by a few bytes in 66,709 and the chained replay diverged from there. The same header serves the H3 vision tower. References are resized to about one megapixel with sides up to 2048 (4:1), independent of the output's `max_side`; the worker's memory is 38.2 / 41.3 / 44.1 GiB at `edit_refs` 1 / 2 / 3 (31.2 GiB without the pack). + +## macOS (Apple Silicon): MLX-Serve worker + +On a Mac the image tool is a `cmd` worker, `tools/image/qwen-mlx/qwen_mlx_worker.py`, not the Zig/CUDA engine above. It +starts MLX-Serve (a native Zig/MLX server) on 127.0.0.1 and is its client, the "model's own server" shape of +docs/ARCHITECTURE.md. Text to image and edits with one to five input images, the same checkpoint (Qwen-Image 2.1 +Turbo, its own unshifted eight-node `sample_sigmas` schedule). + +| Part | Pin (`tools/image/qwen-mlx/pins-macos.json`) | +| --- | --- | +| Runtime | MLX-Serve 26.10.3 release binary `mlx-serve-bin-macos-arm64.tar.gz`, sha256 `d1781c89…652b` (source `046b6476`); 26.10.2 is the first that runs Turbo's own schedule, so the worker refuses older | +| Weights | `ddalcu/Qwen-Image-2.1-Turbo-MLX-Serve-8bit@214ca346`, 17.7 GB: MLX affine 8-bit (group 64) DiT and text-encoder linears, dense f32 VAE, the Qwen3-VL vision tower; every file's size and digest checked by `verify_assets.py` | +| 4-bit weights (the Mac profile's) | `ddalcu/Qwen-Image-2.1-Turbo-MLX-Serve-4bit@394f9477`, 10.5 GB (`pins-macos-q4.json`): 4-bit DiT and text-encoder linears, the same tower, VAE and schedule; `fetch.sh ROOT --pack q4`, option `pack: "q4"`. The profile and `config.example.json` install and run this pack; its bookings are below | +| Worker Python | Pillow 13.0.0 (`requirements-macos.txt`, hashes required) | + +Setup: `tools/image/qwen-mlx/fetch.sh ROOT --pack q4` (binary, pack, venv; without `--pack` the 8-bit one), then a tool like +`tools/image/qwen-mlx/config.example.json`. Tests without a model: `ROOT/venv/bin/python -I +tools/image/qwen-mlx/test_worker.py` (lifecycle, memory, refusals, cancellation, failures, outputs, against a fake +MLX-Serve). + +**Licence.** The weights are Qwen's, under the Qwen Research License Agreement: non-commercial research or evaluation +only; commercial use needs a separate licence from Qwen. The 8-bit conversion does not change that. MLX-Serve is MIT, +its bundled MLX libraries Apache-2.0. + +**Same request, same meaning.** The worker applies the Spark tool's rules itself (the daemon does not run a `cmd` +tool's `check`): `max_side` / `max_pixels` on the output, steps 0 or 8, `n` images with seeds `seed`, `seed+1`, …, +up to `edit_refs` (at most 5) input images in order (the first is the one edited), an edit without a size gets the last +image's shape at about one megapixel within the limits (`editSize`), inputs under 32 pixels a side or past 4:1 are +refused, and a mask (alpha 0 where the edit goes) keeps the first image where it is opaque (the `Keep` composite, done +with Pillow: premultiplied Lanczos, then `Image.composite`; at the soft edges its rounding may differ from the Zig +blend by 1). `guidance` and `negative_prompt` are not used, as on the Spark. Differences, refused rather than changed +quietly: an explicit edit size must be a multiple of 32 (MLX-Serve floors it); WebP inputs are sent as lossless PNG +(MLX-Serve decodes PNG and JPEG). Output is RGBA (the VAE's own alpha, MLX-Serve's `transparent`), with `prompt` and +`seed` text chunks; the daemon adds the machine-generated label. Cancel: SIGUSR1 closes the stream, MLX-Serve stops at +its next step and stays loaded. A dead server takes the worker with it (a 503; the next request reloads). + +**Numerics.** This is a deliberate numerical change from the Spark engine (8-bit MLX weights, MLX kernels), not a port +of it: no bit-parity with the Spark or the twin is claimed. The Mac reference is the baseline below, MLX-Serve 26.10.3 +with its lossy `--w8a8` (int8 x int8 DiT linears, on by default on M5 with 64 GiB or more) turned **off** +(option `w8a8`, default false). + +### Baseline (2026-10-30, Mac Studio M5 Ultra, 96 GiB, macOS 27.2) + +`tools/image/qwen-mlx/baseline.py ROOT OUT` drives the worker through the protocol as the daemon does, under the +machine's GPU lock. Its inputs are made first by the same model (so nothing comes from outside): `woman` 1024x1024 seed +11, `hat` 768x768 seed 12, `cup` 1024x768 seed 13, `dog` 768x1024 seed 14, `room` 1344x768 seed 15 (prompts in the +script). Every case runs with seeds 7 and 8; edits leave the size to the model (the last input's shape: 1024x1024, or +1344x768 for five inputs). Memory: the server's resident set sampled every 100 ms (unified memory, MLX's buffers are in +it; MLX's own `/props` peak agreed within 0.1 GB). + +| Case | Inputs | Seconds a image (seed 7 / 8) | Peak resident | +| --- | --- | --- | --- | +| load (server start, model, a 256x256 warm-up image) | | 2.3 (weights in the file cache; 11.9 s server start cold) | 16.9 GB resident | +| text to image 1024x1024 | | 8.48 / 8.50 | 16.3 GB | +| edit "make it night time…" | woman | 12.92 / 11.26 | 16.8 GB | +| edit "…wearing the red beanie from image 2" | woman, hat | 14.25 / 14.22 | 16.6 GB | +| edit "…sitting in the armchair of the room in image 5…" | all five | 25.59 / 24.44 | 17.3 GB | +| masked edit "give her a bright yellow sweater" | woman + mask | 11.39 / 11.40 | 16.7 GB | +| exit | | 1.9 | | + +Peak over the run 17.33 GB; an 8-bit tool books 16 GiB resident plus 2 GiB working. For comparison the Spark (GB10, FP8) +makes the 1024x1024 image in 6.3 s and edits one input in 11 to 12.5 s. + +**Bookings (the 4-bit pack).** The scheduler books an image job at `resident_bytes + working_bytes`, and an edit at +`working_bytes + reference_working_bytes × inputs` on top of the resident set (`reference_working_bytes` is an option +the daemon reads; the worker leaves it alone). The integration run measured the 4-bit pack under load: 9.27 GiB +resident after the load, a 15.1 GiB peak for a 1-megapixel image, 13.2 GiB for an edit of one input and 24.3 GiB for +one of five. The profile and `config.example.json` book it at: + +| Field | Bytes | | +| --- | --- | --- | +| `resident_bytes` | 10200547328 | 9.5 GiB (9.27 measured at load) | +| `working_bytes` | 6979321856 | 6.5 GiB (9.5 + 6.5 = 16 GiB against the 15.1 GiB peak of a 1 MP image) | +| `options.reference_working_bytes` | 2147483648 | 2 GiB an input: 18 GiB booked for a 1-input edit (13.2 measured), 26 GiB for 5 inputs (24.3 measured) | + +On the Mac profile the image model loads on demand (`warm: false`, `keep_loaded: false`): chat, Drama 3 and Whisper +stay warm (about 66 of the GPU's 77.8 GiB), and the 8-bit pack's 18 GiB does not fit beside them, so the profile +installs the 4-bit pack. Still to do: the 4-bit pack against the 8-bit baseline above on quality and speed. + +**Why MLX-Serve and not mflux.** Same Turbo checkpoint at 8 bits (`mflux-community/qwen-image-2-1-turbo-mflux-q8@19d2068a`, +which also has the vision tower), mflux `5171ae15` `QwenImage21Edit`, the same prompts, inputs and seeds, one load: + +| | MLX-Serve 26.10.3 | mflux 5171ae15 | +| --- | --- | --- | +| text to image 1024x1024 (s, seed 7 / 8) | 8.48 / 8.50 | 11.13 / 8.33 | +| edit, 1 input | 12.92 / 11.26 | 44.01 / 52.89 | +| edit, 2 inputs | 14.25 / 14.22 | 32.31 / 25.43 | +| edit, 5 inputs | 25.59 / 24.44 | 36.46 / 39.11 | +| peak memory (MLX) | 16.8 to 17.3 GB | 29.1 to 34.0 GB | + +**Through the daemon** (the same Mac, LocalRouter built with `zig build -Doptimize=ReleaseSafe`, the tool of +`config.example.json`; seconds are the HTTP round trip): `/v1/images/generations` 1024x1024 10.8 s for the first +request (loading the tool included), `n` 2 16.8 s; `/v1/images/edits` with 1, 2 and 5 inputs 11.3, 14.0 and 24.0 s, a +mask 11.2 s; MCP `generate_image` 8.6 s, `edit_image` with 1 and 2 images 11.3 and 14.5 s. Refused with a 400 and +the reason: six inputs (the API), steps 25, 1600x1024, and an edit at 1008x1008 (the worker). Outputs carry the +machine-generated label and the `prompt` / `seed` chunks; the scheduler measured 16,089 MiB resident. A SIGKILLed +daemon leaves no worker or MLX-Serve behind (`localrouter guard`). + +Both produce correct edits with all five inputs. MLX-Serve is as fast for text to image, two to four times faster for +edits, and needs half the memory; it is also a server that cancels at the next step, where mflux is a library a worker +would have to host itself. diff --git a/tools/image/qwen-mlx/baseline.py b/tools/image/qwen-mlx/baseline.py new file mode 100644 index 0000000..b50e8f5 --- /dev/null +++ b/tools/image/qwen-mlx/baseline.py @@ -0,0 +1,173 @@ +#!/usr/bin/env python3 +"""The macOS Qwen-Image baseline: the real model through the worker, as the daemon drives it, with times and memory. + + ROOT/venv/bin/python -I tools/image/qwen-mlx/baseline.py ROOT OUT [--only t2i,edit1,edit2,edit5,mask] [--w8a8] + +ROOT is fetch.sh's; OUT a directory outside the repository (the images are not committed). The worker loads (timed), +then makes the fixed prompts and seeds below: 1024x1024 text to image, and edits with 1, 2 and 5 input images (the +inputs are themselves made here, so the baseline needs nothing from outside) and one with a mask. Every case runs +twice with two seeds. A sampler reads the server's resident set every 100 ms (unified memory: MLX's Metal buffers are +in it). OUT/baseline.json holds the numbers and OUT/-.png the images. Hold the machine's GPU lock while it +runs. docs/dev/IMAGE-EDIT.md, "macOS", has the results. +""" +import argparse +import json +import os +import platform +import shutil +import subprocess +import sys +import threading +import time + +HERE = os.path.dirname(os.path.abspath(__file__)) + +# The inputs, made first: (name, prompt, width, height, seed). Sizes differ so the edits see mixed shapes. +INPUTS = [ + ("woman", "a photo of a smiling woman with short black hair in a plain grey sweater, studio portrait, soft light", 1024, 1024, 11), + ("hat", "a red knitted wool beanie with a white pompom, product photo on a white background", 768, 768, 12), + ("cup", "a turquoise ceramic coffee cup with steam rising, product photo on a white background", 1024, 768, 13), + ("dog", "a golden retriever sitting on grass, looking at the camera, daylight photo", 768, 1024, 14), + ("room", "a cosy reading room with a green armchair, a bookshelf and a big window, warm evening light", 1344, 768, 15), +] +SEEDS = (7, 8) +# (case, prompt, inputs, size or None (the last input's shape), mask) +CASES = [ + ("t2i", "a red fox standing in fresh snow at golden hour, a wooden sign in the background that reads \"LocalRouter\"", [], (1024, 1024), False), + ("edit1", "make it night time: dark blue sky through the window light, moonlight on her face", ["woman"], None, False), + ("edit2", "the woman from image 1 wearing the red beanie from image 2, same portrait and lighting", ["woman", "hat"], None, False), + ("edit5", "the woman from image 1 wearing the beanie from image 2, sitting in the armchair of the room in image 5, " + "holding the cup from image 3, with the dog from image 4 at her feet", ["woman", "hat", "cup", "dog", "room"], None, False), + ("mask", "give her a bright yellow sweater", ["woman"], (1024, 1024), True), +] + + +class Sampler: + """The peak resident set of the worker's mlx-serve, sampled every 100 ms; `mark` starts a new window.""" + + def __init__(self, worker_pid): + self.worker_pid, self.pid = worker_pid, 0 + self.peak = self.window = 0 + self.stop = False + threading.Thread(target=self.run, daemon=True).start() + + def run(self): + while not self.stop: + if not self.pid: + out = subprocess.run(["pgrep", "-P", str(self.worker_pid), "mlx-serve"], capture_output=True, text=True).stdout.split() + self.pid = int(out[0]) if out else 0 + if self.pid: + out = subprocess.run(["ps", "-o", "rss=", "-p", str(self.pid)], capture_output=True, text=True).stdout.strip() + rss = int(out or 0) * 1024 + self.peak, self.window = max(self.peak, rss), max(self.window, rss) + time.sleep(0.1) + + def mark(self): + w, self.window = self.window, 0 + return w + + +class Worker: + def __init__(self, root, w8a8): + tool = {"id": "qwen-image-2.1-turbo-mlx", "kind": "image", "weights": os.path.join(root, "qwen-image-2.1-turbo-mlx-q8"), + "options": {"w8a8": w8a8}} + self.p = subprocess.Popen([os.path.join(root, "venv", "bin", "python"), "-I", os.path.join(HERE, "qwen_mlx_worker.py")], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, text=True, + env=dict(os.environ, LOCALROUTER_TOOL=json.dumps(tool))) + + def call(self, obj): + self.p.stdin.write(json.dumps(obj) + "\n") + self.p.stdin.flush() + while True: + m = json.loads(self.p.stdout.readline()) + if "progress" not in m: + return m + + +def gb(n): + return round(n / 1e9, 2) + + +def main(): + a = argparse.ArgumentParser(description=__doc__.split("\n")[0]) + a.add_argument("root") + a.add_argument("out") + a.add_argument("--only", default="t2i,edit1,edit2,edit5,mask") + a.add_argument("--w8a8", action="store_true", help="MLX-Serve's lossy int8 DiT linears (not the baseline)") + args = a.parse_args() + out = os.path.abspath(args.out) + os.makedirs(out, exist_ok=True) + only = set(args.only.split(",")) + report = {"machine": f"{platform.machine()} {platform.mac_ver()[0]}", "w8a8": args.w8a8, "cases": []} + + w = Worker(args.root, args.w8a8) + s = Sampler(w.p.pid) + t = time.monotonic() + r = w.call({"op": "load"}) + if not r.get("ok"): + sys.exit(f"load failed: {r}") + report["load_s"] = round(time.monotonic() - t, 2) + report["resident_bytes"] = r["resident"] + report["load_peak_rss"] = s.mark() + print(f"load {report['load_s']} s, resident {gb(r['resident'])} GB", flush=True) + + def run(case, prompt, refs, size, seed, mask=None): + job = os.path.join(out, "jobs", f"{case}-{seed}") + shutil.rmtree(job, ignore_errors=True) + os.makedirs(job) + names = [] + for i, ref in enumerate(refs): + shutil.copy(os.path.join(out, f"input-{ref}.png"), os.path.join(job, f"ref_{i}.png")) + names.append(f"ref_{i}.png") + img = {"prompt": prompt, "n": 1, "seed": seed, "references": names, + "width": size[0] if size else 0, "height": size[1] if size else 0} + if mask: + shutil.copy(mask, os.path.join(job, "mask_0.png")) + img["mask"] = "mask_0.png" + s.mark() + t = time.monotonic() + r = w.call({"op": "generate", "id": f"{case}-{seed}", "dir": job, "request": {"image": img}}) + secs = round(time.monotonic() - t, 2) + peak = s.mark() + if not r.get("ok"): + sys.exit(f"{case} seed {seed}: {r}") + return os.path.join(job, r["files"][0]), secs, peak + + if only - {"t2i"}: + for name, prompt, width, height, seed in INPUTS: + path, secs, peak = run(f"input-{name}", prompt, [], (width, height), seed) + shutil.copy(path, os.path.join(out, f"input-{name}.png")) + report["cases"].append({"case": f"input-{name}", "size": f"{width}x{height}", "seed": seed, "s": secs, "peak_rss": peak}) + print(f"input {name} {width}x{height}: {secs} s, peak {gb(peak)} GB", flush=True) + for case, prompt, refs, size, masked in CASES: + if case not in only: + continue + mask = None + if masked: + from PIL import Image # the worker's venv has it + m = Image.new("RGBA", (1024, 1024), (0, 0, 0, 255)) + m.paste((0, 0, 0, 0), (256, 560, 768, 1024)) # clear over the sweater: only there may it change + mask = os.path.join(out, "mask.png") + m.save(mask) + for seed in SEEDS: + path, secs, peak = run(case, prompt, refs, size, seed, mask) + dest = os.path.join(out, f"{case}-{seed}.png") + shutil.copy(path, dest) + report["cases"].append({"case": case, "inputs": refs, "seed": seed, "s": secs, "peak_rss": peak, "file": dest}) + print(f"{case} seed {seed}: {secs} s, peak {gb(peak)} GB -> {dest}", flush=True) + + t = time.monotonic() + w.p.stdin.write(json.dumps({"op": "exit"}) + "\n") # no reply: the worker stops the server and exits + w.p.stdin.flush() + w.p.wait(timeout=120) + report["exit_s"] = round(time.monotonic() - t, 2) + s.stop = True + report["peak_rss"] = s.peak + shutil.rmtree(os.path.join(out, "jobs"), ignore_errors=True) + with open(os.path.join(out, "baseline.json"), "w") as f: + json.dump(report, f, indent=1) + print(f"peak {gb(s.peak)} GB; {os.path.join(out, 'baseline.json')}") + + +if __name__ == "__main__": + main() diff --git a/tools/image/qwen-mlx/config.example.json b/tools/image/qwen-mlx/config.example.json new file mode 100644 index 0000000..18d8c9f --- /dev/null +++ b/tools/image/qwen-mlx/config.example.json @@ -0,0 +1,26 @@ +{ + "host": "127.0.0.1", + "defaults": {"text_to_image": "qwen-image-2.1-turbo", "image_edit": "qwen-image-2.1-turbo"}, + "tools": [ + { + "id": "qwen-image-2.1-turbo", + "kind": "image", + "name": "Qwen-Image 2.1 Turbo (macOS, MLX-Serve 4-bit, 8 steps)", + "capabilities": ["text_to_image", "image_edit"], + "cmd": ["/opt/localrouter-mac/image/venv/bin/python", "-I", "/opt/localrouter/tools/image/qwen-mlx/qwen_mlx_worker.py"], + "weights": "/opt/localrouter-mac/image/qwen-image-2.1-turbo-mlx-q4", + "license": {"name": "Qwen Research License Agreement", "url": "https://huggingface.co/Qwen/Qwen-Image-2.1/blob/main/LICENSE", + "use": "non-commercial: research or evaluation only; commercial use needs a separate licence from Qwen", + "commercial": false}, + "resident_bytes": 10200547328, + "working_bytes": 6979321856, + "keep_loaded": false, + "priority": 10, + "idle_ttl_s": 60, + "load_timeout_s": 600, + "generate_timeout_s": 600, + "options": {"pack": "q4", "max_side": 1360, "max_pixels": 1048576, "edit_refs": 5, + "reference_working_bytes": 2147483648} + } + ] +} diff --git a/tools/image/qwen-mlx/fetch.sh b/tools/image/qwen-mlx/fetch.sh new file mode 100755 index 0000000..ff625cd --- /dev/null +++ b/tools/image/qwen-mlx/fetch.sh @@ -0,0 +1,69 @@ +#!/usr/bin/env bash +# Downloads what the macOS Qwen-Image worker runs (tools/image/qwen-mlx/pins-macos.json) into ROOT and checks it. +# tools/image/qwen-mlx/fetch.sh ROOT [--pack q8|q4] +# ROOT/mlx-serve-macos-arm64/ the MLX-Serve release binary (its tarball's sha256 checked; quarantine cleared) +# ROOT/qwen-image-2.1-turbo-mlx-PACK/ the MLX-Serve pack, q8 (8-bit, 17.7 GB, pins-macos.json; the default) or q4 +# (4-bit, 10.5 GB, pins-macos-q4.json); every file's size and digest checked by +# tools/package/verify_assets.py: the tool's `weights` (and its option `pack`) +# ROOT/venv/ Python for the worker, with Pillow from requirements-macos.txt (hashes required) +# Files already present with the pinned size are kept (and checked again). Apple Silicon macOS only. The model is under +# the Qwen Research License: non-commercial research or evaluation only. +set -euo pipefail +here=$(cd "$(dirname "$0")" && pwd) +repo=$(cd "$here/../../.." && pwd) +root=${1:?usage: fetch.sh ROOT [--pack q8|q4]} +pack=q8 +[ "${2:-}" = --pack ] && pack=${3:?--pack q8 or q4} +case $pack in + q8) packpins=$here/pins-macos.json ;; + q4) packpins=$here/pins-macos-q4.json ;; + *) echo "fetch.sh: --pack q8 or q4, not $pack" >&2; exit 2 ;; +esac +[ "$(uname -s)/$(uname -m)" = Darwin/arm64 ] || { echo "fetch.sh: macOS on Apple Silicon only" >&2; exit 1; } +mkdir -p "$root" +root=$(cd "$root" && pwd) +pins=$here/pins-macos.json +pin() { python3 -I -c 'import json,sys; v=json.load(open(sys.argv[1])); [v := v[k] for k in sys.argv[2:]]; print(v)' "$@"; } # FILE KEY... + +# MLX-Serve +url=$(pin "$pins" mlx_serve url) digest=$(pin "$pins" mlx_serve sha256) dir=$(pin "$pins" mlx_serve dir) version=$(pin "$pins" mlx_serve version) +if ! "$root/$dir/mlx-serve" --version 2>/dev/null | grep -q "^mlx-serve $version\$"; then + tmp="$root/.mlx-serve.tar.gz.part" + curl -fsSL --retry 3 -o "$tmp" "$url" + got=$(shasum -a 256 "$tmp" | cut -d' ' -f1) + if [ "$got" != "$digest" ]; then + rm -f "$tmp" + echo "fetch.sh: $url sha256 $got, pinned $digest" >&2 + exit 1 + fi + rm -rf "$root/$dir" + tar -xzf "$tmp" -C "$root" + rm -f "$tmp" + xattr -dr com.apple.quarantine "$root/$dir" 2>/dev/null || true +fi +echo "fetch.sh: $("$root/$dir/mlx-serve" --version 2>/dev/null | grep '^mlx-serve ') in $root/$dir" + +# The pack: each pinned file from its revision, then every digest +repo_id=$(pin "$packpins" repo_id) rev=$(pin "$packpins" revision) +weights=$root/qwen-image-2.1-turbo-mlx-$pack +python3 -I -c 'import json,sys; [print(f["path"], f["size_bytes"]) for f in json.load(open(sys.argv[1]))["selected_files"]]' "$packpins" | +while read -r path size; do + dest=$weights/$path + if [ -f "$dest" ] && [ "$(stat -f %z "$dest")" = "$size" ]; then continue; fi + mkdir -p "$(dirname "$dest")" + echo "fetch.sh: $repo_id/$path ($size bytes)" + curl -fSL --retry 3 -C - -o "$dest.part" "https://huggingface.co/$repo_id/resolve/$rev/$path" + mv "$dest.part" "$dest" +done +python3 -I "$repo/tools/package/verify_assets.py" "$packpins" "$weights" + +# The worker's Python +if [ ! -x "$root/venv/bin/python" ]; then + if command -v uv >/dev/null; then uv venv -q -p 3.12 "$root/venv"; else python3 -m venv "$root/venv"; fi +fi +if command -v uv >/dev/null; then + VIRTUAL_ENV=$root/venv uv pip install -q --require-hashes -r "$here/requirements-macos.txt" +else + "$root/venv/bin/python" -m pip install -q --require-hashes -r "$here/requirements-macos.txt" +fi +echo "fetch.sh: done. The tool: \"cmd\": [\"$root/venv/bin/python\", \"-I\", \"$here/qwen_mlx_worker.py\"], \"weights\": \"$weights\", \"options\": {\"pack\": \"$pack\"}" diff --git a/tools/image/qwen-mlx/pins-macos-q4.json b/tools/image/qwen-mlx/pins-macos-q4.json new file mode 100644 index 0000000..feb9d61 --- /dev/null +++ b/tools/image/qwen-mlx/pins-macos-q4.json @@ -0,0 +1,87 @@ +{ + "note": "macOS only: the 4-bit alternative to pins-macos.json's 8-bit pack, for a Mac whose warm models leave no room for 18 GiB (fetch.sh --pack q4; the worker's option pack: q4). The runtime is pins-macos.json's. tools/package/verify_assets.py form (selected_files).", + "repo_id": "ddalcu/Qwen-Image-2.1-Turbo-MLX-Serve-4bit", + "revision": "394f9477a956a0c4fd70ebcbe3a326d74ef0add4", + "base_model": "Qwen/Qwen-Image-2.1-Turbo", + "license": { + "name": "Qwen Research License Agreement: non-commercial research or evaluation only", + "url": "https://huggingface.co/Qwen/Qwen-Image-2.1/blob/main/LICENSE" + }, + "selected_files": [ + { + "path": "LICENSE", + "size_bytes": 7831, + "git_blob_sha1": "13ae08d5a5828f508cf0e852b1b316c0c0bbb9b3" + }, + { + "path": "NOTICE", + "size_bytes": 434, + "git_blob_sha1": "853eb90e79b0ddea49f7edc11c12da326bbea664" + }, + { + "path": "README.md", + "size_bytes": 4137, + "git_blob_sha1": "e70fbf2fe18f0769444df05fd04e44489b8750a3" + }, + { + "path": "config.json", + "size_bytes": 117, + "git_blob_sha1": "623ef1e3fa7bd18ff495f114c4357d142c7ff1ed" + }, + { + "path": "model_index.json", + "size_bytes": 578, + "git_blob_sha1": "cdc1dae2997ff637225c42c388196b38a9bbaae8" + }, + { + "path": "processor/tokenizer.json", + "size_bytes": 11422650, + "lfs_sha256": "be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506" + }, + { + "path": "processor/tokenizer_config.json", + "size_bytes": 705, + "git_blob_sha1": "6577e86ca277ad0367551e280ba1d58ad4d42c52" + }, + { + "path": "scheduler/scheduler_config.json", + "size_bytes": 486, + "git_blob_sha1": "72837792c1944182b3707415cdaeb825e6bf5598" + }, + { + "path": "text_encoder/config.json", + "size_bytes": 1542, + "git_blob_sha1": "b5881fa79f7c988a9ac2f91e4f1192deed7ac034" + }, + { + "path": "text_encoder/model.safetensors", + "size_bytes": 5676130033, + "lfs_sha256": "6d326352bd9468e77c3788120b5cd7c50beafec379578aa5c172819c2f31e746" + }, + { + "path": "transformer/config.json", + "size_bytes": 370, + "git_blob_sha1": "b3f8d1cf08a94a1d4ec09e2bfeb5344af25f2983" + }, + { + "path": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors", + "size_bytes": 2974367042, + "lfs_sha256": "11b3e959d1b9031d15d3b8fabb7b13b8cac571c3cdf22452960523a9c3adcb81" + }, + { + "path": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors", + "size_bytes": 1223193789, + "lfs_sha256": "242eb483f30a10446758255694a23a518c171cd544303a0d681aaf4c9fdb9a33" + }, + { + "path": "vae/config.json", + "size_bytes": 2172, + "git_blob_sha1": "9b780c2dd9601cb0b93ae8f462b58dd3a7468347" + }, + { + "path": "vae/diffusion_pytorch_model.safetensors", + "size_bytes": 662000745, + "lfs_sha256": "3ecb67f7df44bef3a74a65fa2a90d477122e45c7214906f04222a59ca4ca8aeb" + } + ] +} diff --git a/tools/image/qwen-mlx/pins-macos.json b/tools/image/qwen-mlx/pins-macos.json new file mode 100644 index 0000000..6044f40 --- /dev/null +++ b/tools/image/qwen-mlx/pins-macos.json @@ -0,0 +1,96 @@ +{ + "note": "macOS (Apple Silicon) only: the Qwen-Image 2.1 Turbo media worker (qwen_mlx_worker.py). fetch.sh downloads and checks these; the model part is in tools/package/verify_assets.py form (selected_files).", + "mlx_serve": { + "version": "26.10.3", + "url": "https://github.com/ddalcu/mlx-serve/releases/download/v26.10.3/mlx-serve-bin-macos-arm64.tar.gz", + "sha256": "d1781c89c206f0d077d6d1abbcab43aa9ec65280202c83eac908f957956b652b", + "source_commit": "046b647686419df06c6c7132fb1d15d76185c3e5", + "dir": "mlx-serve-macos-arm64", + "license": "MIT (LICENSE), bundled MLX and libraries Apache-2.0 (LICENSE-APACHE-2.0, NOTICE)", + "min_version": "26.10.2" + }, + "repo_id": "ddalcu/Qwen-Image-2.1-Turbo-MLX-Serve-8bit", + "revision": "214ca346520a59597b89bd77408afa49d526bde5", + "base_model": "Qwen/Qwen-Image-2.1-Turbo", + "license": { + "name": "Qwen Research License Agreement: non-commercial research or evaluation only", + "url": "https://huggingface.co/Qwen/Qwen-Image-2.1/blob/main/LICENSE" + }, + "selected_files": [ + { + "path": "LICENSE", + "size_bytes": 7831, + "git_blob_sha1": "13ae08d5a5828f508cf0e852b1b316c0c0bbb9b3" + }, + { + "path": "NOTICE", + "size_bytes": 434, + "git_blob_sha1": "64ce1369d01ae65234d232e24247a46d2690ac83" + }, + { + "path": "README.md", + "size_bytes": 4086, + "git_blob_sha1": "c6885534c89d452982efa92197253a68e2366e56" + }, + { + "path": "config.json", + "size_bytes": 117, + "git_blob_sha1": "8e2235069c6d88a8e614e97011cc44771badc283" + }, + { + "path": "model_index.json", + "size_bytes": 578, + "git_blob_sha1": "cdc1dae2997ff637225c42c388196b38a9bbaae8" + }, + { + "path": "processor/tokenizer.json", + "size_bytes": 11422650, + "lfs_sha256": "be75606093db2094d7cd20f3c2f385c212750648bd6ea4fb2bf507a6a4c55506" + }, + { + "path": "processor/tokenizer_config.json", + "size_bytes": 705, + "git_blob_sha1": "6577e86ca277ad0367551e280ba1d58ad4d42c52" + }, + { + "path": "scheduler/scheduler_config.json", + "size_bytes": 486, + "git_blob_sha1": "72837792c1944182b3707415cdaeb825e6bf5598" + }, + { + "path": "text_encoder/config.json", + "size_bytes": 1542, + "git_blob_sha1": "b5881fa79f7c988a9ac2f91e4f1192deed7ac034" + }, + { + "path": "text_encoder/model.safetensors", + "size_bytes": 9367829459, + "lfs_sha256": "4db3757e59b070b827e1285e506df8874527206bcbc708fdcf02b79bb9a915a2" + }, + { + "path": "transformer/config.json", + "size_bytes": 370, + "git_blob_sha1": "b3f8d1cf08a94a1d4ec09e2bfeb5344af25f2983" + }, + { + "path": "transformer/diffusion_pytorch_model-00001-of-00002.safetensors", + "size_bytes": 5407063648, + "lfs_sha256": "cff3e7601fa701d189ce6b013a3f7ca7a318ba49c8b4c8b820faf443adeb1ec8" + }, + { + "path": "transformer/diffusion_pytorch_model-00002-of-00002.safetensors", + "size_bytes": 2280158619, + "lfs_sha256": "1a4117365d8e6a13b72d71abd13f2a3c60edcb9128a44c1018d0128182bf3728" + }, + { + "path": "vae/config.json", + "size_bytes": 2172, + "git_blob_sha1": "9b780c2dd9601cb0b93ae8f462b58dd3a7468347" + }, + { + "path": "vae/diffusion_pytorch_model.safetensors", + "size_bytes": 662000745, + "lfs_sha256": "3ecb67f7df44bef3a74a65fa2a90d477122e45c7214906f04222a59ca4ca8aeb" + } + ] +} diff --git a/tools/image/qwen-mlx/qwen_mlx_worker.py b/tools/image/qwen-mlx/qwen_mlx_worker.py new file mode 100644 index 0000000..b95a671 --- /dev/null +++ b/tools/image/qwen-mlx/qwen_mlx_worker.py @@ -0,0 +1,565 @@ +#!/usr/bin/env python3 +"""LocalRouter worker for Qwen-Image 2.1 Turbo on macOS (Apple Silicon): text to image and image edits on MLX. + +Speaks the worker protocol (docs/ARCHITECTURE.md): one JSON object a line on stdin and stdout, logs on stderr. + + {"op":"load"} -> {"ok":true,"resident":,"cancel":"signal"} + {"op":"generate","id","dir","request":{"image":{"prompt","width","height","n","seed","steps", + "references":["ref_0.png",…],"mask":"mask_0.png"}}} + -> {"ok":true,"files":["0.png",…],"seed":7,"ms":…} + SIGUSR1 during a generate -> {"ok":false,"type":"cancelled"} (still loaded) + {"op":"exit"} + +The model runs in MLX-Serve (github.com/ddalcu/mlx-serve, a native Zig/MLX server; release and pack pinned in +pins-macos.json, fetched by fetch.sh), which this worker starts on load on 127.0.0.1 and is the client of: it is the +"models that wrap their own server" shape of docs/ARCHITECTURE.md. The pack is ddalcu/Qwen-Image-2.1-Turbo-MLX-Serve-8bit +(MLX affine 8-bit DiT and text-encoder linears, dense f32 VAE, the Qwen3-VL vision tower kept for edits), or its 4-bit +sibling ddalcu/Qwen-Image-2.1-Turbo-MLX-Serve-4bit (option `pack`: q4; the same schedule, tower and VAE, pins-macos-q4.json). It runs the +checkpoint's own eight-step `sample_sigmas` schedule (unshifted), which MLX-Serve reads from model_index.json since +26.10.2; an older server ran the base model's schedule on it, so load refuses one. + +What the worker owns, so the daemon's request means the same here as on the Spark engine (src/engines/qwen_image_tool.zig): + +- Limits: `max_side` (default 1664) and `max_pixels` (default max_side squared) on the output, up to `edit_refs` input + images (default 5, at most 5), steps 0 or the schedule's 8. Anything else is refused (invalid_request_error), never + changed quietly. `guidance` and `negative_prompt` are not used, as on the Spark engine (Turbo runs without CFG). +- n images: one request a seed, `seed`, `seed+1`, … (`0.png`, `1.png`, …). +- Edits: input images in order (the first is the one edited, the others guide it), sent as MLX-Serve's `image` and + `ref_images`. Without a size, the last image's shape at about one megapixel within the limits (the Spark tool's + `editSize`). MLX-Serve floors an explicit edit size to multiples of 32, so other sizes are refused. Each input is + resized to about one megapixel for conditioning; one past 4:1 or under 32 pixels a side is refused. WebP inputs + are converted to PNG losslessly (MLX-Serve decodes PNG and JPEG). +- A mask (alpha 0 where the edit goes, as OpenAI's) keeps the first image wherever it is opaque: both are resized to + the output (Lanczos, premultiplied for RGBA) and the mask's alpha blends the two (the Spark tool's `Keep`). +- Output: MLX-Serve's PNG with `prompt` and `seed` text chunks added (pixels untouched); the daemon adds the + machine-generated label. RGBA when the option `alpha` is on (default: the VAE's own alpha, as the Spark engine). +- Memory: `resident` is what MLX-Serve holds once loaded and warmed (the larger of MLX's active bytes and the server's + resident set). LocalRouter's scheduler is the only admission authority, so the server's own load pre-flight is off + (--skip-mem-preflight); its per-request working-set check still refuses an edit that cannot fit (a 503 here). +- Cancel: SIGUSR1 closes the connection; MLX-Serve stops at the next denoise step and stays loaded. +- A server that dies takes the worker with it (no reply): the job fails with a 503 and the next request starts it again. + The server is in the worker's process group, so the daemon's kill of the group stops it too. + +Options (the tool's `options`): `pack` (q8 | q4: which pinned pack `weights` holds, default q8), `server` (the mlx-serve binary; default /../mlx-serve-macos-arm64/mlx-serve), +`max_side`, `max_pixels`, `edit_refs`, `steps` (8), `alpha` (true), `w8a8` (false: MLX-Serve's lossy int8 x int8 DiT +linears stay off, the documented Mac baseline), `warmup` (true: a 256x256 image at load, so the first request is +not slower), `port` (0: a free one), `load_timeout_s` (600), `verify` (size: every pinned file present with its size; +none). `weights` is the pack's directory (fetch.sh). +""" + +import base64 +import http.client +import io +import json +import math +import os +import platform +import re +import signal +import socket +import struct +import subprocess +import sys +import threading +import time +import zlib + +from PIL import Image + +HERE = os.path.dirname(os.path.abspath(__file__)) +# each pack's pinned files (verify_assets.py form); the runtime's pin is in pins-macos.json +PACKS = {"q8": os.path.join(HERE, "pins-macos.json"), "q4": os.path.join(HERE, "pins-macos-q4.json")} + +MAX_REFS = 5 # LocalRouter's API limit (request.Limits.max_references) +REF_PIXELS = 1024 * 1024 # each input is conditioned at about this many pixels (the pipeline's calculate_dimensions) +REF_SIDE = 2048 # the longest conditioned side: inputs past 4:1 are refused +MIN_VERSION = (26, 10, 2) # the first MLX-Serve that runs the pack's own sample_sigmas schedule +OPTIONS = {"server", "pack", "max_side", "max_pixels", "edit_refs", "steps", "alpha", "w8a8", "warmup", "port", + "load_timeout_s", "verify"} +# options the daemon reads and the worker leaves alone: not reported as unused +DAEMON_OPTIONS = {"reference_working_bytes"} # the scheduler books it per input image of an edit (src/tool/process.zig) + + +def log(msg): + print(f"qwen-mlx: {msg}", file=sys.stderr, flush=True) + + +class Refused(Exception): + """A request this model cannot serve: invalid_request_error (the API's 400).""" + + +class Unavailable(Exception): + """The server is busy or short of memory for this request: service_unavailable (503), retry later.""" + + +class Cancelled(BaseException): + """SIGUSR1 during a generate. Not an Exception, so no handler turns it into a server error.""" + + +GENERATING = False + + +def on_cancel(signum, frame): + if GENERATING: + raise Cancelled() + + +# -- sizes, as the Spark tool computes them (src/engines/qwen_image_tool.zig, qwen_image/edit_prepare.zig) -------------- + +def dimensions(width, height, area): + """calculate_dimensions: the shape of width x height at about `area` pixels, sides multiples of 32 (round half even).""" + if width <= 0 or height <= 0 or area <= 0: + raise ValueError("invalid dimensions") + ratio = width / height + w = round(math.sqrt(area * ratio) / 32) * 32 + h = round(w / ratio / 32) * 32 + if w < 32 or h < 32: + raise ValueError("empty grid") + return w, h + + +def edit_size(width, height, max_side, max_pixels): + """An edit without a size: the last input's shape at about one megapixel, within max_side and max_pixels.""" + area = min(max_pixels, REF_PIXELS) + while area >= 32 * 32: + w, h = dimensions(width, height, area) + if max(w, h) <= max_side and w * h <= max_pixels: + return w, h + area -= 32 * 32 + raise Refused("no size this model takes keeps the input's shape: send size") + + +# -- PNG text chunks -------------------------------------------------------------------------------------------------- + +def add_text(png, items): + """The PNG with tEXt chunks (key, value) inserted right after IHDR; every other byte is kept.""" + if png[:8] != b"\x89PNG\r\n\x1a\n" or png[12:16] != b"IHDR": + raise ValueError("not a PNG") + end = 8 + 8 + struct.unpack(">I", png[8:12])[0] + 4 + chunks = b"" + for key, value in items: + body = key.encode("latin-1") + b"\0" + value.encode("utf-8") + chunks += struct.pack(">I", len(body)) + b"tEXt" + body + chunks += struct.pack(">I", zlib.crc32(b"tEXt" + body) & 0xFFFFFFFF) + return png[:end] + chunks + png[end:] + + +# -- the worker ---------------------------------------------------------------------------------------------------- + +class Options: + def __init__(self, tool): + o = tool.get("options") or {} + unknown = sorted(set(o) - OPTIONS - DAEMON_OPTIONS) + if unknown: + log(f"options not used by the worker: {unknown}") + self.id = tool.get("id") or "qwen-image-2.1-turbo" + self.weights = tool.get("weights") or "" + if not self.weights: + raise ValueError("weights (the MLX-Serve pack's directory) is required") + self.server = o.get("server") or os.path.join(os.path.dirname(os.path.abspath(self.weights)), + "mlx-serve-macos-arm64", "mlx-serve") + self.max_side = int(o.get("max_side", 1664)) // 16 * 16 + self.max_pixels = min(int(o.get("max_pixels", self.max_side * self.max_side)), self.max_side * self.max_side) // 256 * 256 + self.edit_refs = min(int(o.get("edit_refs", MAX_REFS)), MAX_REFS) + self.steps = int(o.get("steps", 8)) + self.alpha = bool(o.get("alpha", True)) + self.w8a8 = bool(o.get("w8a8", False)) + self.warmup = bool(o.get("warmup", True)) + self.port = int(o.get("port", 0)) + self.load_timeout_s = float(o.get("load_timeout_s", 600)) + self.pack = o.get("pack", "q8") + if self.pack not in PACKS: + raise ValueError(f"pack must be one of {sorted(PACKS)}") + self.verify = o.get("verify", "size") + if self.verify not in ("size", "none"): + raise ValueError("verify must be size or none") + if self.max_side < 256 or self.edit_refs < 0 or self.steps <= 0: + raise ValueError("max_side must be at least 256, edit_refs and steps positive") + + +class Worker: + def __init__(self, tool): + self.opts = Options(tool) + self.proc = None + self.port = 0 + self.conn = None + + # -- load ---------------------------------------------------------------------------------------------------- + + def load(self): + try: + self._check_host() + self._check_weights() + self._check_server() + self._start() + if self.opts.warmup: + t = time.monotonic() + self._image({"size": "256x256", "seed": 0, "steps": self.opts.steps, "prompt": "warm up"}, None) + log(f"warm-up image in {time.monotonic() - t:.1f} s") + resident = self._resident() + log(f"loaded: resident {resident >> 20} MiB") + return {"ok": True, "resident": resident, "cancel": "signal"} + except (Refused, ValueError, OSError, RuntimeError) as e: + log(f"load failed: {e}") + self._stop() + return {"ok": False, "error": str(e), "type": "server_error"} + + def _check_host(self): + if sys.platform != "darwin" or platform.machine() != "arm64": + raise RuntimeError("this worker runs on macOS on Apple Silicon only (MLX)") + + def _check_weights(self): + """The pinned files are there with their sizes (fetch.sh checked their digests), and the schedule has `steps` nodes.""" + w = self.opts.weights + if self.opts.verify == "size": + pins = read_json(PACKS[self.opts.pack]) + for f in pins["selected_files"]: + p = os.path.join(w, f["path"]) + if not os.path.isfile(p) or os.path.getsize(p) != f["size_bytes"]: + raise RuntimeError(f"{p}: missing or not the pinned size ({f['size_bytes']} bytes): run fetch.sh") + try: + sigmas = read_json(os.path.join(w, "model_index.json")).get("sample_sigmas") + except (OSError, ValueError) as e: + raise RuntimeError(f"{w}/model_index.json: {e}") + if not isinstance(sigmas, list) or len(sigmas) != self.opts.steps: + raise RuntimeError(f"{w}: the pack's sample_sigmas has {len(sigmas) if isinstance(sigmas, list) else 0} " + f"steps, the tool's steps option {self.opts.steps}") + + def _check_server(self): + try: + out = subprocess.run([self.opts.server, "--version"], capture_output=True, text=True, timeout=30).stdout + except (OSError, subprocess.TimeoutExpired) as e: + raise RuntimeError(f"{self.opts.server}: {e}: run fetch.sh, or set the option server") + m = re.search(r"mlx-serve (\d+)\.(\d+)\.(\d+)", out) + if not m: + raise RuntimeError(f"{self.opts.server} --version said {out.strip()!r}, not an MLX-Serve version") + version = tuple(int(x) for x in m.groups()) + if version < MIN_VERSION: + raise RuntimeError(f"MLX-Serve {'.'.join(map(str, version))} runs the base model's schedule on Turbo: " + f"{'.'.join(map(str, MIN_VERSION))} or newer is needed (fetch.sh)") + + def _start(self): + self.port = self.opts.port or free_port() + argv = [self.opts.server, "--serve", "--model", self.opts.weights, "--host", "127.0.0.1", + "--port", str(self.port), "--skip-mem-preflight", "--w8a8" if self.opts.w8a8 else "--no-w8a8"] + log(f"starting {' '.join(argv)}") + t = time.monotonic() + # stdout is the protocol: the server's output goes to the log. Same process group as the worker. + self.proc = subprocess.Popen(argv, stdin=subprocess.DEVNULL, stdout=sys.stderr, stderr=sys.stderr) + deadline = t + self.opts.load_timeout_s + while True: + if self.proc.poll() is not None: + raise RuntimeError(f"mlx-serve exited with {self.proc.returncode} while loading (see the log)") + if time.monotonic() > deadline: + raise RuntimeError(f"mlx-serve was not ready within {self.opts.load_timeout_s:.0f} s") + try: + status, body = self._get("/v1/models") + if status == 200 and any(m.get("loaded") and m.get("state") == "ready" for m in json.loads(body).get("data", [])): + break + except (OSError, ValueError, http.client.HTTPException): + pass + time.sleep(0.25) + log(f"mlx-serve ready on port {self.port} in {time.monotonic() - t:.1f} s") + + def _resident(self): + active = 0 + try: + status, body = self._get("/props") + if status == 200: + active = int(json.loads(body).get("memory", {}).get("active_bytes", 0)) + except (OSError, ValueError, http.client.HTTPException): + pass + return max(active, rss(self.proc.pid)) + + # -- generate ------------------------------------------------------------------------------------------------ + + def generate(self, cmd, progress): + start = time.monotonic() + req = (cmd.get("request") or {}).get("image") + if req is None: + raise Refused("this model makes images: send an image request") + job = cmd.get("dir") or "" + o = self.opts + refs = req.get("references") or [] + mask = req.get("mask") + width, height = int(req.get("width", 1024)), int(req.get("height", 1024)) + n, seed, steps = int(req.get("n", 1)), int(req.get("seed", 0)), int(req.get("steps", 0)) + prompt = req.get("prompt") or "" + + if len(refs) > o.edit_refs: + raise Refused(f"model '{o.id}' edits with up to {o.edit_refs} input image{'' if o.edit_refs == 1 else 's'}") + if mask is not None and not refs: + raise Refused("a mask needs an input image to edit") + if steps not in (0, o.steps): + raise Refused(f"model '{o.id}' runs its checkpoint's {o.steps}-step schedule: omit steps or send {o.steps}") + if not prompt: + raise Refused("prompt is required") + if n < 1 or n > 4: + raise Refused("n must be 1 to 4") + + inputs = [self._input(job, name) for name in refs] + if width == 0 and height == 0: + if not inputs: + raise Refused(f"model '{o.id}' needs a size") + width, height = edit_size(*inputs[-1][1].size, o.max_side, o.max_pixels) + if max(width, height) > o.max_side: + raise Refused(f"model '{o.id}' takes sizes up to {o.max_side} pixels a side") + if width * height > o.max_pixels: + raise Refused(f"model '{o.id}' takes at most {o.max_pixels} pixels (width x height), such as 1024x1024 or 1360x768") + if width % 16 or height % 16 or width < 256 or height < 256: + raise Refused(f"size {width}x{height}: each side must be a multiple of 16, at least 256") + if inputs and (width % 32 or height % 32): + raise Refused(f"size {width}x{height}: an edit's sides must be multiples of 32 on this model (MLX-Serve)") + keep = Keep(job, mask, inputs[0][1], width, height) if mask is not None else None + + body = {"size": f"{width}x{height}", "steps": o.steps} + if o.alpha: + body["transparent"] = True + if inputs: + body["mode"] = "edit" + body["image"] = inputs[0][0] + if len(inputs) > 1: + body["ref_images"] = [b for b, _ in inputs[1:]] + files = [] + for i in range(n): + s = (seed + i) % (1 << 64) + # the prompt last: MLX-Serve finds the other fields by their first key match in the body + png = self._image({**body, "seed": s, "prompt": prompt}, progress) + if keep is not None: + png = keep.apply(png) + name = f"{i}.png" + with open(os.path.join(job, name), "wb") as f: + f.write(add_text(png, [("prompt", prompt), ("seed", str(s))])) + files.append(name) + return {"ok": True, "files": files, "seed": seed, "ms": int((time.monotonic() - start) * 1000)} + + def _input(self, job, name): + """An input image as MLX-Serve takes it (base64 PNG or JPEG) and as Pillow reads it, after the size checks.""" + path = os.path.join(job, name) + if os.path.dirname(os.path.normpath(name)) or name.startswith("."): + raise Refused("references must name files in the job directory") + try: + with open(path, "rb") as f: + data = f.read() + im = Image.open(io.BytesIO(data)) + im.load() + except (OSError, ValueError, Image.DecompressionBombError) as e: + raise Refused(f"an input image could not be read ({e}): send a PNG, JPEG or WebP") + try: + if min(im.size) < 32: + raise ValueError("too small") + w, h = dimensions(im.width, im.height, REF_PIXELS) + except ValueError: + raise Refused(f"an input image of {im.width}x{im.height} cannot be edited: send one with sides of at least 32 pixels") + if max(w, h) > REF_SIDE: + raise Refused(f"an input image of {im.width}x{im.height} is too long and thin to edit: at most 4:1") + if im.format not in ("PNG", "JPEG"): + buf = io.BytesIO() + im.save(buf, "PNG") # lossless, alpha kept + data = buf.getvalue() + return base64.b64encode(data).decode("ascii"), im + + def _image(self, body, progress): + """One image from MLX-Serve: its PNG bytes. Streams progress; a cancel closes the connection.""" + conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=None) + self.conn = conn + try: + # streamed: progress events, and a cancel MLX-Serve sees at its next step + conn.request("POST", "/v1/images/generations", body=json.dumps({"stream": True, **body}), + headers={"Content-Type": "application/json"}) + r = conn.getresponse() + if r.status != 200: + msg = server_message(r.read()) + if r.status == 400 and "GPU memory" in msg: + raise Unavailable(msg) + if r.status == 400: + raise Refused(msg) + if r.status == 503: + raise Unavailable(msg) + raise RuntimeError(f"mlx-serve answered {r.status}: {msg}") + if not r.getheader("Content-Type", "").startswith("text/event-stream"): + return base64.b64decode(json.loads(r.read())["data"][0]["b64_json"]) + for raw in r: + line = raw.strip() + if not line.startswith(b"data:"): + continue + ev = json.loads(line[5:]) + kind = ev.get("type") + if kind == "progress" and progress is not None: + progress({"phase": "denoise", "step": ev.get("step", 0), "of": ev.get("total", 0)}) + elif kind == "complete": + return base64.b64decode(ev["data"][0]["b64_json"]) + elif kind == "error": + raise RuntimeError(f"mlx-serve: {ev.get('message', 'generation failed')}") + raise RuntimeError("mlx-serve ended the stream without an image") + finally: + self.conn = None + conn.close() + + def _get(self, path): + conn = http.client.HTTPConnection("127.0.0.1", self.port, timeout=10) + try: + conn.request("GET", path) + r = conn.getresponse() + return r.status, r.read() + finally: + conn.close() + + def server_alive(self): + return self.proc is not None and self.proc.poll() is None + + # -- exit ------------------------------------------------------------------------------------------------------ + + def _stop(self): + p, self.proc = self.proc, None + if p is None or p.poll() is not None: + return + p.terminate() + try: + p.wait(timeout=30) + except subprocess.TimeoutExpired: + log("mlx-serve did not stop in 30 s; killing it") + p.kill() + p.wait() + + +class Keep: + """A mask's composite: the first input at the output's size, kept wherever the mask is opaque (OpenAI's edits: + alpha 0 marks where the edit goes). The mask must be the first image's size.""" + + def __init__(self, job, name, first, width, height): + path = os.path.join(job, name) + if os.path.dirname(os.path.normpath(name)) or name.startswith("."): + raise Refused("mask must name a file in the job directory") + try: + m = Image.open(path) + m.load() + except (OSError, ValueError, Image.DecompressionBombError) as e: + raise Refused(f"the mask could not be used ({e}): send a PNG of the first image's size, transparent where the edit goes") + if m.size != first.size: + raise Refused(f"the mask is {m.width}x{m.height}, the first image {first.width}x{first.height}: " + "send a PNG of the first image's size, transparent where the edit goes") + self.original = first.convert("RGBA").resize((width, height), Image.LANCZOS) # Pillow premultiplies RGBA + self.alpha = m.convert("RGBA").getchannel("A").resize((width, height), Image.LANCZOS) + + def apply(self, png): + edited = Image.open(io.BytesIO(png)) + edited.load() + original = self.original if edited.mode == "RGBA" else self.original.convert(edited.mode) + out = Image.composite(original, edited, self.alpha) + buf = io.BytesIO() + out.save(buf, "PNG", compress_level=1) + return buf.getvalue() + + +def read_json(path): + with open(path) as f: + return json.load(f) + + +def server_message(raw): + try: + e = json.loads(raw).get("error") + return (e.get("message") if isinstance(e, dict) else str(e)) or "error" + except (ValueError, AttributeError): + return raw.decode("utf-8", "replace").strip() or "error" + + +def free_port(): + with socket.socket() as s: + s.bind(("127.0.0.1", 0)) + return s.getsockname()[1] + + +def rss(pid): + """A process's resident set in bytes (unified memory: MLX's Metal buffers are in it), 0 when unreadable.""" + try: + out = subprocess.run(["ps", "-o", "rss=", "-p", str(pid)], capture_output=True, text=True, timeout=5).stdout + return int(out.strip() or 0) * 1024 + except (OSError, ValueError, subprocess.TimeoutExpired): + return 0 + + +def send(obj): + sys.stdout.write(json.dumps(obj) + "\n") + sys.stdout.flush() + + +def watch_parent(worker): + """Stop the server and exit when the daemon that started us is gone (LOCALROUTER_PARENT; 0 or unset: no check).""" + parent = int(os.environ.get("LOCALROUTER_PARENT", "0") or 0) + if parent <= 0: + return + + def run(): + while True: + time.sleep(2) + if os.getppid() != parent: + log("the daemon is gone; stopping") + worker._stop() + os._exit(0) + + threading.Thread(target=run, daemon=True).start() + + +def generate(worker, cmd): + global GENERATING + try: + try: + GENERATING = True # inside the try: a signal right after it is still caught + return worker.generate(cmd, lambda p: send({"progress": p})) + finally: + GENERATING = False + except Cancelled: + log("generate cancelled; mlx-serve stops at its next step") + return {"ok": False, "type": "cancelled", "error": "cancelled"} + except Refused as e: + return {"ok": False, "type": "invalid_request_error", "error": str(e)} + except Unavailable as e: + return {"ok": False, "type": "service_unavailable", "error": str(e)} + except (OSError, RuntimeError, ValueError, KeyError, http.client.HTTPException) as e: + if not worker.server_alive(): + return {"gone": True, "error": f"mlx-serve is gone ({e})"} + log(f"generate failed: {e}") + return {"ok": False, "type": "server_error", "error": str(e)} + + +def main(): + signal.signal(signal.SIGUSR1, on_cancel) + try: + worker = Worker(json.loads(os.environ.get("LOCALROUTER_TOOL", "{}"))) + except (ValueError, TypeError) as e: + worker, problem = None, str(e) + if worker is not None: + watch_parent(worker) + for line in sys.stdin: + cmd = json.loads(line) + op = cmd.get("op") + if worker is None: + if op == "exit": + return + send({"ok": False, "error": f"bad tool config: {problem}", "type": "server_error"}) + continue + if op == "load": + send(worker.load()) + elif op == "generate": + reply = generate(worker, cmd) + if reply.pop("gone", False): + log(f"{reply['error']}: exiting so the daemon restarts the model") + worker._stop() + os._exit(1) + send(reply) + elif op == "exit": + stopping = threading.Thread(target=worker._stop) + stopping.start() + while stopping.is_alive(): + stopping.join(timeout=5) + if stopping.is_alive(): + send({"progress": {"phase": "unload"}}) + return + if worker is not None: + worker._stop() + + +if __name__ == "__main__": + main() diff --git a/tools/image/qwen-mlx/requirements-macos.txt b/tools/image/qwen-mlx/requirements-macos.txt new file mode 100644 index 0000000..8973e30 --- /dev/null +++ b/tools/image/qwen-mlx/requirements-macos.txt @@ -0,0 +1,84 @@ +pillow==13.0.0 \ + --hash=sha256:02464da86f561f410e13b061ec37c3068a9cb557a095494157742d07b3c24a27 \ + --hash=sha256:04a63720f79e22fd54d897c233a68314c50ca0603360f54ce2bd5db8f0f43782 \ + --hash=sha256:05cd2cd39055d0b9650060e4513ef4a9a0aea870083950bd046a0d6490f93a22 \ + --hash=sha256:06903aac074d00c988e77a1fc435286dc498506d4ca8d6f69df9babddf4ecc15 \ + --hash=sha256:08b826a9bf1265c76de04a42061d61ea6b109418b7826f28ff1ecdfaa26c5502 \ + --hash=sha256:0d1a9cf7bc6917538b20f8d06b9ab2147e542f68849e96fc58eb66bdf6acc692 \ + --hash=sha256:10b5019883c98f1dd819ca48ffc6f33a91cee04948fe637c605f4345c7cafd6f \ + --hash=sha256:12b0fa15879fe5a9d9eebe66e4e2d74128a773b1a4448010c9c714f45e04b9d8 \ + --hash=sha256:12fbb851cf323cc19aac3eea1df5ad17b57e91e9fe3cfd06100dba7bb8208c13 \ + --hash=sha256:149a1c323ba366f90b48e6f1a6db6582eb50512df374588de56fc909abc0f0f3 \ + --hash=sha256:19b66ca0a4d25cc0200b381ce824c7ea090db61427315e73cfb05dbf1b35c369 \ + --hash=sha256:1c774b9c0ba0d0b356f71530677b39028878a2cae5a93e0a13779cfd17120f88 \ + --hash=sha256:20722cd5375533ad754053363a80b8fb30c79d37930a6f7228030e237ce3bd5d \ + --hash=sha256:20822140c9a2ba9c5e42df342b0a892ab2e102fc0e3db967a133bb9ce3f898bb \ + --hash=sha256:21c32061a5512dc13f87cc7303e120e4d77c2f70df89a35af0e3ce246e5ac57f \ + --hash=sha256:2310e2adfc13c6d457d7d0f2ffdd5c56a74a317650cf765887df8d515a647ac0 \ + --hash=sha256:23edf4ec4b6f4f240e8df6e4985e9ee2ced86474632433e55fb1e2548744bf40 \ + --hash=sha256:2d50500a4a3bc6469adcdb319cf05d688b08ca3bf0f22a05339f3b25b433a599 \ + --hash=sha256:2e31c4e68cffd7b1f116127626813b82fcba19edf6b940a7b0c266d6484d0055 \ + --hash=sha256:34fb3db1149c879343861831728c243c432d4a5afc8f1878ca7612ef8de2a033 \ + --hash=sha256:37290fbef0c73e2d51d594eb78fded7dc7170d1c3ad0a890269e6fd8c8c4e061 \ + --hash=sha256:405be6058e206a86fedce2f715ecc2343f94072d7925251632dffa2c05f70dca \ + --hash=sha256:40a266371ece581e5de13b6635ef78d7487824af80561b6ca27694a88c19bc0f \ + --hash=sha256:42d63be3af4f9b951b199b0a9d23365d37ab14687e6d77b74940432c8515c4ca \ + --hash=sha256:44629d4498c9cb6b2051806155f22141df0d8a929c56134b166655232dfb2f8c \ + --hash=sha256:49c12af515073e9bf511e4ef736743bf70d531a643668722279314b937742ba4 \ + --hash=sha256:4a25756244bdcce3b50580baadf320fa48297f0eeb8e114ba64139a40d4d0fc8 \ + --hash=sha256:4bd024d567eca931b668978d672dba914d4b8e2628e2c3073a2ad178de8a2620 \ + --hash=sha256:4f8db54177035a2b5f717f94475b74907e12a96dc74f6d487f3966b0a4694071 \ + --hash=sha256:51ebb529fc5e8abdd55302f6c3f8943037348ec2384329a3e55677f69ae2fdfc \ + --hash=sha256:567f20f1c75897d7663830da6ea8af719e03d712b1d9ee7f813319f9ece42e5a \ + --hash=sha256:5933fcea5a5ba64c686f55c886a115cbfd92f2900460784a1e79909fdbafd7b6 \ + --hash=sha256:5bdb38ba220b61b4ec31dd1046dea6b016d3c6850a23f43ed8b58332b2fc668e \ + --hash=sha256:5f56499c1a74d2b862beb0c33b8016f99d095858455df9d21583bbbc1c15b582 \ + --hash=sha256:67871a13cf47275201553b536c025da777f997fa562a19eb798af4ad4b97c18f \ + --hash=sha256:67c8d44ad89fa7eb88de5d6ef6199b1e03dd24974f85a67639fc8badb801fe79 \ + --hash=sha256:6a0920cb55df5a4d3f6862964908078c0bd59332bdd1bcd20d14e8c70ef27cc5 \ + --hash=sha256:6bce6a29c3898191f4c9998e4d4e8f69790caeabde46e54da4bc235ce5bde683 \ + --hash=sha256:6f38f3f885a09cafc93fd667b1173a92f379e43785d635e7983c893f0da319fb \ + --hash=sha256:73ee25cf8cf8b75d5a5aeed7cb6588f5e5dbfbb0fde51707b630bbb23897645c \ + --hash=sha256:7b999816b5850a35ada08443e6e03945f1f6488144a83afe2d37b240a552aafa \ + --hash=sha256:7c88c10dd432f3b39ccdafc9ed81d0c21b8fd634502574da76d1bfa4fe2c76d0 \ + --hash=sha256:7d5e097eabd5d4915c27ac9f6e9325bb1e01a7c94d7348b5c4451d3fc30e8521 \ + --hash=sha256:7e91e97a929232cdb828889ee6ad0e48d7f40f4ec38818f4e4a8968b2b1af22e \ + --hash=sha256:7f6299e494e46253d40c34653cac369412f44ca53b5f85d848d3f95b1ed0546c \ + --hash=sha256:816752a5ef76a80484c8c441c85593b7ea9304f6231fc5776ad1fbff7e773058 \ + --hash=sha256:83b37abd87b590acd8f47ca8a61f70893322b11834d36e07c2298dacad4d14d2 \ + --hash=sha256:8619685cad8efc4c38c66d3b36aad854523297c43d748b3b6f0cc03c90ba3463 \ + --hash=sha256:8bf295fe459343bc6093c8e5e0df60bfa189e0fcb775a22e2c2488fdeccc6c0d \ + --hash=sha256:8e459b613cf31f8980f584387b4fe47bf08e076e288d61aa8eba1f7fd4ddbb81 \ + --hash=sha256:8f36aa77f74b14a2103622b196402c71330e4e7cf4122ee45c854e9a01e0eecf \ + --hash=sha256:93ff9c9c3663ce23aaa5d0b6b6abac55359e2a0c06c7e30f3771829af816be63 \ + --hash=sha256:9a635068d9793c7f0d929f39f18b58f1bd86c236e8e7b64167b86af9232ebca0 \ + --hash=sha256:9dda0788a9932776a75d8bcc020076edc308ce2508b51f6739a45e3158b3fc42 \ + --hash=sha256:a192360d512305fe92d3053e4f7627bbad1f7708f1f4ad38cccd9336acabe9ad \ + --hash=sha256:a2a14cd73d5e6c927f644b0110ee81f59c4ef506259ec9183904747f2caacc81 \ + --hash=sha256:a625df9a2a5597fa5e5266e01aefc82604bb85d291e3547981e863560d75da6c \ + --hash=sha256:a6aaa7871aebca081eb5c45f83e845bbd9712ae99fa9286dbfdcee835e63f4b4 \ + --hash=sha256:aa694c49787040943570d61c9b8fc01271bd2a2cfdd55f0d896f8aecdce1a981 \ + --hash=sha256:adb3c2124c68e480f4632d06cb7caa3b48a51678475f483b4edb1c6dee29ea43 \ + --hash=sha256:b84be41387fc00db28c13c02e30d42b02d8d8f6a4a741a1aee6012ca7f1f6039 \ + --hash=sha256:bf97ba1b5542050e9b0c74e632323b6d8136cc99d64e193abebee22fe6013c76 \ + --hash=sha256:c3bc8ee0271387e0b1e3ad0ff893645990a656e79bf9c42f4f734a3642b11d51 \ + --hash=sha256:c50c80e3591e5cf9a8dfa6007d4f10732f02e32263850999286e56ccd650cfb2 \ + --hash=sha256:c6fd031333dedcad7ce3e03d244093f43504f209e86925ebdb2b868fa4037630 \ + --hash=sha256:c87a04e7f152ee4db5dbe15f0bfa3fe628a1aeedbe66b81c43b87f52c813f9ac \ + --hash=sha256:cdcd1ba258178f7356d74170c9a81ba4eab6a5d55748bc402a5a2a6fe6be3510 \ + --hash=sha256:d4737f50d23ddd11d6685faacbb9252887dde0bc7997127757247983d74399d6 \ + --hash=sha256:d7710eb676e9bb68965cb045be0d333b4c0d7f9dfbaa7fbf121385b1588683d3 \ + --hash=sha256:d9b0a2bee3f3e03e9f29057aedb807dbf1e1d96ee1b9d0f422d9db6c8583271d \ + --hash=sha256:e65dd6e75de67b2362a4fa3e16aa627595f65054dc4ec73d64680c84d592fa61 \ + --hash=sha256:e892354a424b1e82a78ba16785822bdd40de06f19446a3d4a28047df1a4f8932 \ + --hash=sha256:e8ca8f69b7c77fa96dad99f8638c1fc4f1c0bed4acaa1baf7a3c31bdb80a0ab1 \ + --hash=sha256:e91c2b1cbeae768c98752e67e717e283d68dae45717b41f2d91eaee648e37b7d \ + --hash=sha256:e98b4f91f99b24b1da72e10621ba181f7281922384a7403094debb4199de6c00 \ + --hash=sha256:ef5fe6b7db3f501c56000ba8fe7c7e1a35ca1d45824fbeb95a67daa85cff4dbe \ + --hash=sha256:f11a3c1358bb3c68012a7b1f41b223c3db9f7bfde90b481df028cfe1c3363313 \ + --hash=sha256:f391dc92a015c478745192fe4f2e120c2af2aba18d9fe0616341cbf537599806 \ + --hash=sha256:f7cbb7a64e7ec664775195e4f81923bc8a93d9451b963f844e59118f63147a7a \ + --hash=sha256:f9bd7d260d986f3553518c5d723544346fa48031f8aef3265d786aeb59845bc0 \ + --hash=sha256:fc481f742ccda20e9ed760b8f978dff50f2a4e9bc253ac167fb0a05b0783ef52 \ + --hash=sha256:ff3fe9eb621fbb9db4c4be33ab946f6bac5f98ab50ca8d3260503653656c961b \ + --hash=sha256:ffea6b2986a797c84b6d1d27d4f43e0b83c73fc16281a930891cd237cf0d08fe diff --git a/tools/image/qwen-mlx/test_worker.py b/tools/image/qwen-mlx/test_worker.py new file mode 100644 index 0000000..a745a1e --- /dev/null +++ b/tools/image/qwen-mlx/test_worker.py @@ -0,0 +1,476 @@ +"""The macOS Qwen-Image worker against a fake MLX-Serve: lifecycle, memory, refusals, cancellation, failures, outputs. + + ROOT/venv/bin/python -I tools/image/qwen-mlx/test_worker.py (fetch.sh's venv: Pillow; no model, no GPU) + +The fake (this file run as `fake`) answers MLX-Serve's routes the way v26.10.3 does: /v1/models, /props, and +/v1/images/generations streaming progress events then the PNG; it logs each request's fields, without the images' +bytes, to FAKE_LOG. A prompt picks a failure: "slow" (many steps, for a cancel), "fail" (a 500), "refuse" (a 400), +"oom" (MLX-Serve's working-set 400), "die" (the server exits). +""" +import base64 +import io +import json +import os +import signal +import struct +import subprocess +import sys +import tempfile +import threading +import time +import unittest + +HERE = os.path.dirname(os.path.abspath(__file__)) +sys.path.insert(0, HERE) +from PIL import Image # noqa: E402 + +import qwen_mlx_worker as qw # noqa: E402 + +SIGMAS = [1.0, 0.978453, 0.95418, 0.926626, 0.89508, 0.845148, 0.704534, 0.414568] +ACTIVE = 17 << 30 # the fake's MLX active bytes + + +# -- the fake server ------------------------------------------------------------------------------------------------- + +def fake_main(argv): + if "--version" in argv: + print(f"mlx-serve {os.environ.get('FAKE_VERSION', '26.10.3')}\nmlx 0.32.3") + return + from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + port = int(argv[argv.index("--port") + 1]) + flog = os.environ["FAKE_LOG"] + + def record(obj): + with open(flog, "a") as f: + f.write(json.dumps(obj) + "\n") + + record({"started": os.getpid(), "argv": argv}) + + class H(BaseHTTPRequestHandler): + def log_message(self, *a): + pass + + def reply(self, status, obj): + body = json.dumps(obj).encode() + self.send_response(status) + self.send_header("Content-Type", "application/json") + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path == "/v1/models": + self.reply(200, {"data": [{"id": "x", "loaded": True, "state": "ready"}]}) + elif self.path == "/props": + self.reply(200, {"memory": {"active_bytes": ACTIVE, "peak_bytes": ACTIVE}}) + else: + self.reply(404, {"error": {"message": "no"}}) + + def do_POST(self): + b = json.loads(self.rfile.read(int(self.headers["Content-Length"]))) + keys = list(b) + refs = [b["image"]] + b.get("ref_images", []) if "image" in b else [] + formats = [Image.open(io.BytesIO(base64.b64decode(r))).format for r in refs] + sizes = [Image.open(io.BytesIO(base64.b64decode(r))).size for r in refs] + record({k: v for k, v in b.items() if k not in ("image", "ref_images")} | {"keys": keys, "ref_formats": formats, "ref_sizes": sizes}) + p = b["prompt"] + if p == "die": + os._exit(3) + if p == "fail": + return self.reply(500, {"error": {"message": "generation failed"}}) + if p == "refuse": + return self.reply(400, {"error": {"message": "'mode' must be \"edit\" or \"variation\""}}) + if p == "oom": + return self.reply(400, {"error": {"message": "this edit's working set (references + target) needs more GPU memory than is free"}}) + if b.get("stream") is not True: # MLX-Serve answers one JSON body: no progress, no prompt cancel + record({"unstreamed": p}) + return self.reply(500, {"error": {"message": "the fake streams only"}}) + self.send_response(200) + self.send_header("Content-Type", "text/event-stream") + self.end_headers() + total = 400 if p == "slow" else 8 + try: + for i in range(1, total + 1): + self.wfile.write(f'data: {{"type":"progress","stage":"denoise","step":{i},"total":{total}}}\n\n'.encode()) + self.wfile.flush() + if p == "slow": + time.sleep(0.05) + except (BrokenPipeError, ConnectionResetError): + record({"cancelled": p}) + return + w, h = (int(x) for x in b["size"].split("x")) + mode = "RGBA" if b.get("transparent") else "RGB" + colour = (b["seed"] % 256, 100, 200, 255)[:len(mode)] + buf = io.BytesIO() + Image.new(mode, (w, h), colour).save(buf, "PNG") + ev = {"type": "complete", "data": [{"b64_json": base64.b64encode(buf.getvalue()).decode()}]} + self.wfile.write(f"data: {json.dumps(ev)}\n\n".encode()) + + class Server(ThreadingHTTPServer): + def server_bind(self): # HTTPServer's own looks the host's name up (getfqdn), which can hang without DNS + import socketserver + socketserver.TCPServer.server_bind(self) + self.server_name, self.server_port = "127.0.0.1", port + + Server(("127.0.0.1", port), H).serve_forever() + + +# -- helpers --------------------------------------------------------------------------------------------------------- + +def png(path, size, colour=(10, 20, 30, 255), fmt="PNG", mode="RGBA"): + Image.new(mode, size, colour[:len(mode)]).save(path, fmt) + + +class Env: + """A temp dir with fake weights, a server wrapper, a job dir and the fake's log.""" + + def __init__(self, options=None, version="26.10.3", sigmas=SIGMAS): + self.tmp = tempfile.TemporaryDirectory() + t = self.tmp.name + self.weights = os.path.join(t, "weights") + os.makedirs(self.weights) + with open(os.path.join(self.weights, "model_index.json"), "w") as f: + json.dump({"sample_sigmas": sigmas}, f) + self.server = os.path.join(t, "mlx-serve") + with open(self.server, "w") as f: + f.write(f"#!/bin/sh\nexec '{sys.executable}' -I '{os.path.abspath(__file__)}' fake \"$@\"\n") + os.chmod(self.server, 0o755) + self.job = os.path.join(t, "job") + os.makedirs(self.job) + self.log = os.path.join(t, "fake.log") + open(self.log, "w").close() # noqa: SIM115 + self.env = dict(os.environ, FAKE_LOG=self.log, FAKE_VERSION=version) + self.tool = {"id": "qwen-mac", "kind": "image", "weights": self.weights, + "options": {"server": self.server, "verify": "none", "warmup": False, **(options or {})}} + + def records(self): + with open(self.log) as f: + return [json.loads(x) for x in f] + + def requests(self): + return [r for r in self.records() if "prompt" in r] + + def close(self): + self.tmp.cleanup() + + +class Proc: + """The worker as the daemon runs it: LOCALROUTER_TOOL, one JSON line a message each way.""" + + def __init__(self, env): + self.p = subprocess.Popen([sys.executable, "-I", os.path.join(HERE, "qwen_mlx_worker.py")], + stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.DEVNULL, text=True, + env=dict(env.env, LOCALROUTER_TOOL=json.dumps(env.tool)), start_new_session=True) + + def send(self, obj): + self.p.stdin.write(json.dumps(obj) + "\n") + self.p.stdin.flush() + + def line(self, timeout=30): + out = [] + t = threading.Thread(target=lambda: out.append(self.p.stdout.readline())) + t.start() + t.join(timeout) + if not out: + raise TimeoutError("no line from the worker") + return json.loads(out[0]) if out[0] else None + + def reply(self, timeout=30): + """The next line that is not progress, and the progress seen before it.""" + progress = [] + while True: + m = self.line(timeout) + if m is None or "progress" not in m: + return m, progress + progress.append(m["progress"]) + + def gen(self, env, image, id="img_1"): + self.send({"op": "generate", "id": id, "dir": env.job, "request": {"image": image}}) + return self.reply() + + def close(self): + if self.p.poll() is None: + os.killpg(self.p.pid, signal.SIGKILL) + self.p.wait() + self.p.stdin.close() + self.p.stdout.close() + + +def alive(pid): + try: + os.kill(pid, 0) + return True + except ProcessLookupError: + return False + + +def text_chunks(data): + out, i = {}, 8 + while i < len(data): + n = struct.unpack(">I", data[i:i + 4])[0] + kind = data[i + 4:i + 8] + if kind == b"tEXt": + k, v = data[i + 8:i + 8 + n].split(b"\0", 1) + out[k.decode()] = v.decode() + i += 12 + n + return out + + +# -- tests ----------------------------------------------------------------------------------------------------------- + +class Sizes(unittest.TestCase): + def test_an_edit_without_a_size_keeps_the_last_images_shape_within_the_limits(self): + # the Spark tool's own cases (qwen_image_tool.zig "an edit without a size keeps the image's shape") + self.assertEqual(qw.edit_size(512, 512, 1360, 1048576), (1024, 1024)) + self.assertEqual(qw.edit_size(1920, 1080, 1360, 1048576), (1376 - 32, 768)) + for wh in [(928, 1184), (3000, 1000), (100, 4000), (1, 1)]: + w, h = qw.edit_size(*wh, 1360, 1048576) + self.assertTrue(max(w, h) <= 1360 and w * h <= 1048576 and w % 32 == 0 and h % 32 == 0, (wh, w, h)) + + def test_dimensions_round_half_to_even_as_the_engine(self): + self.assertEqual(qw.dimensions(1024, 768, 1024 * 1024), (1184, 896)) + self.assertRaises(ValueError, qw.dimensions, 1, 5000, 1024 * 1024) + + +class Lifecycle(unittest.TestCase): + def setUp(self): + self.env = Env({"warmup": True}) + self.w = Proc(self.env) + + def tearDown(self): + self.w.close() + self.env.close() + + def test_load_generate_exit(self): + self.w.send({"op": "load"}) + r, _ = self.w.reply() + self.assertEqual((r["ok"], r["cancel"]), (True, "signal")) + self.assertGreaterEqual(r["resident"], ACTIVE) # memory: MLX's active bytes or the server's RSS, the larger + started = [x for x in self.env.records() if "started" in x][0] + argv = started["argv"] + self.assertEqual(argv[argv.index("--host") + 1], "127.0.0.1") + self.assertIn("--skip-mem-preflight", argv) + self.assertIn("--no-w8a8", argv) + + warm = self.env.requests() + self.assertEqual([(q["prompt"], q["size"]) for q in warm], [("warm up", "256x256")]) + r, progress = self.w.gen(self.env, {"prompt": "a fox", "width": 1024, "height": 768, "n": 2, "seed": 7}) + self.assertEqual((r["ok"], r["files"], r["seed"]), (True, ["0.png", "1.png"], 7)) + self.assertEqual(progress[:2], [{"phase": "denoise", "step": 1, "of": 8}, {"phase": "denoise", "step": 2, "of": 8}]) + reqs = self.env.requests()[1:] + self.assertEqual([q["seed"] for q in reqs], [7, 8]) + self.assertEqual({(q["size"], q["steps"], q["transparent"], q["stream"]) for q in reqs}, {("1024x768", 8, True, True)}) + self.assertEqual(reqs[0]["keys"][-1], "prompt") # MLX-Serve takes each field's first match: the prompt goes last + self.assertNotIn("mode", reqs[0]) + for i, name in enumerate(r["files"]): + with open(os.path.join(self.env.job, name), "rb") as f: + data = f.read() + self.assertEqual(text_chunks(data), {"prompt": "a fox", "seed": str(7 + i)}) + im = Image.open(io.BytesIO(data)) + self.assertEqual((im.size, im.mode, im.getpixel((0, 0))[0]), ((1024, 768), "RGBA", 7 + i)) + + self.w.send({"op": "exit"}) + self.assertEqual(self.w.p.wait(timeout=30), 0) + self.assertFalse(alive(started["started"]), "mlx-serve outlived the worker") + + def test_steps_and_seed_edges(self): + self.w.send({"op": "load"}) + self.w.reply() + r, _ = self.w.gen(self.env, {"prompt": "x", "width": 512, "height": 512, "steps": 8, "seed": (1 << 64) - 1, "n": 2}) + self.assertTrue(r["ok"], r) + self.assertEqual([q["seed"] for q in self.env.requests()[1:]], [(1 << 64) - 1, 0]) # wraps as the Spark engine's u64 + + +class Refusals(unittest.TestCase): + """Refused before anything reaches the server (invalid_request_error), as the Spark tool's check.""" + + def setUp(self): + self.env = Env({"max_side": 1360, "max_pixels": 1048576}) + self.w = qw.Worker(self.env.tool) + self.w.port = 1 # nothing listens: a request that got through would fail differently + png(os.path.join(self.env.job, "a.png"), (512, 512)) + png(os.path.join(self.env.job, "m.png"), (256, 256)) + png(os.path.join(self.env.job, "thin.png"), (2500, 500)) + png(os.path.join(self.env.job, "tiny.png"), (8, 8)) + with open(os.path.join(self.env.job, "bad.png"), "wb") as f: + f.write(b"not an image") + + def tearDown(self): + self.env.close() + + def refused(self, image, why): + with self.assertRaises(qw.Refused) as c: + self.w.generate({"dir": self.env.job, "request": {"image": {"prompt": "p", "width": 1024, "height": 1024, **image}}}, None) + self.assertIn(why, str(c.exception)) + + def test_requests_the_model_cannot_serve(self): + self.refused({"references": ["a.png"] * 6}, "up to 5 input images") + self.refused({"steps": 25}, "8-step schedule: omit steps or send 8") + self.refused({"width": 1376, "height": 768}, "up to 1360 pixels a side") + self.refused({"width": 1360, "height": 1024}, "at most 1048576 pixels") + self.refused({"mask": "m.png"}, "a mask needs an input image") + self.refused({"references": ["a.png"], "width": 1008, "height": 1008}, "multiples of 32") + self.refused({"references": ["a.png"], "mask": "m.png"}, "the mask is 256x256, the first image 512x512") + self.refused({"references": ["thin.png"]}, "too long and thin") + self.refused({"references": ["tiny.png"]}, "sides of at least 32 pixels") + self.refused({"references": ["bad.png"]}, "could not be read") + self.refused({"references": ["../a.png"]}, "files in the job directory") + self.refused({"width": 0, "height": 0}, "needs a size") + self.refused({"prompt": ""}, "prompt is required") + self.w.opts.edit_refs = 2 + self.refused({"references": ["a.png"] * 3}, "up to 2 input images") + + def test_not_an_image_request(self): + with self.assertRaises(qw.Refused): + self.w.generate({"dir": self.env.job, "request": {"speech": {"input": "hi"}}}, None) + + +class Edits(unittest.TestCase): + def setUp(self): + self.env = Env({"max_side": 1360, "max_pixels": 1048576}) + self.w = Proc(self.env) + self.w.send({"op": "load"}) + self.w.reply() + + def tearDown(self): + self.w.close() + self.env.close() + + def test_inputs_in_order_and_the_size_from_the_last(self): + png(os.path.join(self.env.job, "ref_0.png"), (512, 512)) + png(os.path.join(self.env.job, "ref_1.jpg"), (640, 480), fmt="JPEG", mode="RGB") + png(os.path.join(self.env.job, "ref_2.webp"), (1920, 1080), fmt="WEBP") + r, _ = self.w.gen(self.env, {"prompt": "put them together", "width": 0, "height": 0, "seed": 3, + "references": ["ref_0.png", "ref_1.jpg", "ref_2.webp"]}) + self.assertTrue(r["ok"], r) + q = self.env.requests()[-1] + self.assertEqual(q["mode"], "edit") + self.assertEqual(q["ref_sizes"], [[512, 512], [640, 480], [1920, 1080]]) # first is `image`, the rest `ref_images` + self.assertEqual(q["ref_formats"], ["PNG", "JPEG", "PNG"]) # WebP sent as PNG, the others as they came + self.assertEqual(q["size"], "1344x768") # the last image's shape within max_side 1360 / max_pixels 1 MP + + def test_a_mask_keeps_the_first_image_where_it_is_opaque(self): + png(os.path.join(self.env.job, "ref_0.png"), (512, 512), colour=(10, 20, 30, 255)) + m = Image.new("RGBA", (512, 512), (0, 0, 0, 255)) + m.paste((0, 0, 0, 0), (0, 0, 256, 512)) # the left half clear: the edit goes there + m.save(os.path.join(self.env.job, "mask_0.png")) + r, _ = self.w.gen(self.env, {"prompt": "make it night", "width": 512, "height": 512, "seed": 200, + "references": ["ref_0.png"], "mask": "mask_0.png"}) + self.assertTrue(r["ok"], r) + with open(os.path.join(self.env.job, "0.png"), "rb") as f: + data = f.read() + im = Image.open(io.BytesIO(data)) + self.assertEqual(im.getpixel((10, 10)), (200, 100, 200, 255)) # the edit (the fake's colour from the seed) + self.assertEqual(im.getpixel((500, 10)), (10, 20, 30, 255)) # the original, exactly + self.assertEqual(text_chunks(data), {"prompt": "make it night", "seed": "200"}) + + +class Cancel(unittest.TestCase): + def test_sigusr1_ends_a_generate_and_the_model_stays_loaded(self): + env = Env() + w = Proc(env) + try: + w.send({"op": "load"}) + w.reply() + pid = [x for x in env.records() if "started" in x][0]["started"] + w.send({"op": "generate", "id": "a", "dir": env.job, "request": {"image": {"prompt": "slow", "width": 512, "height": 512}}}) + self.assertIn("progress", w.line()) + w.p.send_signal(signal.SIGUSR1) + r, _ = w.reply() + self.assertEqual((r["ok"], r["type"]), (False, "cancelled")) + r, _ = w.gen(env, {"prompt": "after", "width": 512, "height": 512}) + self.assertTrue(r["ok"], r) + self.assertTrue(alive(pid)) + self.assertEqual(len([x for x in env.records() if "started" in x]), 1) # no reload + deadline = time.time() + 10 + while not any(x.get("cancelled") == "slow" for x in env.records()) and time.time() < deadline: + time.sleep(0.1) + self.assertTrue(any(x.get("cancelled") == "slow" for x in env.records()), "the server never saw the cancel") + finally: + w.close() + env.close() + + +class Failures(unittest.TestCase): + def setUp(self): + self.env = Env() + self.w = Proc(self.env) + self.w.send({"op": "load"}) + self.w.reply() + + def tearDown(self): + self.w.close() + self.env.close() + + def test_server_answers_map_to_the_daemons_types(self): + for prompt, kind in [("fail", "server_error"), ("refuse", "invalid_request_error"), ("oom", "service_unavailable")]: + r, _ = self.w.gen(self.env, {"prompt": prompt, "width": 512, "height": 512}) + self.assertEqual((r["ok"], r["type"]), (False, kind), prompt) + r, _ = self.w.gen(self.env, {"prompt": "still here", "width": 512, "height": 512}) + self.assertTrue(r["ok"]) + + def test_a_server_that_dies_takes_the_worker_with_it(self): + self.w.send({"op": "generate", "id": "x", "dir": self.env.job, "request": {"image": {"prompt": "die", "width": 512, "height": 512}}}) + self.assertIsNone(self.w.line()) # no reply: the daemon fails the job with a 503 and starts the model again + self.assertEqual(self.w.p.wait(timeout=30), 1) + + +class LoadFailures(unittest.TestCase): + def load(self, env): + w = Proc(env) + try: + w.send({"op": "load"}) + r, _ = w.reply() + return r + finally: + w.close() + env.close() + + def test_an_old_server_runs_the_base_schedule_on_turbo(self): + r = self.load(Env(version="26.10.1")) + self.assertFalse(r["ok"]) + self.assertIn("26.10.2 or newer", r["error"]) + + def test_the_schedule_must_have_the_tools_steps(self): + r = self.load(Env({"steps": 6})) + self.assertIn("sample_sigmas has 8 steps, the tool's steps option 6", r["error"]) + + def test_pinned_files_must_be_there(self): + r = self.load(Env({"verify": "size"})) + self.assertIn("missing or not the pinned size", r["error"]) + + def test_each_pack_is_checked_against_its_own_pins(self): + env = Env({"verify": "size", "pack": "q4"}) + os.makedirs(os.path.join(env.weights, "text_encoder")) + with open(os.path.join(env.weights, "text_encoder", "model.safetensors"), "wb") as f: + f.truncate(17) # not the 4-bit text encoder's pinned size + r = self.load(env) + self.assertFalse(r["ok"]) + self.assertIn("not the pinned size", r["error"]) + self.assertEqual(qw.PACKS["q4"], os.path.join(HERE, "pins-macos-q4.json")) + + def test_the_daemons_own_options_are_not_reported_as_unused(self): + import contextlib + err = io.StringIO() + with contextlib.redirect_stderr(err): + qw.Options({"weights": "/w", "options": {"reference_working_bytes": 2147483648}}) + self.assertEqual(err.getvalue(), "") + qw.Options({"weights": "/w", "options": {"typo_bytes": 1}}) + self.assertIn("options not used by the worker: ['typo_bytes']", err.getvalue()) + + def test_an_unknown_pack_is_a_bad_config(self): + r = self.load(Env({"pack": "q2"})) + self.assertIn("bad tool config: pack must be one of ['q4', 'q8']", r["error"]) + + def test_a_bad_config_answers_every_load(self): + env = Env() + env.tool["weights"] = "" + r = self.load(env) + self.assertIn("bad tool config", r["error"]) + + +if __name__ == "__main__": + if len(sys.argv) > 1 and sys.argv[1] == "fake": + fake_main(sys.argv[2:]) + else: + unittest.main()