zimage-mnn 0.1.1__tar.gz

zimage_mnn-0.1.1/PKG-INFO

@@ -0,0 +1,99 @@
+Metadata-Version: 2.4
+Name: zimage-mnn
+Version: 0.1.1
+Summary: Z-Image-Turbo text-to-image generation on MNN (Apple Silicon Metal / CPU)
+Author: ws
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/alibaba/MNN
+Keywords: mnn,diffusion,text-to-image,z-image,apple-silicon
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: MacOS
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy
+Requires-Dist: pillow
+Requires-Dist: jinja2
+Requires-Dist: mnn>=3.6.0
+Requires-Dist: transformers>=4.40
+Requires-Dist: huggingface_hub
+
+# zimage-mnn
+
+Run [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) text-to-image locally with
+[MNN](https://github.com/alibaba/MNN) — optimized for Apple Silicon (Metal) with a CPU fallback.
+
+The 12.3B-parameter DiT is quantized to int8 weights and runs the full pipeline
+(text encoder → 8-step denoiser → VAE) entirely on-device. No cloud calls.
+
+## Install & run (one line)
+
+```bash
+uvx zimage-mnn "a cute cat sitting on a windowsill"
+```
+
+Models (~11 GB) download automatically on first run and are cached in `~/.cache/zimage_mnn`.
+
+## CLI usage
+
+```bash
+# basic
+uvx zimage-mnn "a red sports car on a coastal highway"
+
+# options
+uvx zimage-mnn "a panda eating bamboo" \
+    --size 1024 \        # resolution (multiple of 16; model tuned for 1024)
+    --steps 8 \          # denoising steps (2-12)
+    --seed 42 \          # random seed
+    --num-images 2 \     # generate several (seed increments)
+    --output out.png \   # output file / prefix
+    --backend metal      # auto | metal | cpu
+```
+
+## HTTP API
+
+```bash
+uvx zimage-mnn --api --port 8000
+```
+
+```bash
+curl -X POST http://127.0.0.1:8000/generate \
+     -H "Content-Type: application/json" \
+     -d '{"prompt":"a cute cat","size":1024,"steps":8,"seed":42}' \
+     -o cat.png
+```
+
+`GET /health` returns `{"status":"ok"}`.
+
+## Python API
+
+```python
+from zimage_mnn import ZImageMNNPipeline
+
+pipe = ZImageMNNPipeline.from_pretrained()          # downloads models
+image = pipe("a cute cat", size=1024, steps=8, seed=42)
+image.save("cat.png")
+
+images = pipe.generate("a dog", num_images=3, seed=7)
+```
+
+## Performance (Apple M5 Pro, 1024×1024, 8 steps)
+
+| Backend | per step | end-to-end |
+|---------|----------|------------|
+| Metal (fp32, int8 weights) | ~13 s | ~2 min |
+| CPU (int8, 12 threads) | ~27 s | ~4 min |
+
+Quality matches the original PyTorch pipeline (single-step cosine ≈ 0.9996).
+
+## Requirements
+
+- Python ≥ 3.9
+- macOS with Apple Silicon recommended (Metal). CPU works anywhere but is slower.
+- ~16 GB RAM recommended; ~11 GB disk for the models.
+
+## License
+
+Apache-2.0. Model weights derive from
+[Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) — see its license.

zimage_mnn-0.1.1/README.md

@@ -0,0 +1,78 @@
+# zimage-mnn
+
+Run [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) text-to-image locally with
+[MNN](https://github.com/alibaba/MNN) — optimized for Apple Silicon (Metal) with a CPU fallback.
+
+The 12.3B-parameter DiT is quantized to int8 weights and runs the full pipeline
+(text encoder → 8-step denoiser → VAE) entirely on-device. No cloud calls.
+
+## Install & run (one line)
+
+```bash
+uvx zimage-mnn "a cute cat sitting on a windowsill"
+```
+
+Models (~11 GB) download automatically on first run and are cached in `~/.cache/zimage_mnn`.
+
+## CLI usage
+
+```bash
+# basic
+uvx zimage-mnn "a red sports car on a coastal highway"
+
+# options
+uvx zimage-mnn "a panda eating bamboo" \
+    --size 1024 \        # resolution (multiple of 16; model tuned for 1024)
+    --steps 8 \          # denoising steps (2-12)
+    --seed 42 \          # random seed
+    --num-images 2 \     # generate several (seed increments)
+    --output out.png \   # output file / prefix
+    --backend metal      # auto | metal | cpu
+```
+
+## HTTP API
+
+```bash
+uvx zimage-mnn --api --port 8000
+```
+
+```bash
+curl -X POST http://127.0.0.1:8000/generate \
+     -H "Content-Type: application/json" \
+     -d '{"prompt":"a cute cat","size":1024,"steps":8,"seed":42}' \
+     -o cat.png
+```
+
+`GET /health` returns `{"status":"ok"}`.
+
+## Python API
+
+```python
+from zimage_mnn import ZImageMNNPipeline
+
+pipe = ZImageMNNPipeline.from_pretrained()          # downloads models
+image = pipe("a cute cat", size=1024, steps=8, seed=42)
+image.save("cat.png")
+
+images = pipe.generate("a dog", num_images=3, seed=7)
+```
+
+## Performance (Apple M5 Pro, 1024×1024, 8 steps)
+
+| Backend | per step | end-to-end |
+|---------|----------|------------|
+| Metal (fp32, int8 weights) | ~13 s | ~2 min |
+| CPU (int8, 12 threads) | ~27 s | ~4 min |
+
+Quality matches the original PyTorch pipeline (single-step cosine ≈ 0.9996).
+
+## Requirements
+
+- Python ≥ 3.9
+- macOS with Apple Silicon recommended (Metal). CPU works anywhere but is slower.
+- ~16 GB RAM recommended; ~11 GB disk for the models.
+
+## License
+
+Apache-2.0. Model weights derive from
+[Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) — see its license.

zimage_mnn-0.1.1/pyproject.toml

@@ -0,0 +1,39 @@
+[build-system]
+requires = ["setuptools>=61.0"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "zimage-mnn"
+version = "0.1.1"
+description = "Z-Image-Turbo text-to-image generation on MNN (Apple Silicon Metal / CPU)"
+readme = "README.md"
+requires-python = ">=3.9"
+license = { text = "Apache-2.0" }
+authors = [{ name = "ws" }]
+keywords = ["mnn", "diffusion", "text-to-image", "z-image", "apple-silicon"]
+classifiers = [
+    "Programming Language :: Python :: 3",
+    "License :: OSI Approved :: Apache Software License",
+    "Operating System :: MacOS",
+    "Topic :: Scientific/Engineering :: Artificial Intelligence",
+]
+dependencies = [
+    "numpy",
+    "pillow",
+    "jinja2",
+    "mnn>=3.6.0",
+    "transformers>=4.40",
+    "huggingface_hub",
+]
+
+[project.scripts]
+zimage-mnn = "zimage_mnn.cli:main"
+
+[project.urls]
+Homepage = "https://github.com/alibaba/MNN"
+
+[tool.setuptools]
+packages = ["zimage_mnn"]
+
+[tool.setuptools.package-data]
+zimage_mnn = ["bin/*"]

zimage_mnn-0.1.1/setup.cfg

@@ -0,0 +1,4 @@
+[egg_info]
+tag_build = 
+tag_date = 0
+

zimage_mnn-0.1.1/zimage_mnn/__init__.py

@@ -0,0 +1,5 @@
+"""zimage-mnn: Z-Image-Turbo text-to-image on MNN."""
+from .pipeline import DEFAULT_REPO, ZImageMNNPipeline
+
+__all__ = ["ZImageMNNPipeline", "DEFAULT_REPO"]
+__version__ = "0.1.0"

zimage_mnn-0.1.1/zimage_mnn/api.py

@@ -0,0 +1,90 @@
+"""Minimal HTTP API server for zimage-mnn (stdlib only, no extra deps).
+
+Endpoints:
+    GET  /health                 -> {"status": "ok"}
+    POST /generate               -> returns PNG bytes
+         body (JSON): {"prompt": str, "size": int, "steps": int, "seed": int, "num_images": int}
+"""
+from __future__ import annotations
+
+import io
+import json
+from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer
+
+
+def serve(args):
+    from .pipeline import DEFAULT_REPO, ZImageMNNPipeline
+
+    repo = args.repo or DEFAULT_REPO
+    print(f"Loading pipeline (models from {repo})...", flush=True)
+    pipe = ZImageMNNPipeline.from_pretrained(repo, model_dir=args.model_dir, backend=args.backend)
+    print("Pipeline ready.", flush=True)
+
+    class Handler(BaseHTTPRequestHandler):
+        def log_message(self, fmt, *a):  # quieter logs
+            print("[api]", fmt % a, flush=True)
+
+        def _send_json(self, obj, code=200):
+            body = json.dumps(obj).encode()
+            self.send_response(code)
+            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 in ("/health", "/"):
+                self._send_json({"status": "ok"})
+            else:
+                self._send_json({"error": "not found"}, 404)
+
+        def do_POST(self):
+            if self.path != "/generate":
+                self._send_json({"error": "not found"}, 404)
+                return
+            try:
+                length = int(self.headers.get("Content-Length", 0))
+                req = json.loads(self.rfile.read(length) or b"{}")
+                prompt = req.get("prompt")
+                if not prompt:
+                    self._send_json({"error": "prompt is required"}, 400)
+                    return
+                size = int(req.get("size", args.size))
+                steps = int(req.get("steps", args.steps))
+                seed = int(req.get("seed", args.seed))
+                num_images = int(req.get("num_images", 1))
+
+                images = pipe.generate(prompt, num_images=num_images, size=size,
+                                       steps=steps, seed=seed, verbose=False)
+                if num_images == 1:
+                    buf = io.BytesIO()
+                    images[0].save(buf, format="PNG")
+                    data = buf.getvalue()
+                    self.send_response(200)
+                    self.send_header("Content-Type", "image/png")
+                    self.send_header("Content-Length", str(len(data)))
+                    self.end_headers()
+                    self.wfile.write(data)
+                else:
+                    # multiple images -> return JSON with base64
+                    import base64
+
+                    out = []
+                    for img in images:
+                        buf = io.BytesIO()
+                        img.save(buf, format="PNG")
+                        out.append(base64.b64encode(buf.getvalue()).decode())
+                    self._send_json({"images": out})
+            except Exception as e:  # noqa: BLE001
+                self._send_json({"error": str(e)}, 500)
+
+    server = ThreadingHTTPServer((args.host, args.port), Handler)
+    print(f"Serving on http://{args.host}:{args.port}", flush=True)
+    print('Try: curl -X POST http://%s:%d/generate -H "Content-Type: application/json" '
+          '-d \'{"prompt":"a cute cat"}\' -o out.png' % (args.host, args.port), flush=True)
+    try:
+        server.serve_forever()
+    except KeyboardInterrupt:
+        pass
+    finally:
+        server.server_close()

zimage_mnn-0.1.1/zimage_mnn/cli.py

@@ -0,0 +1,101 @@
+"""Command-line interface for zimage-mnn.
+
+Usage:
+    uvx zimage-mnn "a cute cat" --size 1024 --steps 8 --num-images 2
+    uvx zimage-mnn --api --port 8000
+"""
+from __future__ import annotations
+
+import argparse
+import sys
+import time
+
+
+def build_parser() -> argparse.ArgumentParser:
+    p = argparse.ArgumentParser(
+        prog="zimage-mnn",
+        description="Z-Image-Turbo text-to-image on MNN (Apple Silicon / CPU).")
+    p.add_argument("prompt", nargs="?", help="Text prompt for image generation.")
+    p.add_argument("--repo", default=None, help="HF repo id to download models from.")
+    p.add_argument("--model-dir", default=None, help="Use a local model directory instead of downloading.")
+    p.add_argument("--size", type=int, default=1024, help="Image resolution (width=height). Must be a multiple of 16.")
+    p.add_argument("--width", type=int, default=None, help="Image width (overrides --size).")
+    p.add_argument("--height", type=int, default=None, help="Image height (overrides --size).")
+    p.add_argument("--steps", type=int, default=8, help="Number of denoising steps (2-12).")
+    p.add_argument("--seed", type=int, default=42, help="Random seed.")
+    p.add_argument("--num-images", "-n", type=int, default=1, help="Number of images to generate.")
+    p.add_argument("--output", "-o", default="zimage.png", help="Output file (or prefix for multiple images).")
+    p.add_argument("--backend", default="auto", choices=["auto", "metal", "cpu"], help="Compute backend.")
+    p.add_argument("--api", action="store_true", help="Run an HTTP API server instead of generating.")
+    p.add_argument("--host", default="127.0.0.1", help="API host.")
+    p.add_argument("--port", type=int, default=8000, help="API port.")
+    p.add_argument("--quiet", "-q", action="store_true", help="Suppress progress output.")
+    return p
+
+
+def _save_images(images, output: str):
+    from pathlib import Path
+
+    saved = []
+    if len(images) == 1:
+        Path(output).parent.mkdir(parents=True, exist_ok=True)
+        images[0].save(output)
+        saved.append(output)
+    else:
+        stem = Path(output)
+        stem.parent.mkdir(parents=True, exist_ok=True)
+        base = stem.with_suffix("")
+        ext = stem.suffix or ".png"
+        for i, img in enumerate(images):
+            p = f"{base}_{i}{ext}"
+            img.save(p)
+            saved.append(p)
+    return saved
+
+
+def run_api(args):
+    from .api import serve
+
+    serve(args)
+
+
+def main(argv=None):
+    args = build_parser().parse_args(argv)
+
+    if args.api:
+        run_api(args)
+        return 0
+
+    if not args.prompt:
+        print("Error: a prompt is required (or use --api).", file=sys.stderr)
+        return 1
+
+    from .pipeline import DEFAULT_REPO, ZImageMNNPipeline
+
+    repo = args.repo or DEFAULT_REPO
+    t0 = time.time()
+    pipe = ZImageMNNPipeline.from_pretrained(
+        repo, model_dir=args.model_dir, backend=args.backend)
+    if not args.quiet:
+        print(f"Pipeline ready in {time.time() - t0:.1f}s", flush=True)
+
+    width = args.width or args.size
+    height = args.height or args.size
+    if width % 16 or height % 16:
+        print("Error: width/height must be multiples of 16.", file=sys.stderr)
+        return 1
+
+    images = []
+    for i in range(args.num_images):
+        img = pipe(args.prompt, size=width, steps=args.steps, seed=args.seed + i,
+                   verbose=not args.quiet)
+        images.append(img)
+
+    saved = _save_images(images, args.output)
+    for s in saved:
+        print(f"Saved: {s}")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

zimage_mnn-0.1.1/zimage_mnn/pipeline.py

@@ -0,0 +1,262 @@
+"""Z-Image-Turbo MNN inference pipeline.
+
+Runs the full text-to-image pipeline on MNN:
+  text encoder (Qwen3) -> DiT denoiser (FlowMatch Euler) -> VAE decoder.
+
+Models are downloaded from HuggingFace on first use and cached locally.
+"""
+from __future__ import annotations
+
+import os
+import struct
+import time
+from pathlib import Path
+
+import numpy as np
+
+DEFAULT_REPO = "yunfengwang/z-image-turbo-mnn"
+CACHE_DIR = Path(os.environ.get("ZIMAGE_MNN_CACHE", Path.home() / ".cache" / "zimage_mnn"))
+
+# VAE post-processing constants (from the original model config)
+VAE_SCALING_FACTOR = 0.3611
+VAE_SHIFT_FACTOR = 0.1159
+SCHEDULER_SHIFT = 3.0
+CAP_LEN = 256  # fixed text-embedding length the unet was exported with
+
+
+def _download_models(repo_id: str, model_dir: Path) -> Path:
+    """Download model files from HuggingFace if not already cached."""
+    needed = [
+        "unet.mnn", "unet.mnn.weight",
+        "text_encoder.mnn", "text_encoder.mnn.weight",
+        "vae_decoder.mnn",
+        "tokenizer/tokenizer.json", "tokenizer/tokenizer_config.json",
+    ]
+    missing = [f for f in needed if not (model_dir / f).exists()]
+    if missing:
+        from huggingface_hub import snapshot_download
+
+        print(f"Downloading models from {repo_id} ...")
+        snapshot_download(repo_id, local_dir=str(model_dir))
+    return model_dir
+
+
+class ZImageMNNPipeline:
+    """MNN pipeline for Z-Image-Turbo text-to-image."""
+
+    def __init__(self, model_dir: str | Path, backend: str = "auto"):
+        import MNN.expr as F  # noqa: F401
+        import MNN.nn as nn  # noqa: F401
+
+        self.model_dir = Path(model_dir)
+        self.backend = backend
+        self._tokenizer = None
+
+    @classmethod
+    def from_pretrained(cls, repo_id: str = DEFAULT_REPO, model_dir: str | Path | None = None,
+                        backend: str = "auto") -> "ZImageMNNPipeline":
+        model_dir = Path(model_dir) if model_dir else CACHE_DIR / repo_id.replace("/", "--")
+        _download_models(repo_id, model_dir)
+        return cls(model_dir, backend=backend)
+
+    # -- helpers ---------------------------------------------------------
+    def _load_tokenizer(self):
+        if self._tokenizer is None:
+            from transformers import AutoTokenizer
+
+            self._tokenizer = AutoTokenizer.from_pretrained(str(self.model_dir / "tokenizer"))
+        return self._tokenizer
+
+    def _metal_available(self) -> bool:
+        if self.backend == "cpu":
+            return False
+        if self.backend == "metal":
+            return True
+        # auto: try Metal, fall back to CPU
+        import platform
+
+        return platform.system() == "Darwin"
+
+    def _encode_prompt(self, prompt: str) -> np.ndarray:
+        import MNN.expr as F
+        import MNN.nn as nn
+
+        tokenizer = self._load_tokenizer()
+        max_length = 512
+        messages = [{"role": "user", "content": prompt}]
+        formatted = tokenizer.apply_chat_template(
+            messages, tokenize=False, add_generation_prompt=True, enable_thinking=True)
+        inputs = tokenizer([formatted], padding="max_length", max_length=max_length,
+                           truncation=True, return_tensors="np")
+        input_ids = inputs.input_ids.astype(np.int64)
+        mask1d = inputs.attention_mask[0].astype(bool)
+        actual_len = int(mask1d.sum())
+
+        # 4D causal+padding mask: 0 = attend, -inf = masked
+        min_dtype = np.finfo(np.float32).min
+        causal = np.tril(np.ones((max_length, max_length), dtype=bool))
+        mask4d = np.where(causal & mask1d[np.newaxis, :], 0.0, min_dtype).astype(np.float32)
+        mask4d = mask4d[np.newaxis, np.newaxis, :, :]
+
+        rtmgr = nn.create_runtime_manager(
+            ({"backend": "CPU", "precision": "normal", "numThread": os.cpu_count() or 4},))
+        te = nn.load_module_from_file(
+            str(self.model_dir / "text_encoder.mnn"),
+            ["input_ids", "attention_mask"], ["last_hidden_state"], runtime_manager=rtmgr)
+        ids_var = F.const(input_ids, list(input_ids.shape), F.NCHW, F.int64)
+        mask_var = F.const(mask4d, list(mask4d.shape), F.NCHW)
+        out = te.forward([ids_var, mask_var])
+        hidden = np.array(out[0].read())
+        return hidden[0, :actual_len, :].astype(np.float32)
+
+    def _pad_cap(self, cap_feats: np.ndarray) -> np.ndarray:
+        n = min(cap_feats.shape[0], CAP_LEN)
+        padded = np.zeros((CAP_LEN, cap_feats.shape[1]), dtype=np.float32)
+        padded[:n] = cap_feats[:n]
+        padded[n:] = cap_feats[n - 1]
+        return padded[np.newaxis, :, :]
+
+    @staticmethod
+    def _make_sigmas(num_steps: int) -> np.ndarray:
+        sigmas = np.linspace(1.0, 1.0 / num_steps, num_steps)
+        sigmas = SCHEDULER_SHIFT * sigmas / (1 + (SCHEDULER_SHIFT - 1) * sigmas)
+        return np.append(sigmas, 0.0)
+
+    def _denoise(self, cap_input: np.ndarray, num_steps: int, size: int, seed: int,
+                 verbose: bool = True) -> np.ndarray:
+        # Prefer the bundled C++ binary (fp16 Metal, ~2x faster than the Python API).
+        binary = self._find_binary()
+        if binary is not None and size == 1024:
+            result = self._denoise_binary(binary, cap_input, num_steps, seed, verbose)
+            if result is not None:
+                return result
+        return self._denoise_python(cap_input, num_steps, size, seed, verbose)
+
+    def _find_binary(self):
+        """Locate the bundled zimage_mnn binary (macOS arm64 only)."""
+        import platform
+
+        if platform.system() != "Darwin" or platform.machine() != "arm64":
+            return None
+        candidate = Path(__file__).parent / "bin" / "zimage_mnn"
+        if candidate.exists() and os.access(candidate, os.X_OK):
+            return candidate
+        return None
+
+    def _denoise_binary(self, binary: Path, cap_input: np.ndarray, num_steps: int, seed: int,
+                        verbose: bool):
+        import subprocess
+        import tempfile
+
+        latent = 1024 // 8
+        # The binary reads embeddings (ZIMAGE_EMB) and initial noise (ZIMAGE_NOISE) from
+        # temp files, and writes latents to <output>.latents.bin.
+        with tempfile.TemporaryDirectory() as td:
+            emb_path = os.path.join(td, "emb.bin")
+            noise_path = os.path.join(td, "noise.bin")
+            out_prefix = os.path.join(td, "out")
+
+            cap_flat = cap_input[0].astype(np.float32)  # (256, 2560)
+            with open(emb_path, "wb") as f:
+                f.write(struct.pack("i", cap_flat.shape[0]))
+                f.write(cap_flat.tobytes())
+
+            # Same RNG as the Python path so seeds match.
+            rng = np.random.default_rng(seed)
+            rng.standard_normal((16 * latent * latent), dtype=np.float32).tofile(noise_path)
+
+            env = dict(os.environ, ZIMAGE_EMB=emb_path, ZIMAGE_NOISE=noise_path)
+            # backend: 4 = Metal fp16 (fast, model is fp16-safe), 0 = CPU
+            backend = "4" if self._metal_available() else "0"
+            cmd = [str(binary), str(self.model_dir), backend, str(num_steps), str(seed),
+                   out_prefix, "generate"]
+            t0 = time.time()
+            r = subprocess.run(cmd, capture_output=True, text=True, env=env)
+            if r.returncode != 0:
+                if verbose:
+                    print(f"  binary denoise failed, falling back to Python: {r.stderr[-200:]}")
+                return None
+            with open(out_prefix + ".latents.bin", "rb") as f:
+                f.read(4)  # skip the int count header written by the C++ binary
+                latents = np.frombuffer(f.read(), dtype=np.float32)
+        if verbose:
+            print(f"  denoise (C++ binary): {time.time() - t0:.1f}s", flush=True)
+        return latents.reshape(1, 16, latent, latent)
+
+    def _denoise_python(self, cap_input: np.ndarray, num_steps: int, size: int, seed: int,
+                        verbose: bool = True) -> np.ndarray:
+        import MNN.expr as F
+        import MNN.nn as nn
+
+        latent = size // 8
+        rng = np.random.default_rng(seed)
+        latents = rng.standard_normal((1, 16, latent, latent), dtype=np.float32)
+        sigmas = self._make_sigmas(num_steps)
+
+        use_metal = self._metal_available()
+        cfg = ({"backend": "METAL", "precision": "high"} if use_metal
+               else {"backend": "CPU", "precision": "normal", "numThread": os.cpu_count() or 4})
+        rtmgr = nn.create_runtime_manager((cfg,))
+        unet = nn.load_module_from_file(
+            str(self.model_dir / "unet.mnn"),
+            ["sample", "timestep", "encoder_hidden_states"], ["out_sample"], runtime_manager=rtmgr)
+
+        sv = F.const(latents, list(latents.shape), F.NCHW)
+        tv = F.const(np.array([0.5], dtype=np.float32), [1], F.NCHW)
+        cv = F.const(cap_input, list(cap_input.shape), F.NCHW)
+        unet.forward([sv, tv, cv])[0].read()  # warmup / shader compile
+
+        for i in range(num_steps):
+            t0 = time.time()
+            t_input = (1000.0 - sigmas[i] * 1000.0) / 1000.0
+            sv = F.const(latents, list(latents.shape), F.NCHW)
+            tv = F.const(np.array([t_input], dtype=np.float32), [1], F.NCHW)
+            out = unet.forward([sv, tv, cv])
+            noise_pred = -np.array(out[0].read())
+            latents = latents + (sigmas[i + 1] - sigmas[i]) * noise_pred
+            if verbose:
+                print(f"  step {i + 1}/{num_steps}: {time.time() - t0:.2f}s", flush=True)
+        return latents
+
+    def _decode(self, latents: np.ndarray) -> np.ndarray:
+        import MNN.expr as F
+        import MNN.nn as nn
+
+        latents = latents / VAE_SCALING_FACTOR + VAE_SHIFT_FACTOR
+        rtmgr = nn.create_runtime_manager(
+            ({"backend": "CPU", "precision": "normal", "numThread": os.cpu_count() or 4},))
+        vae = nn.load_module_from_file(
+            str(self.model_dir / "vae_decoder.mnn"),
+            ["latent_sample"], ["sample"], runtime_manager=rtmgr)
+        lv = F.const(latents.astype(np.float32), list(latents.shape), F.NCHW)
+        out = vae.forward([lv])
+        image = np.array(out[0].read())
+        image = image[0].transpose(1, 2, 0)
+        return np.clip((image + 1) / 2 * 255, 0, 255).astype(np.uint8)
+
+    # -- public API ------------------------------------------------------
+    def __call__(self, prompt: str, size: int = 1024, steps: int = 8, seed: int = 42,
+                 verbose: bool = True):
+        """Generate one image. Returns a PIL.Image."""
+        from PIL import Image
+
+        t0 = time.time()
+        if verbose:
+            print(f"Encoding prompt: {prompt!r}", flush=True)
+        cap_feats = self._encode_prompt(prompt)
+        cap_input = self._pad_cap(cap_feats)
+        if verbose:
+            print(f"Denoising ({steps} steps, {size}x{size}, seed={seed})...", flush=True)
+        latents = self._denoise(cap_input, steps, size, seed, verbose=verbose)
+        if verbose:
+            print("VAE decoding...", flush=True)
+        pixels = self._decode(latents)
+        if verbose:
+            print(f"Done in {time.time() - t0:.1f}s", flush=True)
+        return Image.fromarray(pixels)
+
+    def generate(self, prompt: str, num_images: int = 1, size: int = 1024, steps: int = 8,
+                 seed: int = 42, verbose: bool = True):
+        """Generate multiple images (varying seed). Returns a list of PIL.Image."""
+        return [self(prompt, size=size, steps=steps, seed=seed + i, verbose=verbose)
+                for i in range(num_images)]

zimage_mnn-0.1.1/zimage_mnn.egg-info/PKG-INFO

@@ -0,0 +1,99 @@
+Metadata-Version: 2.4
+Name: zimage-mnn
+Version: 0.1.1
+Summary: Z-Image-Turbo text-to-image generation on MNN (Apple Silicon Metal / CPU)
+Author: ws
+License: Apache-2.0
+Project-URL: Homepage, https://github.com/alibaba/MNN
+Keywords: mnn,diffusion,text-to-image,z-image,apple-silicon
+Classifier: Programming Language :: Python :: 3
+Classifier: License :: OSI Approved :: Apache Software License
+Classifier: Operating System :: MacOS
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Requires-Python: >=3.9
+Description-Content-Type: text/markdown
+Requires-Dist: numpy
+Requires-Dist: pillow
+Requires-Dist: jinja2
+Requires-Dist: mnn>=3.6.0
+Requires-Dist: transformers>=4.40
+Requires-Dist: huggingface_hub
+
+# zimage-mnn
+
+Run [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) text-to-image locally with
+[MNN](https://github.com/alibaba/MNN) — optimized for Apple Silicon (Metal) with a CPU fallback.
+
+The 12.3B-parameter DiT is quantized to int8 weights and runs the full pipeline
+(text encoder → 8-step denoiser → VAE) entirely on-device. No cloud calls.
+
+## Install & run (one line)
+
+```bash
+uvx zimage-mnn "a cute cat sitting on a windowsill"
+```
+
+Models (~11 GB) download automatically on first run and are cached in `~/.cache/zimage_mnn`.
+
+## CLI usage
+
+```bash
+# basic
+uvx zimage-mnn "a red sports car on a coastal highway"
+
+# options
+uvx zimage-mnn "a panda eating bamboo" \
+    --size 1024 \        # resolution (multiple of 16; model tuned for 1024)
+    --steps 8 \          # denoising steps (2-12)
+    --seed 42 \          # random seed
+    --num-images 2 \     # generate several (seed increments)
+    --output out.png \   # output file / prefix
+    --backend metal      # auto | metal | cpu
+```
+
+## HTTP API
+
+```bash
+uvx zimage-mnn --api --port 8000
+```
+
+```bash
+curl -X POST http://127.0.0.1:8000/generate \
+     -H "Content-Type: application/json" \
+     -d '{"prompt":"a cute cat","size":1024,"steps":8,"seed":42}' \
+     -o cat.png
+```
+
+`GET /health` returns `{"status":"ok"}`.
+
+## Python API
+
+```python
+from zimage_mnn import ZImageMNNPipeline
+
+pipe = ZImageMNNPipeline.from_pretrained()          # downloads models
+image = pipe("a cute cat", size=1024, steps=8, seed=42)
+image.save("cat.png")
+
+images = pipe.generate("a dog", num_images=3, seed=7)
+```
+
+## Performance (Apple M5 Pro, 1024×1024, 8 steps)
+
+| Backend | per step | end-to-end |
+|---------|----------|------------|
+| Metal (fp32, int8 weights) | ~13 s | ~2 min |
+| CPU (int8, 12 threads) | ~27 s | ~4 min |
+
+Quality matches the original PyTorch pipeline (single-step cosine ≈ 0.9996).
+
+## Requirements
+
+- Python ≥ 3.9
+- macOS with Apple Silicon recommended (Metal). CPU works anywhere but is slower.
+- ~16 GB RAM recommended; ~11 GB disk for the models.
+
+## License
+
+Apache-2.0. Model weights derive from
+[Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo) — see its license.

zimage_mnn-0.1.1/zimage_mnn.egg-info/SOURCES.txt

@@ -0,0 +1,15 @@
+README.md
+pyproject.toml
+zimage_mnn/__init__.py
+zimage_mnn/api.py
+zimage_mnn/cli.py
+zimage_mnn/pipeline.py
+zimage_mnn.egg-info/PKG-INFO
+zimage_mnn.egg-info/SOURCES.txt
+zimage_mnn.egg-info/dependency_links.txt
+zimage_mnn.egg-info/entry_points.txt
+zimage_mnn.egg-info/requires.txt
+zimage_mnn.egg-info/top_level.txt
+zimage_mnn/bin/libMNN.dylib
+zimage_mnn/bin/libMNN_Express.dylib
+zimage_mnn/bin/zimage_mnn

zimage_mnn-0.1.1/zimage_mnn.egg-info/dependency_links.txt

@@ -0,0 +1 @@
+

zimage_mnn-0.1.1/zimage_mnn.egg-info/entry_points.txt

@@ -0,0 +1,2 @@
+[console_scripts]
+zimage-mnn = zimage_mnn.cli:main

zimage_mnn-0.1.1/zimage_mnn.egg-info/requires.txt

@@ -0,0 +1,6 @@
+numpy
+pillow
+jinja2
+mnn>=3.6.0
+transformers>=4.40
+huggingface_hub

zimage_mnn-0.1.1/zimage_mnn.egg-info/top_level.txt

@@ -0,0 +1 @@
+zimage_mnn