Upload folder using huggingface_hub

README.md

@@ -0,0 +1,163 @@
+---
+license: apache-2.0
+language:
+  - en
+  - zh
+library_name: mlx
+tags:
+  - mlx
+  - text-to-image
+  - apple-silicon
+  - image-generation
+  - diffusion
+  - dit
+base_model: Tongyi-MAI/Z-Image-Turbo
+pipeline_tag: text-to-image
+---
+
+# Z-Image-Turbo MLX
+
+**Pure MLX (Apple Silicon) inference pipeline for [Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo)** — a 10.26B parameter text-to-image model by Tongyi-MAI.
+
+Zero PyTorch dependency. Runs natively on Apple Silicon via Metal GPU.
+
+## Highlights
+
+- **100% MLX native** — no torch, no diffusers, no transformers needed
+- **bfloat16 inference** with `mx.compile()` kernel fusion
+- **Fused attention** via `mx.fast.scaled_dot_product_attention`
+- **Optional quantization** (4-bit / 8-bit) for low-memory devices
+- **Pixel-identical quality** to the PyTorch reference (verified VAE pixel diff = 0.00 on same latent input)
+- **Chinese & English** prompts supported (Qwen3 text encoder)
+
+## Performance (Apple Silicon)
+
+| Resolution | Steps | MLX (bf16) | PyTorch MPS | Ratio |
+|-----------|-------|-----------|-------------|-------|
+| 512×512 | 4 | 5.6s | 5.4s | 1.04× |
+| 512×512 | 8 | 10.6s | 10.0s | 1.06× |
+| 768×768 | 8 | 26.5s | 24.6s | 1.08× |
+
+| Metric | MLX | PyTorch MPS |
+|--------|-----|-------------|
+| Load time | **6.3s** | 16.8s |
+| Memory (loaded) | 19.1 GB | ~19 GB |
+| Dependencies disk | ~200 MB | ~5 GB |
+
+## Quick Start
+
+### Install
+
+```bash
+pip install mlx safetensors tokenizers pillow numpy
+```
+
+### Download
+
+```bash
+# Clone this repo (includes inference code)
+git clone https://huggingface.co/illusion615/Z-Image-Turbo-MLX
+
+# Or download the original weights (the inference code handles both)
+huggingface-cli download Tongyi-MAI/Z-Image-Turbo
+```
+
+### Generate
+
+```python
+from pipeline import ZImageMLXPipeline
+
+pipe = ZImageMLXPipeline()
+pipe.load()
+
+# Generate an image
+image = pipe.generate(
+    prompt="a red cat sitting on a wooden table, detailed fur, soft lighting",
+    width=512,
+    height=512,
+    num_steps=8,
+    seed=42,
+)
+
+# Save
+from PIL import Image
+Image.fromarray(image).save("output.png")
+
+pipe.unload()
+```
+
+### With Quantization (low memory)
+
+```python
+import mlx.nn as nn
+
+pipe = ZImageMLXPipeline()
+pipe.load()
+
+# Quantize DiT to 4-bit (reduces memory from 19 GB → 11 GB)
+def large_only(path, module):
+    return isinstance(module, nn.Linear) and module.weight.shape[-1] >= 1024
+
+nn.quantize(pipe._dit, bits=4, group_size=64, class_predicate=large_only)
+```
+
+## Architecture
+
+```
+ZImageMLXPipeline
+├── Qwen3 Text Encoder (4.02B params, bfloat16)
+│   └── 36-layer decoder-only transformer, hidden_size=2560
+├── ZImage DiT Transformer (6.15B params, bfloat16)
+│   ├── 2 noise_refiner blocks (with AdaLN)
+│   ├── 2 context_refiner blocks (no AdaLN)
+│   ├── 30 main DiT blocks (with AdaLN + RoPE)
+│   └── Final layer (adaLN + Linear)
+├── VAE Decoder (84M params, float32 force_upcast)
+│   └── 4 UpDecoderBlock2D + MidBlock2D with attention
+└── FlowMatch Euler Scheduler (shift=3.0)
+```
+
+## Files
+
+```
+├── pipeline.py          # Main inference pipeline
+├── zimage_dit.py        # DiT transformer (S3-DiT architecture)
+├── qwen3_encoder.py     # Qwen3 text encoder
+├── autoencoder.py       # VAE decoder
+├── scheduler.py         # FlowMatch Euler sampler
+├── tokenizer.py         # Fast BPE tokenizer
+├── weight_loader.py     # Safetensors loader with key mapping
+└── config.json          # Pipeline configuration
+```
+
+## Model Source
+
+This is a pure-MLX inference adaptation of [Tongyi-MAI/Z-Image-Turbo](https://huggingface.co/Tongyi-MAI/Z-Image-Turbo). The original model weights are loaded from the upstream HuggingFace repository. All inference code is original work.
+
+## Verified Quality
+
+The MLX pipeline has been validated against the PyTorch/diffusers reference at four levels:
+
+| Component | Validation | Result |
+|-----------|-----------|--------|
+| Tokenizer | Token-by-token comparison | Exact match |
+| Text Encoder | Cosine similarity | 0.999989 |
+| Scheduler | Sigma schedule diff | max 1e-7 |
+| VAE Decoder | Pixel difference (same latent) | 0.00 |
+
+## License
+
+Apache 2.0 (same as upstream Z-Image-Turbo)
+
+## Citation
+
+If you use this MLX adaptation, please also cite the original model:
+
+```bibtex
+@misc{z-image-turbo,
+  title={Z-Image-Turbo},
+  author={Tongyi-MAI},
+  year={2025},
+  url={https://huggingface.co/Tongyi-MAI/Z-Image-Turbo}
+}
+```

__init__.py

@@ -0,0 +1 @@
+# Z-Image-Turbo MLX native backend

autoencoder.py

@@ -0,0 +1,217 @@
+"""AutoencoderKL Decoder — pure MLX implementation.
+
+Decodes latent representations to RGB images without PyTorch/diffusers
+dependency.  Architecture matches diffusers AutoencoderKL with the
+Z-Image-Turbo VAE config:
+
+    latent_channels  = 16
+    block_out_channels = [128, 256, 512, 512]
+    layers_per_block = 2       (decoder uses layers_per_block + 1 = 3)
+    norm_num_groups  = 32
+    mid_block_add_attention = true
+    force_upcast     = true    (all ops in float32)
+    scaling_factor   = 0.3611
+    shift_factor     = 0.1159
+
+Data format: NHWC throughout (MLX convention).
+"""
+
+from __future__ import annotations
+
+import math
+
+import mlx.core as mx
+import mlx.nn as nn
+
+# Match diffusers VAE GroupNorm: eps=1e-6, pytorch_compatible=True
+_GN_EPS = 1e-6
+
+
+def _gn(groups: int, channels: int) -> nn.GroupNorm:
+    return nn.GroupNorm(groups, channels, eps=_GN_EPS, pytorch_compatible=True)
+
+
+# ── Building blocks ──────────────────────────────────────────────
+
+
+class ResnetBlock2D(nn.Module):
+    """Residual block: GroupNorm → SiLU → Conv → GroupNorm → SiLU → Conv + skip."""
+
+    def __init__(self, in_channels: int, out_channels: int, groups: int = 32):
+        super().__init__()
+        self.norm1 = _gn(groups, in_channels)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1)
+        self.norm2 = _gn(groups, out_channels)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1)
+
+        self.conv_shortcut = None
+        if in_channels != out_channels:
+            self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        residual = x
+        x = nn.silu(self.norm1(x))
+        x = self.conv1(x)
+        x = nn.silu(self.norm2(x))
+        x = self.conv2(x)
+        if self.conv_shortcut is not None:
+            residual = self.conv_shortcut(residual)
+        return x + residual
+
+
+class AttentionBlock(nn.Module):
+    """Single-head self-attention over spatial positions (NHWC)."""
+
+    def __init__(self, channels: int, groups: int = 32):
+        super().__init__()
+        self.group_norm = _gn(groups, channels)
+        self.to_q = nn.Linear(channels, channels)
+        self.to_k = nn.Linear(channels, channels)
+        self.to_v = nn.Linear(channels, channels)
+        # diffusers wraps out-proj in a list (Sequential): to_out.0
+        self.to_out = [nn.Linear(channels, channels)]
+
+    def __call__(self, x: mx.array) -> mx.array:
+        residual = x
+        B, H, W, C = x.shape
+        x = self.group_norm(x)
+        x = x.reshape(B, H * W, C)
+
+        q = self.to_q(x)
+        k = self.to_k(x)
+        v = self.to_v(x)
+
+        scale = 1.0 / math.sqrt(C)
+        attn = (q @ k.transpose(0, 2, 1)) * scale
+        attn = mx.softmax(attn, axis=-1)
+        x = attn @ v
+
+        x = self.to_out[0](x)
+        x = x.reshape(B, H, W, C)
+        return x + residual
+
+
+class Upsample2D(nn.Module):
+    """2× nearest-neighbour upsample followed by a 3×3 conv."""
+
+    def __init__(self, channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, kernel_size=3, padding=1)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # Nearest-neighbour 2× in NHWC
+        B, H, W, C = x.shape
+        x = mx.repeat(x, 2, axis=1)
+        x = mx.repeat(x, 2, axis=2)
+        x = self.conv(x)
+        return x
+
+
+class UpDecoderBlock2D(nn.Module):
+    """Decoder up-block: N resnet blocks + optional 2× upsample."""
+
+    def __init__(
+        self,
+        in_channels: int,
+        out_channels: int,
+        num_layers: int = 3,
+        add_upsample: bool = True,
+        groups: int = 32,
+    ):
+        super().__init__()
+        self.resnets = []
+        for i in range(num_layers):
+            res_in = in_channels if i == 0 else out_channels
+            self.resnets.append(ResnetBlock2D(res_in, out_channels, groups))
+
+        self.upsamplers = []
+        if add_upsample:
+            self.upsamplers.append(Upsample2D(out_channels))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        for resnet in self.resnets:
+            x = resnet(x)
+        for up in self.upsamplers:
+            x = up(x)
+        return x
+
+
+class MidBlock2D(nn.Module):
+    """Mid block: resnet → self-attention → resnet."""
+
+    def __init__(self, channels: int, groups: int = 32):
+        super().__init__()
+        self.resnets = [
+            ResnetBlock2D(channels, channels, groups),
+            ResnetBlock2D(channels, channels, groups),
+        ]
+        self.attentions = [AttentionBlock(channels, groups)]
+
+    def __call__(self, x: mx.array) -> mx.array:
+        x = self.resnets[0](x)
+        x = self.attentions[0](x)
+        x = self.resnets[1](x)
+        return x
+
+
+# ── Decoder ──────────────────────────────────────────────────────
+
+
+class Decoder(nn.Module):
+    """AutoencoderKL Decoder (NHWC, pure MLX).
+
+    Module hierarchy matches diffusers weight-key paths after stripping
+    the ``decoder.`` prefix, so weights can be loaded directly.
+    """
+
+    def __init__(
+        self,
+        latent_channels: int = 16,
+        block_out_channels: tuple[int, ...] = (128, 256, 512, 512),
+        layers_per_block: int = 2,
+        norm_num_groups: int = 32,
+    ):
+        super().__init__()
+        reversed_ch = list(reversed(block_out_channels))  # [512, 512, 256, 128]
+
+        # Input projection
+        self.conv_in = nn.Conv2d(latent_channels, reversed_ch[0], kernel_size=3, padding=1)
+
+        # Mid block
+        self.mid_block = MidBlock2D(reversed_ch[0], norm_num_groups)
+
+        # Up blocks (3 upsamples → total 8× spatial increase)
+        self.up_blocks = []
+        for i, out_ch in enumerate(reversed_ch):
+            in_ch = reversed_ch[i - 1] if i > 0 else reversed_ch[0]
+            add_upsample = i < len(reversed_ch) - 1
+            self.up_blocks.append(
+                UpDecoderBlock2D(
+                    in_channels=in_ch,
+                    out_channels=out_ch,
+                    num_layers=layers_per_block + 1,
+                    add_upsample=add_upsample,
+                    groups=norm_num_groups,
+                )
+            )
+
+        # Output
+        self.conv_norm_out = _gn(norm_num_groups, reversed_ch[-1])
+        self.conv_out = nn.Conv2d(reversed_ch[-1], 3, kernel_size=3, padding=1)
+
+    def __call__(self, z: mx.array) -> mx.array:
+        """Decode latents → image.
+
+        Args:
+            z: (B, H, W, C) latent tensor in NHWC, **already scaled**.
+
+        Returns:
+            (B, 8H, 8W, 3) decoded image.
+        """
+        x = self.conv_in(z)
+        x = self.mid_block(x)
+        for block in self.up_blocks:
+            x = block(x)
+        x = nn.silu(self.conv_norm_out(x))
+        x = self.conv_out(x)
+        return x

config.json

@@ -0,0 +1,57 @@
+{
+  "pipeline_type": "ZImageMLXPipeline",
+  "base_model": "Tongyi-MAI/Z-Image-Turbo",
+  "framework": "mlx",
+  "model": {
+    "total_params": "10.26B",
+    "text_encoder": {
+      "type": "Qwen3",
+      "params": "4.02B",
+      "hidden_size": 2560,
+      "num_layers": 36,
+      "num_attention_heads": 32,
+      "num_key_value_heads": 8,
+      "dtype": "bfloat16"
+    },
+    "transformer": {
+      "type": "ZImageTransformer (S3-DiT)",
+      "params": "6.15B",
+      "dim": 3840,
+      "n_heads": 30,
+      "head_dim": 128,
+      "n_layers": 30,
+      "n_refiner_layers": 2,
+      "ffn_dim": 10240,
+      "in_channels": 16,
+      "patch_size": 2,
+      "dtype": "bfloat16"
+    },
+    "vae": {
+      "type": "AutoencoderKL Decoder",
+      "params": "84M",
+      "latent_channels": 16,
+      "block_out_channels": [128, 256, 512, 512],
+      "scaling_factor": 0.3611,
+      "shift_factor": 0.1159,
+      "dtype": "float32"
+    },
+    "scheduler": {
+      "type": "FlowMatchEulerDiscrete",
+      "shift": 3.0,
+      "num_train_timesteps": 1000
+    }
+  },
+  "quantization": {
+    "supported_bits": [4, 8, 16],
+    "default_bits": 16,
+    "group_size": 64,
+    "min_quantize_dim": 1024
+  },
+  "generation_defaults": {
+    "width": 512,
+    "height": 512,
+    "num_steps": 8,
+    "guidance_scale": 0.0,
+    "max_text_len": 256
+  }
+}

pipeline.py

@@ -0,0 +1,323 @@
+"""Z-Image-Turbo MLX Pipeline — end-to-end text-to-image generation.
+
+Flow:
+    1. Tokenize prompt → token IDs
+    2. Qwen3 Encoder → text hidden states  (MLX)
+    3. Initialize random latents
+    4. Denoise loop: DiT forward pass × N steps  (MLX)
+    5. VAE decode latents → RGB image  (MLX native)
+    6. Save to PNG
+"""
+
+from __future__ import annotations
+
+import logging
+import time
+from pathlib import Path
+
+import mlx.core as mx
+import numpy as np
+from PIL import Image
+
+from .autoencoder import Decoder
+from .qwen3_encoder import Qwen3Encoder, Qwen3EncoderConfig
+from .zimage_dit import ZImageTransformer, ZImageDiTConfig
+from .scheduler import FlowMatchEulerScheduler
+from .tokenizer import Qwen2Tokenizer
+from .weight_loader import (
+    _find_model_path,
+    _log_memory,
+    load_text_encoder_weights,
+    load_transformer_weights,
+    load_vae_decoder_weights,
+)
+
+logger = logging.getLogger("zimage-mlx")
+
+
+def _cast_to_bf16(model):
+    """Cast all parameters of an nn.Module to bfloat16 in-place.
+
+    This halves memory and speeds up Metal compute for the DiT.
+    """
+    from mlx.utils import tree_map
+    params = model.parameters()
+    bf16_params = tree_map(lambda x: x.astype(mx.bfloat16) if isinstance(x, mx.array) else x, params)
+    model.update(bf16_params)
+    return model
+
+
+class ZImageMLXPipeline:
+    """End-to-end Z-Image-Turbo inference pipeline — 100% MLX.
+
+    All stages run on Apple Silicon via MLX: text encoding,
+    DiT denoising, and VAE decoding.  No PyTorch dependency.
+    """
+
+    def __init__(self, model_id: str = "Tongyi-MAI/Z-Image-Turbo"):
+        self.model_id = model_id
+        self._model_path: Path | None = None
+        self._tokenizer: Qwen2Tokenizer | None = None
+        self._encoder: Qwen3Encoder | None = None
+        self._dit: ZImageTransformer | None = None
+        self._dit_compiled = None  # mx.compile'd forward pass
+        self._scheduler = FlowMatchEulerScheduler(shift=3.0)
+        self._vae: Decoder | None = None
+        self._loaded = False
+
+    def load(self, model_path: Path | None = None):
+        """Load all model components.
+
+        Memory strategy (staged loading):
+        - Encoder, DiT, VAE are loaded sequentially.
+        - During generation, encoder is released after text encoding
+          to reduce peak memory (see ``generate()``).
+        """
+        t0 = time.monotonic()
+        self._model_path = model_path or _find_model_path(self.model_id)
+        _log_memory("before load")
+
+        # 1. Tokenizer
+        logger.info("[ZImage-MLX] Loading tokenizer...")
+        self._tokenizer = Qwen2Tokenizer(self._model_path)
+
+        # 2. Text encoder (Qwen3)
+        logger.info("[ZImage-MLX] Loading text encoder (Qwen3, 36 layers)...")
+        self._encoder = Qwen3Encoder(Qwen3EncoderConfig())
+        te_weights = load_text_encoder_weights(self._model_path)
+        self._encoder.load_weights(list(te_weights.items()))
+        # Weights are already bfloat16 on disk; keep them as-is for memory savings
+        mx.eval(self._encoder.parameters())
+        del te_weights  # release weight dict immediately
+        logger.info("[ZImage-MLX] Text encoder loaded (bfloat16)")
+        _log_memory("after text encoder")
+
+        # 3. DiT (ZImageTransformer)
+        logger.info("[ZImage-MLX] Loading transformer (S3-DiT, 30+2+2 layers)...")
+        self._dit = ZImageTransformer(ZImageDiTConfig())
+        dit_weights = load_transformer_weights(self._model_path)
+        self._dit.load_weights(list(dit_weights.items()))
+        # Cast DiT to bfloat16 for faster inference (~2× speedup, ~50% memory)
+        # PyTorch diffusers also runs at bfloat16 on MPS
+        self._dit = _cast_to_bf16(self._dit)
+        mx.eval(self._dit.parameters())
+        del dit_weights  # release weight dict immediately
+        # Compile DiT forward for additional Metal kernel fusion speedup
+        self._dit_compiled = mx.compile(self._dit)
+        logger.info("[ZImage-MLX] Transformer loaded (bfloat16 + compiled)")
+        _log_memory("after transformer")
+
+        # 4. VAE decoder (MLX native)
+        logger.info("[ZImage-MLX] Loading VAE decoder...")
+        self._vae = Decoder()
+        vae_weights = load_vae_decoder_weights(self._model_path)
+        self._vae.load_weights(vae_weights)
+        mx.eval(self._vae.parameters())
+        del vae_weights  # release weight list immediately
+        logger.info("[ZImage-MLX] VAE decoder loaded")
+
+        elapsed = time.monotonic() - t0
+        self._loaded = True
+        _log_memory("after full load")
+        logger.info("[ZImage-MLX] Pipeline loaded in %.1fs", elapsed)
+
+
+
+    def generate(
+        self,
+        prompt: str,
+        width: int = 768,
+        height: int = 768,
+        num_steps: int = 8,
+        seed: int | None = None,
+        guidance_scale: float = 0.0,  # Z-Image-Turbo typically uses 0
+        max_text_len: int = 256,
+    ) -> np.ndarray:
+        """Generate an image from a text prompt.
+
+        Args:
+            prompt: Text description (Chinese or English)
+            width: Output width (must be divisible by 16)
+            height: Output height (must be divisible by 16)
+            num_steps: Number of denoising steps
+            seed: Random seed (None for random)
+            guidance_scale: CFG scale (0 = no guidance)
+            max_text_len: Max text token length
+
+        Returns:
+            RGB image as numpy array (H, W, 3) uint8
+        """
+        if not self._loaded:
+            raise RuntimeError("Pipeline not loaded. Call load() first.")
+
+        t0 = time.monotonic()
+        if seed is None:
+            seed = int(time.time()) % (2**31)
+
+        # Ensure encoder is available (may have been released after prev gen)
+        self._reload_encoder()
+
+        # ── 1. Tokenize (with chat template like diffusers) ──
+        chat_result = self._tokenizer.apply_chat_template(prompt, max_length=max_text_len)
+        token_ids = chat_result["input_ids"]   # list[int]
+        attn_mask = chat_result["attention_mask"]  # list[int]
+
+        input_ids = mx.array([token_ids])      # (1, L)
+
+        # ── 2. Text encode ──
+        t_enc = time.monotonic()
+        if self._encoder is None:
+            raise RuntimeError("Text encoder not loaded. Call load() first.")
+        all_hidden = self._encoder(input_ids)    # (1, L, 2560) — bfloat16
+        cap_feats = all_hidden  # (1, L, 2560)
+        mx.eval(cap_feats)
+        logger.info("[ZImage-MLX] Text encoded in %.2fs, %d tokens", time.monotonic() - t_enc, cap_feats.shape[1])
+
+        # Release encoder to free memory before DiT denoising.
+        self._release_encoder()
+
+        # ── 3. Initialize latents ──
+        latent_h = height // 8
+        latent_w = width // 8
+        mx.random.seed(seed)
+        # Use bfloat16 latents to match DiT precision
+        latents = mx.random.normal((1, 16, latent_h, latent_w)).astype(mx.bfloat16)
+
+        # Ensure cap_feats is bfloat16 for DiT
+        cap_feats = cap_feats.astype(mx.bfloat16)
+
+        # ── 4. Denoise loop ──
+        sigmas = self._scheduler.get_sigmas(num_steps)
+        mx.eval(sigmas)
+        sigmas_list = sigmas.tolist()
+
+        dit_fn = self._dit_compiled if self._dit_compiled is not None else self._dit
+
+        t_denoise = time.monotonic()
+        for i in range(num_steps):
+            sigma = sigmas_list[i]
+            sigma_next = sigmas_list[i + 1]
+
+            t_step = mx.array([1.0 - sigma], dtype=mx.bfloat16)
+            noise_pred = dit_fn(latents, t_step, cap_feats)
+
+            # Diffusers negates the model output before passing to scheduler:
+            #   noise_pred = -noise_pred
+            # Then scheduler does: prev_sample = sample + (sigma_next - sigma) * model_output
+            noise_pred = -noise_pred
+
+            latents = self._scheduler.step(noise_pred, sigma, sigma_next, latents)
+            mx.eval(latents)
+
+            logger.info("[ZImage-MLX] Step %d/%d (sigma %.4f → %.4f)", i + 1, num_steps, sigma, sigma_next)
+
+        denoise_time = time.monotonic() - t_denoise
+        logger.info("[ZImage-MLX] Denoised in %.2fs (%.2fs/step)", denoise_time, denoise_time / num_steps)
+
+        # ── 5. VAE decode (MLX native) ──
+        t_vae = time.monotonic()
+        image = self._vae_decode(latents)
+        logger.info("[ZImage-MLX] VAE decoded in %.2fs", time.monotonic() - t_vae)
+
+        total = time.monotonic() - t0
+        logger.info("[ZImage-MLX] Total generation: %.2fs", total)
+
+        return image
+
+    def _vae_decode(self, latents: mx.array) -> np.ndarray:
+        """Decode latents → RGB image using MLX VAE.
+
+        Diffusers formula:
+            z = latents / scaling_factor + shift_factor
+            raw = vae.decode(z)          # output in [-1, 1]
+            image = raw / 2 + 0.5        # denormalize to [0, 1]
+        """
+        scaling_factor = 0.3611
+        shift_factor = 0.1159
+
+        # NCHW → NHWC for MLX convolutions
+        z = latents.transpose(0, 2, 3, 1)           # (B,C,H,W) → (B,H,W,C)
+        z = z.astype(mx.float32)                     # force_upcast
+        z = z / scaling_factor + shift_factor
+
+        decoded = self._vae(z)                       # (B,8H,8W,3) in [-1, 1]
+        mx.eval(decoded)
+
+        # Denormalize [-1,1] → [0,1], then clamp → uint8
+        img = decoded[0] / 2.0 + 0.5
+        img = mx.clip(img, 0.0, 1.0)
+        img = np.array(img)
+        img = (img * 255).astype(np.uint8)
+        return img
+
+    def generate_and_save(
+        self,
+        prompt: str,
+        output_path: str,
+        width: int = 768,
+        height: int = 768,
+        num_steps: int = 8,
+        seed: int | None = None,
+    ) -> dict:
+        """Generate an image and save to file.
+
+        Returns:
+            Dict with generation metadata.
+        """
+        t0 = time.monotonic()
+        if seed is None:
+            seed = int(time.time()) % (2**31)
+
+        image = self.generate(
+            prompt=prompt,
+            width=width,
+            height=height,
+            num_steps=num_steps,
+            seed=seed,
+        )
+
+        # Save
+        img = Image.fromarray(image)
+        img.save(output_path)
+
+        elapsed = time.monotonic() - t0
+        return {
+            "image_path": output_path,
+            "width": width,
+            "height": height,
+            "seed": seed,
+            "num_steps": num_steps,
+            "elapsed_s": round(elapsed, 2),
+            "prompt": prompt,
+        }
+
+    def _release_encoder(self):
+        """Release text encoder to free ~5 GB before denoising."""
+        if self._encoder is not None:
+            self._encoder = None
+            mx.clear_cache()
+            _log_memory("after releasing encoder")
+
+    def _reload_encoder(self):
+        """Reload encoder for next generation (lazy, on-demand)."""
+        if self._encoder is None and self._model_path is not None:
+            logger.info("[ZImage-MLX] Reloading text encoder...")
+            self._encoder = Qwen3Encoder(Qwen3EncoderConfig())
+            te_weights = load_text_encoder_weights(self._model_path)
+            self._encoder.load_weights(list(te_weights.items()))
+            # Weights are bfloat16 on disk; keep as-is
+            mx.eval(self._encoder.parameters())
+            del te_weights
+            _log_memory("after reloading encoder")
+
+    def unload(self):
+        """Release all model memory."""
+        self._encoder = None
+        self._dit = None
+        self._vae = None
+        self._tokenizer = None
+        self._loaded = False
+        mx.clear_cache()
+        _log_memory("after full unload")
+        mx.clear_cache()
+        logger.info("[ZImage-MLX] Pipeline unloaded")

qwen3_encoder.py

@@ -0,0 +1,266 @@
+"""Qwen3 Text Encoder — MLX native implementation for Z-Image-Turbo.
+
+Architecture (from model config):
+  - 36 layers, hidden_size=2560, 32 attention heads, 8 KV heads (GQA 4:1)
+  - head_dim=128, intermediate_size=9728
+  - hidden_act=silu (SwiGLU FFN)
+  - RMSNorm (eps=1e-6), QK-Norm on q/k projections
+  - RoPE (theta=1_000_000)
+  - vocab_size=151936
+
+Weight key pattern:
+  model.embed_tokens.weight
+  model.layers.N.input_layernorm.weight
+  model.layers.N.self_attn.{q_proj,k_proj,v_proj,o_proj}.weight
+  model.layers.N.self_attn.{q_norm,k_norm}.weight
+  model.layers.N.post_attention_layernorm.weight
+  model.layers.N.mlp.{gate_proj,up_proj,down_proj}.weight
+  model.norm.weight
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+# ── Config ────────────────────────────────────────────────────────
+
+@dataclass
+class Qwen3EncoderConfig:
+    hidden_size: int = 2560
+    num_hidden_layers: int = 36
+    num_attention_heads: int = 32
+    num_key_value_heads: int = 8
+    head_dim: int = 128
+    intermediate_size: int = 9728
+    rms_norm_eps: float = 1e-6
+    rope_theta: float = 1_000_000.0
+    vocab_size: int = 151936
+    max_position_embeddings: int = 40960
+
+
+# ── RMSNorm ───────────────────────────────────────────────────────
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-6):
+        super().__init__()
+        self.weight = mx.ones((dim,))
+        self.eps = eps
+
+    def __call__(self, x: mx.array) -> mx.array:
+        rms = mx.rsqrt(mx.mean(x * x, axis=-1, keepdims=True) + self.eps)
+        return x * rms * self.weight
+
+
+# ── RoPE ──────────────────────────────────────────────────────────
+
+class RotaryEmbedding(nn.Module):
+    def __init__(self, dim: int, theta: float = 1_000_000.0, max_seq_len: int = 8192):
+        super().__init__()
+        self.dim = dim
+        self.theta = theta
+        inv_freq = 1.0 / (theta ** (mx.arange(0, dim, 2, dtype=mx.float32) / dim))
+        self._inv_freq = inv_freq
+        self._max_cached = 0
+        self._cos_cache = None
+        self._sin_cache = None
+
+    def _update_cache(self, seq_len: int):
+        if seq_len <= self._max_cached and self._cos_cache is not None:
+            return
+        t = mx.arange(seq_len, dtype=mx.float32)
+        freqs = mx.outer(t, self._inv_freq)
+        emb = mx.concatenate([freqs, freqs], axis=-1)
+        self._cos_cache = mx.cos(emb)
+        self._sin_cache = mx.sin(emb)
+        self._max_cached = seq_len
+
+    def __call__(self, seq_len: int) -> tuple[mx.array, mx.array]:
+        self._update_cache(seq_len)
+        return self._cos_cache[:seq_len], self._sin_cache[:seq_len]
+
+
+def _rotate_half(x: mx.array) -> mx.array:
+    x1, x2 = mx.split(x, 2, axis=-1)
+    return mx.concatenate([-x2, x1], axis=-1)
+
+
+def apply_rotary_pos_emb(q: mx.array, k: mx.array, cos: mx.array, sin: mx.array) -> tuple[mx.array, mx.array]:
+    # q/k shape: (B, heads, L, head_dim)
+    # cos/sin: (seq_len, head_dim) → (1, 1, seq_len, head_dim)
+    cos = cos[None, None, :, :]
+    sin = sin[None, None, :, :]
+    q_rot = q * cos + _rotate_half(q) * sin
+    k_rot = k * cos + _rotate_half(k) * sin
+    return q_rot, k_rot
+
+
+# ── Attention ─────────────────────────────────────────────────────
+
+class Qwen3Attention(nn.Module):
+    def __init__(self, cfg: Qwen3EncoderConfig):
+        super().__init__()
+        self.n_heads = cfg.num_attention_heads
+        self.n_kv_heads = cfg.num_key_value_heads
+        self.head_dim = cfg.head_dim
+        self.n_rep = self.n_heads // self.n_kv_heads  # GQA repeat factor
+
+        self.q_proj = nn.Linear(cfg.hidden_size, self.n_heads * self.head_dim, bias=False)
+        self.k_proj = nn.Linear(cfg.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.v_proj = nn.Linear(cfg.hidden_size, self.n_kv_heads * self.head_dim, bias=False)
+        self.o_proj = nn.Linear(self.n_heads * self.head_dim, cfg.hidden_size, bias=False)
+
+        # QK-Norm
+        self.q_norm = RMSNorm(self.head_dim, eps=cfg.rms_norm_eps)
+        self.k_norm = RMSNorm(self.head_dim, eps=cfg.rms_norm_eps)
+
+    def __call__(
+        self,
+        x: mx.array,
+        cos: mx.array,
+        sin: mx.array,
+        mask: mx.array | None = None,
+    ) -> mx.array:
+        B, L, _ = x.shape
+
+        q = self.q_proj(x).reshape(B, L, self.n_heads, self.head_dim)
+        k = self.k_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)
+        v = self.v_proj(x).reshape(B, L, self.n_kv_heads, self.head_dim)
+
+        # QK-Norm (per-head)
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        # Transpose to (B, heads, L, head_dim) for RoPE
+        q = q.transpose(0, 2, 1, 3)
+        k = k.transpose(0, 2, 1, 3)
+        v = v.transpose(0, 2, 1, 3)
+
+        # Apply RoPE
+        q, k = apply_rotary_pos_emb(q, k, cos, sin)
+
+        # GQA: repeat KV heads
+        if self.n_rep > 1:
+            k = mx.repeat(k, self.n_rep, axis=1)
+            v = mx.repeat(v, self.n_rep, axis=1)
+
+        # Scaled dot-product attention
+        scale = 1.0 / math.sqrt(self.head_dim)
+        attn = (q @ k.transpose(0, 1, 3, 2)) * scale
+
+        if mask is not None:
+            attn = attn + mask
+
+        attn = mx.softmax(attn, axis=-1)
+        out = (attn @ v).transpose(0, 2, 1, 3).reshape(B, L, -1)
+
+        return self.o_proj(out)
+
+
+# ── MLP (SwiGLU) ─────────────────────────────────────────────────
+
+class Qwen3MLP(nn.Module):
+    def __init__(self, cfg: Qwen3EncoderConfig):
+        super().__init__()
+        self.gate_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
+        self.up_proj = nn.Linear(cfg.hidden_size, cfg.intermediate_size, bias=False)
+        self.down_proj = nn.Linear(cfg.intermediate_size, cfg.hidden_size, bias=False)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+# ── Transformer Layer ────────────────────────────────────────────
+
+class Qwen3DecoderLayer(nn.Module):
+    def __init__(self, cfg: Qwen3EncoderConfig):
+        super().__init__()
+        self.input_layernorm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps)
+        self.self_attn = Qwen3Attention(cfg)
+        self.post_attention_layernorm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps)
+        self.mlp = Qwen3MLP(cfg)
+
+    def __call__(
+        self,
+        x: mx.array,
+        cos: mx.array,
+        sin: mx.array,
+        mask: mx.array | None = None,
+    ) -> mx.array:
+        # Pre-norm attention
+        h = self.input_layernorm(x)
+        h = self.self_attn(h, cos, sin, mask)
+        x = x + h
+
+        # Pre-norm FFN
+        h = self.post_attention_layernorm(x)
+        h = self.mlp(h)
+        x = x + h
+
+        return x
+
+
+# ── Full Encoder ─────────────────────────────────────────────────
+
+class Qwen3Encoder(nn.Module):
+    """Qwen3 text encoder for Z-Image-Turbo.
+
+    Uses the model as an encoder: runs all 36 layers, returns the
+    final hidden states (no causal mask, no generation).
+    """
+
+    def __init__(self, cfg: Qwen3EncoderConfig | None = None):
+        super().__init__()
+        if cfg is None:
+            cfg = Qwen3EncoderConfig()
+        self.cfg = cfg
+
+        self.embed_tokens = nn.Embedding(cfg.vocab_size, cfg.hidden_size)
+        self.layers = [Qwen3DecoderLayer(cfg) for _ in range(cfg.num_hidden_layers)]
+        self.norm = RMSNorm(cfg.hidden_size, eps=cfg.rms_norm_eps)
+        self.rotary_emb = RotaryEmbedding(
+            dim=cfg.head_dim,
+            theta=cfg.rope_theta,
+            max_seq_len=cfg.max_position_embeddings,
+        )
+
+    def __call__(self, input_ids: mx.array, mask: mx.array | None = None) -> mx.array:
+        """Encode text tokens.
+
+        Returns the second-to-last hidden state (hidden_states[-2]),
+        matching diffusers ZImagePipeline which uses
+        ``text_encoder(..., output_hidden_states=True).hidden_states[-2]``.
+
+        Applies a causal attention mask by default (matching HuggingFace
+        Qwen3Model which uses causal masking internally).
+
+        Args:
+            input_ids: (B, L) token IDs
+            mask: optional attention mask (B, 1, L, L) — None = auto causal mask
+
+        Returns:
+            hidden_states: (B, L, hidden_size) — penultimate layer output
+        """
+        B, L = input_ids.shape
+        x = self.embed_tokens(input_ids)
+
+        cos, sin = self.rotary_emb(L)
+
+        # Build causal mask if none provided (matches HuggingFace Qwen3Model)
+        if mask is None:
+            mask = mx.full((L, L), -1e9)
+            mask = mx.triu(mask, k=1)             # upper triangle = -inf
+            mask = mask[None, None, :, :]          # (1, 1, L, L)
+
+        n_layers = len(self.layers)
+        for i, layer in enumerate(self.layers):
+            x = layer(x, cos, sin, mask)
+            if i == n_layers - 2:
+                # Capture second-to-last layer output (no final norm)
+                penultimate = x
+
+        return penultimate

scheduler.py

@@ -0,0 +1,79 @@
+"""FlowMatch Euler Discrete Scheduler for Z-Image-Turbo.
+
+Implements the ODE-based Flow Matching sampling schedule used by
+Z-Image-Turbo (same as FLUX-schnell).
+
+Config: shift=3.0, num_train_timesteps=1000
+"""
+
+from __future__ import annotations
+
+import mlx.core as mx
+
+
+class FlowMatchEulerScheduler:
+    """Flow Matching Euler sampler.
+
+    The forward diffusion maps data x0 to noise x1 via:
+        x_t = (1 - t) * x0 + t * noise
+
+    Denoising reverses this with Euler steps from t=1 → t=0.
+    """
+
+    def __init__(
+        self,
+        num_train_timesteps: int = 1000,
+        shift: float = 3.0,
+        sigma_min: float = 0.002994012087583542,
+    ):
+        self.num_train_timesteps = num_train_timesteps
+        self.shift = shift
+        self.sigma_min = sigma_min
+
+    def get_sigmas(self, num_steps: int) -> mx.array:
+        """Compute sigma schedule matching diffusers FlowMatchEulerDiscreteScheduler.
+
+        diffusers logic:
+          1. sigma_max=1.0, sigma_min≈0.00299 (not 0!)
+          2. timesteps = linspace(sigma_max*1000, sigma_min*1000, num_steps)
+          3. sigmas_raw = timesteps / 1000
+          4. sigmas = shift * raw / (1 + (shift-1) * raw)
+          5. append terminal sigma = 0
+        """
+        # Match diffusers: linspace from sigma_max to sigma_min, NOT to 0
+        timesteps = mx.linspace(
+            float(self.num_train_timesteps),
+            float(self.sigma_min * self.num_train_timesteps),
+            num_steps,
+        )
+        raw_sigmas = timesteps / float(self.num_train_timesteps)
+        sigmas = self.shift * raw_sigmas / (1.0 + (self.shift - 1.0) * raw_sigmas)
+        # Append terminal zero
+        sigmas = mx.concatenate([sigmas, mx.array([0.0])])
+        return sigmas
+
+    def unshift_sigma(self, shifted_sigma: float) -> float:
+        """Invert the shift to recover the original (unshifted) sigma.
+
+        shifted = shift * s / (1 + (shift-1) * s)  →  s = shifted / (shift - (shift-1) * shifted)
+        """
+        if self.shift == 1.0:
+            return shifted_sigma
+        denom = self.shift - (self.shift - 1.0) * shifted_sigma
+        if denom == 0:
+            return 1.0
+        return shifted_sigma / denom
+
+    def step(
+        self,
+        model_output: mx.array,
+        sigma: float,
+        sigma_next: float,
+        sample: mx.array,
+    ) -> mx.array:
+        """Single Euler step.
+
+        v-prediction: x_{t-1} = x_t + (sigma_next - sigma) * v_pred
+        """
+        dt = sigma_next - sigma
+        return sample + dt * model_output

tokenizer.py

@@ -0,0 +1,70 @@
+"""Qwen2 Tokenizer adapter for Z-Image-Turbo.
+
+Uses the `tokenizers` library directly for fast BPE tokenization,
+avoiding the slow AutoTokenizer.from_pretrained() initialization.
+"""
+
+from __future__ import annotations
+
+import json
+import logging
+from pathlib import Path
+
+logger = logging.getLogger("zimage-mlx")
+
+
+class Qwen2Tokenizer:
+    """Fast BPE tokenizer using tokenizers library."""
+
+    def __init__(self, model_path: Path):
+        from tokenizers import Tokenizer as HFTokenizer
+
+        tokenizer_path = model_path / "tokenizer"
+        json_file = tokenizer_path / "tokenizer.json"
+        if not json_file.exists():
+            json_file = model_path / "tokenizer.json"
+        if not json_file.exists():
+            raise FileNotFoundError(f"tokenizer.json not found in {model_path}")
+
+        self._tokenizer = HFTokenizer.from_file(str(json_file))
+
+        # Load chat template from tokenizer_config.json if available
+        config_file = tokenizer_path / "tokenizer_config.json"
+        if not config_file.exists():
+            config_file = model_path / "tokenizer_config.json"
+        self._chat_template = None
+        if config_file.exists():
+            with open(config_file) as f:
+                cfg = json.load(f)
+            self._chat_template = cfg.get("chat_template")
+
+        logger.info("[ZImage] Tokenizer loaded: vocab_size=%d", self._tokenizer.get_vocab_size())
+
+    def encode(self, text: str, max_length: int = 512) -> list[int]:
+        """Encode text to token IDs."""
+        encoded = self._tokenizer.encode(text)
+        ids = encoded.ids
+        if len(ids) > max_length:
+            ids = ids[:max_length]
+        return ids
+
+    def apply_chat_template(self, prompt: str, max_length: int = 512) -> dict:
+        """Apply Qwen3 chat template format and tokenize.
+
+        Wraps prompt in chat format:
+            <|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n
+
+        Returns dict with 'input_ids' and 'attention_mask'.
+        """
+        # Build chat-formatted text manually (Qwen3 chat template)
+        chat_text = f"<|im_start|>user\n{prompt}<|im_end|>\n<|im_start|>assistant\n"
+        encoded = self._tokenizer.encode(chat_text)
+        ids = encoded.ids
+        if len(ids) > max_length:
+            ids = ids[:max_length]
+        attn_mask = [1] * len(ids)
+        return {"input_ids": ids, "attention_mask": attn_mask}
+
+    @property
+    def vocab_size(self) -> int:
+        return self._tokenizer.get_vocab_size()

weight_loader.py

@@ -0,0 +1,195 @@
+"""Weight loader for Z-Image-Turbo MLX backend.
+
+Loads safetensors weights from HuggingFace cache and maps them
+to the MLX module hierarchy.
+"""
+
+from __future__ import annotations
+
+import glob
+import logging
+from pathlib import Path
+
+import mlx.core as mx
+
+logger = logging.getLogger("zimage-mlx")
+
+# Default HF cache path for Z-Image-Turbo
+_DEFAULT_MODEL_ID = "Tongyi-MAI/Z-Image-Turbo"
+_HF_CACHE = Path.home() / ".cache" / "huggingface" / "hub"
+
+# Local weights directory (project-local, survives HF cache cleanup)
+_LOCAL_WEIGHTS_DIR = Path(__file__).parent / "weights"
+
+
+def _find_model_path(model_id: str = _DEFAULT_MODEL_ID) -> Path:
+    """Find local weight path for a model.
+
+    Priority:
+      1. Project-local ``backends/mlx_zimage/weights/`` (if text_encoder/ exists)
+      2. HF cache ``~/.cache/huggingface/hub/models--Tongyi-MAI--Z-Image-Turbo/``
+    """
+    # 1. Local weights directory
+    if _LOCAL_WEIGHTS_DIR.is_dir() and (_LOCAL_WEIGHTS_DIR / "text_encoder").is_dir():
+        logger.info("[ZImage] Using local weights: %s", _LOCAL_WEIGHTS_DIR)
+        return _LOCAL_WEIGHTS_DIR
+
+    # 2. HF cache
+    safe_id = model_id.replace("/", "--")
+    model_dir = _HF_CACHE / f"models--{safe_id}"
+    if not model_dir.exists():
+        raise FileNotFoundError(
+            f"Model not found. Neither local ({_LOCAL_WEIGHTS_DIR}) "
+            f"nor HF cache ({model_dir}) available."
+        )
+    # Find the latest snapshot
+    snapshots = sorted(model_dir.glob("snapshots/*"), key=lambda p: p.stat().st_mtime, reverse=True)
+    if not snapshots:
+        raise FileNotFoundError(f"No snapshots found in {model_dir}")
+    logger.info("[ZImage] Using HF cache: %s", snapshots[0])
+    return snapshots[0]
+
+
+def _log_memory(label: str) -> None:
+    """Log Metal memory usage (safe no-op if unavailable)."""
+    try:
+        active = mx.metal.get_active_memory() / (1024 ** 3)
+        peak = mx.metal.get_peak_memory() / (1024 ** 3)
+        logger.info("[ZImage] MEM %s: active=%.2f GB, peak=%.2f GB", label, active, peak)
+    except Exception:
+        pass  # mx.metal not available (e.g. CI / non-Apple)
+
+
+def _load_safetensors_shards(
+    shard_dir: Path,
+    pattern: str = "*.safetensors",
+    *,
+    key_filter: str | None = None,
+) -> dict[str, mx.array]:
+    """Load safetensors files via mx.load() — zero-copy, preserves bfloat16.
+
+    Args:
+        shard_dir: Directory containing safetensors shard files.
+        pattern: Glob pattern for shard files.
+        key_filter: If set, only load keys starting with this prefix.
+    """
+    files = sorted(shard_dir.glob(pattern))
+    if not files:
+        raise FileNotFoundError(f"No safetensors files in {shard_dir}")
+
+    params: dict[str, mx.array] = {}
+    for f in files:
+        # mx.load() natively reads safetensors → mx.array (preserves bfloat16)
+        shard = mx.load(str(f))
+        if key_filter:
+            shard = {k: v for k, v in shard.items() if k.startswith(key_filter)}
+        params.update(shard)
+        logger.info("[ZImage] Loaded shard %s (%d keys)", f.name, len(shard))
+
+    logger.info("[ZImage] Total: %d keys from %d files in %s", len(params), len(files), shard_dir.name)
+    _log_memory(f"after loading {shard_dir.name}")
+    return params
+
+
+# ── Text Encoder weight mapping ──────────────────────────────────
+
+def load_text_encoder_weights(model_path: Path | None = None) -> dict[str, mx.array]:
+    """Load and map Qwen3 text encoder weights for MLX.
+
+    The safetensors keys use the pattern:
+        model.embed_tokens.weight
+        model.layers.N.input_layernorm.weight
+        model.layers.N.self_attn.q_proj.weight
+        ...
+        model.norm.weight
+
+    Our MLX module uses:
+        embed_tokens.weight
+        layers.N.input_layernorm.weight
+        layers.N.self_attn.q_proj.weight
+        ...
+        norm.weight
+
+    So we strip the leading "model." prefix.
+    """
+    if model_path is None:
+        model_path = _find_model_path()
+
+    te_dir = model_path / "text_encoder"
+    raw = _load_safetensors_shards(te_dir, "model-*.safetensors")
+
+    mapped: dict[str, mx.array] = {}
+    for key, tensor in raw.items():
+        # Strip "model." prefix
+        if key.startswith("model."):
+            new_key = key[len("model."):]
+        else:
+            new_key = key
+
+        mapped[new_key] = tensor
+
+    logger.info("[ZImage] Text encoder: %d parameters mapped", len(mapped))
+    return mapped
+
+
+# ── Transformer weight mapping ───────────────────────────────────
+
+def load_transformer_weights(model_path: Path | None = None) -> dict[str, mx.array]:
+    """Load ZImageTransformer2DModel weights."""
+    if model_path is None:
+        model_path = _find_model_path()
+
+    dit_dir = model_path / "transformer"
+    raw = _load_safetensors_shards(dit_dir, "diffusion_pytorch_model-*.safetensors")
+
+    # Keys are already flat (no "model." prefix), use as-is
+    logger.info("[ZImage] Transformer: %d parameters loaded", len(raw))
+    return raw
+
+
+# ── VAE weight mapping ───────────────────────────────────────────
+
+def load_vae_weights(model_path: Path | None = None) -> dict[str, mx.array]:
+    """Load AutoencoderKL weights."""
+    if model_path is None:
+        model_path = _find_model_path()
+
+    vae_dir = model_path / "vae"
+    raw = _load_safetensors_shards(vae_dir)
+
+    logger.info("[ZImage] VAE: %d parameters loaded", len(raw))
+    return raw
+
+
+def load_vae_decoder_weights(model_path: Path | None = None) -> list[tuple[str, mx.array]]:
+    """Load VAE decoder weights, mapped for the MLX Decoder module.
+
+    Only loads keys starting with ``decoder.`` (skips encoder weights
+    to avoid wasting memory).  Performs two transformations:
+    1. Strips the ``decoder.`` prefix so keys match the Decoder module tree.
+    2. Transposes Conv2d weights from PyTorch (O,I,kH,kW) → MLX (O,kH,kW,I).
+
+    Returns a list of (key, array) tuples ready for ``Decoder.load_weights()``.
+    """
+    if model_path is None:
+        model_path = _find_model_path()
+
+    vae_dir = model_path / "vae"
+    # Only load decoder.* keys — skip encoder weights entirely
+    raw = _load_safetensors_shards(vae_dir, key_filter="decoder.")
+
+    weights: list[tuple[str, mx.array]] = []
+    for key, val in raw.items():
+        key = key[len("decoder."):]
+
+        # Conv2d weight: (O, I, kH, kW) → (O, kH, kW, I)
+        if val.ndim == 4:
+            val = val.transpose(0, 2, 3, 1)
+
+        # force_upcast: ensure float32 for numerical stability
+        val = val.astype(mx.float32)
+
+        weights.append((key, val))
+
+    logger.info("[ZImage] VAE decoder: %d parameters mapped", len(weights))
+    return weights

zimage_dit.py

@@ -0,0 +1,606 @@
+"""ZImageTransformer2DModel — MLX native S3-DiT for Z-Image-Turbo.
+
+Architecture (from model config + weight shapes):
+  - 30 main DiT layers + 2 context_refiner + 2 noise_refiner
+  - dim=3840, n_heads=30, head_dim=128
+  - Dual-norm (pre+post) for both attention and FFN
+  - SwiGLU FFN (w1/w2/w3), intermediate=10240
+  - QK-Norm (RMSNorm on head_dim=128)
+  - AdaLN modulation: 4 outputs per block (shift_attn, scale_attn, shift_ffn, scale_ffn)
+  - N-dim RoPE: axes_dims=[32,48,48], rope_theta=256
+  - Timestep embedding: sinusoidal(256) → MLP(256→1024→256)
+  - Caption projector: RMSNorm(2560) → Linear(2560→3840)
+  - Patch embed: Linear(64→3840)  (in_channels=16, patch_size=2 → 16×2²=64)
+  - Final layer: adaLN(256→3840) + Linear(3840→64)
+
+Weight key patterns:
+  t_embedder.mlp.{0,2}.{weight,bias}
+  cap_embedder.{0,1}.{weight,bias}     (0=RMSNorm, 1=Linear)
+  cap_pad_token, x_pad_token
+  all_x_embedder.2-1.{weight,bias}
+  layers.N.{adaLN_modulation.0, attention.*, attention_norm*, feed_forward.*, ffn_norm*}
+  context_refiner.N.{attention.*, attention_norm*, feed_forward.*, ffn_norm*}
+  noise_refiner.N.{adaLN_modulation.0, attention.*, attention_norm*, feed_forward.*, ffn_norm*}
+  all_final_layer.2-1.{linear, adaLN_modulation.1}
+"""
+
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+
+import mlx.core as mx
+import mlx.nn as nn
+
+
+# ── Config ────────────────────────────────────────────────────────
+
+@dataclass
+class ZImageDiTConfig:
+    dim: int = 3840
+    n_heads: int = 30
+    n_kv_heads: int = 30
+    n_layers: int = 30
+    n_refiner_layers: int = 2
+    head_dim: int = 128
+    ffn_dim: int = 10240
+    in_channels: int = 16
+    patch_size: int = 2
+    cap_feat_dim: int = 2560          # Qwen3 hidden_size
+    t_embed_dim: int = 256            # timestep embedding dim
+    t_hidden_dim: int = 1024          # timestep MLP hidden
+    axes_dims: list[int] = field(default_factory=lambda: [32, 48, 48])
+    axes_lens: list[int] = field(default_factory=lambda: [1536, 512, 512])
+    rope_theta: float = 256.0
+    norm_eps: float = 1e-5
+    qk_norm: bool = True
+    t_scale: float = 1000.0
+
+
+# ── RMSNorm ───────────────────────────────────────────────────────
+
+class RMSNorm(nn.Module):
+    def __init__(self, dim: int, eps: float = 1e-5):
+        super().__init__()
+        self.weight = mx.ones((dim,))
+        self.eps = eps
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return x * mx.rsqrt(mx.mean(x * x, axis=-1, keepdims=True) + self.eps) * self.weight
+
+
+# ── Timestep Embedding ────────────────────────────────────────────
+
+def timestep_embedding(t: mx.array, dim: int = 256) -> mx.array:
+    """Sinusoidal timestep embedding."""
+    half = dim // 2
+    freqs = mx.exp(-math.log(10000.0) * mx.arange(half, dtype=mx.float32) / half)
+    args = t[:, None].astype(mx.float32) * freqs[None, :]
+    return mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
+
+
+class TimestepEmbedder(nn.Module):
+    """Sinusoidal → MLP timestep embedder: sin(t) → Linear → SiLU → Linear."""
+    def __init__(self, t_embed_dim: int = 256, hidden_dim: int = 1024):
+        super().__init__()
+        self.mlp = [
+            nn.Linear(t_embed_dim, hidden_dim),      # mlp.0
+            None,                                      # SiLU (index 1, not a layer)
+            nn.Linear(hidden_dim, t_embed_dim),        # mlp.2
+        ]
+
+    def __call__(self, t: mx.array) -> mx.array:
+        x = timestep_embedding(t, self.mlp[0].weight.shape[1])
+        x = nn.silu(self.mlp[0](x))
+        x = self.mlp[2](x)
+        return x
+
+
+# ── N-dim RoPE (matches diffusers RopeEmbedder) ──────────────────
+
+class RopeEmbedder:
+    """Precomputed per-axis frequency tables, indexed by position IDs.
+
+    Matches diffusers ``RopeEmbedder``:
+        1. Precompute complex frequencies per axis (as real angle tables here)
+        2. At forward time, gather from tables using integer position IDs
+        3. Concatenate per-axis results → (seq_len, sum(axes_dims)//2)
+
+    The returned angles are used with :func:`apply_rope` which does the
+    equivalent of ``torch.view_as_complex(x) * polar(1, angles)`` using
+    real-valued cos/sin operations.
+    """
+
+    def __init__(
+        self,
+        axes_dims: list[int],
+        axes_lens: list[int],
+        theta: float = 256.0,
+    ):
+        self.axes_dims = axes_dims
+        self.axes_lens = axes_lens
+        self.theta = theta
+        # Precompute per-axis frequency tables
+        self._freq_tables: list[mx.array] = []
+        for d, e in zip(axes_dims, axes_lens):
+            inv_freq = 1.0 / (theta ** (mx.arange(0, d, 2, dtype=mx.float32) / d))
+            timestep = mx.arange(e, dtype=mx.float32)
+            freqs = mx.outer(timestep, inv_freq)  # (e, d/2)
+            self._freq_tables.append(freqs)
+
+    def __call__(self, pos_ids: mx.array) -> mx.array:
+        """Look up RoPE angles from precomputed tables.
+
+        Args:
+            pos_ids: (seq_len, 3) integer position IDs — one per axis.
+
+        Returns:
+            (seq_len, rope_half_dim) rotation angles.
+        """
+        parts = []
+        for i in range(len(self.axes_dims)):
+            idx = pos_ids[:, i].astype(mx.int32)
+            parts.append(self._freq_tables[i][idx])  # (seq_len, d_i/2)
+        return mx.concatenate(parts, axis=-1)
+
+
+def build_position_ids(
+    cap_len: int,
+    pH: int,
+    pW: int,
+) -> tuple[mx.array, mx.array]:
+    """Build position ID grids matching diffusers patchify_and_embed.
+
+    Caption tokens: ``create_coordinate_grid(size=(cap_len, 1, 1), start=(1, 0, 0))``
+        → t-axis = 1..cap_len, h-axis = 0, w-axis = 0
+
+    Image tokens: ``create_coordinate_grid(size=(1, pH, pW), start=(cap_len+1, 0, 0))``
+        → t-axis = cap_len+1, h-axis = 0..pH-1, w-axis = 0..pW-1
+
+    Returns:
+        (img_pos_ids, cap_pos_ids) each of shape (N, 3)
+    """
+    # Caption: (cap_len, 3) — t varies, h=0, w=0
+    cap_t = mx.arange(1, cap_len + 1, dtype=mx.int32)[:, None]  # (cap_len, 1)
+    cap_hw = mx.zeros((cap_len, 2), dtype=mx.int32)
+    cap_pos = mx.concatenate([cap_t, cap_hw], axis=-1)  # (cap_len, 3)
+
+    # Image: (pH*pW, 3) — t=cap_len+1, h and w vary
+    t_val = cap_len + 1
+    img_ids = []
+    for h in range(pH):
+        for w in range(pW):
+            img_ids.append([t_val, h, w])
+    img_pos = mx.array(img_ids, dtype=mx.int32)  # (pH*pW, 3)
+
+    return img_pos, cap_pos
+
+
+def apply_rope(x: mx.array, freqs: mx.array) -> mx.array:
+    """Apply rotary position embedding using interleaved pairing.
+
+    Equivalent to diffusers' complex multiplication:
+        ``x_complex = view_as_complex(x.reshape(..., -1, 2))``
+        ``x_out = view_as_real(x_complex * freqs_cis).flatten()``
+
+    x: (B, n_heads, L, head_dim)
+    freqs: (L, rope_half_dim) where rope_half_dim = sum(axes_dims)//2
+    """
+    rope_half_dim = freqs.shape[-1]
+    rope_dim = rope_half_dim * 2
+    x_rope = x[..., :rope_dim]
+    x_pass = x[..., rope_dim:]
+
+    cos = mx.cos(freqs)[None, None, :, :]  # (1, 1, L, rope_half_dim)
+    sin = mx.sin(freqs)[None, None, :, :]
+
+    # Interleaved pairing: (x[0], x[1]), (x[2], x[3]), ...
+    x_even = x_rope[..., 0::2]  # even indices → "real"
+    x_odd = x_rope[..., 1::2]   # odd indices → "imag"
+
+    out_even = x_even * cos - x_odd * sin
+    out_odd = x_even * sin + x_odd * cos
+
+    # Interleave back: [re0, im0, re1, im1, ...]
+    out = mx.stack([out_even, out_odd], axis=-1)  # (..., rope_half_dim, 2)
+    x_rope = out.reshape(*out.shape[:-2], rope_dim)
+
+    return mx.concatenate([x_rope, x_pass], axis=-1)
+
+
+# ── Attention Block ───────────────────────────────────────────────
+
+class DiTAttention(nn.Module):
+    """Self-attention with QK-Norm and optional RoPE."""
+
+    def __init__(self, dim: int, n_heads: int, head_dim: int, qk_norm: bool = True, norm_eps: float = 1e-5):
+        super().__init__()
+        self.n_heads = n_heads
+        self.head_dim = head_dim
+        self.to_q = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.to_k = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.to_v = nn.Linear(dim, n_heads * head_dim, bias=False)
+        self.to_out = [nn.Linear(n_heads * head_dim, dim, bias=False)]  # to_out.0
+
+        if qk_norm:
+            self.norm_q = RMSNorm(head_dim, eps=norm_eps)
+            self.norm_k = RMSNorm(head_dim, eps=norm_eps)
+        else:
+            self.norm_q = None
+            self.norm_k = None
+
+    def __call__(self, x: mx.array, freqs: mx.array | None = None, mask: mx.array | None = None) -> mx.array:
+        B, L, _ = x.shape
+        q = self.to_q(x).reshape(B, L, self.n_heads, self.head_dim)
+        k = self.to_k(x).reshape(B, L, self.n_heads, self.head_dim)
+        v = self.to_v(x).reshape(B, L, self.n_heads, self.head_dim)
+
+        # QK-Norm
+        if self.norm_q is not None:
+            q = self.norm_q(q)
+            k = self.norm_k(k)
+
+        # (B, n_heads, L, head_dim)
+        q = q.transpose(0, 2, 1, 3)
+        k = k.transpose(0, 2, 1, 3)
+        v = v.transpose(0, 2, 1, 3)
+
+        # RoPE
+        if freqs is not None:
+            q = apply_rope(q, freqs)
+            k = apply_rope(k, freqs)
+
+        # Fused scaled dot-product attention (Metal kernel, no NxN materialization)
+        scale = 1.0 / math.sqrt(self.head_dim)
+        if mask is not None:
+            # Convert boolean mask (B, L) to additive mask for fused attention
+            attn_mask = mask[:, None, None, :].astype(q.dtype)
+            attn_mask = (1.0 - attn_mask) * (-1e9)
+            out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale, mask=attn_mask)
+        else:
+            out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
+
+        out = out.transpose(0, 2, 1, 3).reshape(B, L, -1)
+        return self.to_out[0](out)
+
+
+# ── SwiGLU FFN ────────────────────────────────────────────────────
+
+class SwiGLUFFN(nn.Module):
+    """SwiGLU feed-forward: gate * silu(w1(x)) + w3(x) → w2."""
+    def __init__(self, dim: int, ffn_dim: int):
+        super().__init__()
+        self.w1 = nn.Linear(dim, ffn_dim, bias=False)  # gate
+        self.w2 = nn.Linear(ffn_dim, dim, bias=False)   # down
+        self.w3 = nn.Linear(dim, ffn_dim, bias=False)   # up
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.w2(nn.silu(self.w1(x)) * self.w3(x))
+
+
+# ── AdaLN Modulation ─────────────────────────────────────────────
+
+class AdaLNModulation(nn.Module):
+    """AdaLN-Zero: project conditioning to shift/scale pairs.
+
+    Output dim = dim * n_mods (e.g. 3840 * 4 = 15360 for main blocks).
+    """
+    def __init__(self, cond_dim: int, out_dim: int):
+        super().__init__()
+        # Weight key is adaLN_modulation.0 (index 0 in a Sequential-like list)
+        self._linear = nn.Linear(cond_dim, out_dim)
+
+    # Expose as list for weight loading: adaLN_modulation.0.weight/bias
+    @property
+    def parameters(self):
+        return {"0": {"weight": self._linear.weight, "bias": self._linear.bias}}
+
+    def __call__(self, c: mx.array) -> mx.array:
+        return self._linear(c)
+
+
+# ── DiT Block (main layers + noise_refiner) ──────────────────────
+
+class DiTBlock(nn.Module):
+    """S3-DiT block with AdaLN modulation.
+
+    4 modulations: shift_attn, scale_attn, shift_ffn, scale_ffn
+    Dual-norm: pre-norm + post-norm for both attention and FFN.
+    """
+
+    def __init__(self, cfg: ZImageDiTConfig):
+        super().__init__()
+        self.attention = DiTAttention(cfg.dim, cfg.n_heads, cfg.head_dim, cfg.qk_norm, cfg.norm_eps)
+        self.attention_norm1 = RMSNorm(cfg.dim, eps=cfg.norm_eps)  # pre-attn norm
+        self.attention_norm2 = RMSNorm(cfg.dim, eps=cfg.norm_eps)  # post-attn norm
+        self.feed_forward = SwiGLUFFN(cfg.dim, cfg.ffn_dim)
+        self.ffn_norm1 = RMSNorm(cfg.dim, eps=cfg.norm_eps)        # pre-ffn norm
+        self.ffn_norm2 = RMSNorm(cfg.dim, eps=cfg.norm_eps)        # post-ffn norm
+
+        # AdaLN: 4 modulation signals (shift_a, scale_a, shift_f, scale_f)
+        self.adaLN_modulation = [nn.Linear(cfg.t_embed_dim, cfg.dim * 4)]
+
+    def __call__(self, x: mx.array, c: mx.array, freqs: mx.array | None = None, mask: mx.array | None = None) -> mx.array:
+        """
+        Args:
+            x: (B, L, dim) hidden states
+            c: (B, t_embed_dim) conditioning (timestep embedding)
+            freqs: optional RoPE frequencies for image tokens
+            mask: optional (B, L) boolean attention mask
+        """
+        # Compute modulation from conditioning
+        mod = self.adaLN_modulation[0](c)  # (B, dim*4)
+        scale_msa, gate_msa, scale_mlp, gate_mlp = mx.split(mod, 4, axis=-1)
+
+        gate_msa = mx.tanh(gate_msa)
+        gate_mlp = mx.tanh(gate_mlp)
+        scale_msa = 1.0 + scale_msa
+        scale_mlp = 1.0 + scale_mlp
+
+        scale_msa = scale_msa[:, None, :]
+        gate_msa = gate_msa[:, None, :]
+        scale_mlp = scale_mlp[:, None, :]
+        gate_mlp = gate_mlp[:, None, :]
+
+        attn_out = self.attention(self.attention_norm1(x) * scale_msa, freqs, mask)
+        x = x + gate_msa * self.attention_norm2(attn_out)
+
+        x = x + gate_mlp * self.ffn_norm2(
+            self.feed_forward(self.ffn_norm1(x) * scale_mlp)
+        )
+
+        return x
+
+
+# ── Refiner Block (context_refiner — no AdaLN) ──────────────────
+
+class RefinerBlock(nn.Module):
+    """Refiner block WITHOUT AdaLN modulation (used for context_refiner)."""
+
+    def __init__(self, cfg: ZImageDiTConfig):
+        super().__init__()
+        self.attention = DiTAttention(cfg.dim, cfg.n_heads, cfg.head_dim, cfg.qk_norm, cfg.norm_eps)
+        self.attention_norm1 = RMSNorm(cfg.dim, eps=cfg.norm_eps)
+        self.attention_norm2 = RMSNorm(cfg.dim, eps=cfg.norm_eps)
+        self.feed_forward = SwiGLUFFN(cfg.dim, cfg.ffn_dim)
+        self.ffn_norm1 = RMSNorm(cfg.dim, eps=cfg.norm_eps)
+        self.ffn_norm2 = RMSNorm(cfg.dim, eps=cfg.norm_eps)
+
+    def __call__(self, x: mx.array, freqs: mx.array | None = None, mask: mx.array | None = None) -> mx.array:
+        h = self.attention_norm1(x)
+        h = self.attention(h, freqs, mask)
+        h = self.attention_norm2(h)
+        x = x + h
+
+        h = self.ffn_norm1(x)
+        h = self.feed_forward(h)
+        h = self.ffn_norm2(h)
+        x = x + h
+
+        return x
+
+
+# ── Final Layer ───────────────────────────────────────────────────
+
+class FinalLayer(nn.Module):
+    """Final projection: LayerNorm + adaLN scale + Linear(dim → patch_dim)."""
+    def __init__(self, dim: int, patch_dim: int, t_embed_dim: int):
+        super().__init__()
+        self.linear = nn.Linear(dim, patch_dim)
+        # adaLN_modulation.1 — SiLU + Linear (SiLU at index 0, Linear at index 1)
+        self.adaLN_modulation = [None, nn.Linear(t_embed_dim, dim)]
+
+    def __call__(self, x: mx.array, c: mx.array) -> mx.array:
+        # SiLU is part of FinalLayer's adaLN_modulation (unlike DiTBlock)
+        scale = 1.0 + self.adaLN_modulation[1](nn.silu(c))  # (B, dim)
+        scale = scale[:, None, :]  # (B, 1, dim)
+
+        # LayerNorm (no learnable params) + scale + linear
+        x = mx.fast.layer_norm(x, None, None, eps=1e-6)
+        x = x * scale
+        x = self.linear(x)
+        return x
+
+
+# ── Full ZImage Transformer ──────────────────────────────────────
+
+class ZImageTransformer(nn.Module):
+    """ZImageTransformer2DModel — S3-DiT for Z-Image-Turbo.
+
+    Forward flow:
+        1. Embed timestep → t_emb (B, 256)
+        2. Project caption features: RMSNorm + Linear → cap_emb (B, L_text, 3840)
+        3. Patchify + embed image latents → x_emb (B, L_img, 3840)
+        4. Concatenate [cap_emb, x_emb] → full sequence
+        5. Context refiner (2 blocks, no AdaLN)
+        6. Split → img tokens get RoPE, cap tokens don't
+        7. Main DiT layers (30 blocks, with AdaLN)
+        8. Noise refiner (2 blocks, with AdaLN)
+        9. Extract image tokens → final layer → unpatchify
+    """
+
+    def __init__(self, cfg: ZImageDiTConfig | None = None):
+        super().__init__()
+        if cfg is None:
+            cfg = ZImageDiTConfig()
+        self.cfg = cfg
+
+        # Timestep embedder
+        self.t_embedder = TimestepEmbedder(cfg.t_embed_dim, cfg.t_hidden_dim)
+
+        # Caption projector: cap_embedder.0 = RMSNorm, cap_embedder.1 = Linear
+        self.cap_embedder = [
+            RMSNorm(cfg.cap_feat_dim, eps=cfg.norm_eps),
+            nn.Linear(cfg.cap_feat_dim, cfg.dim),
+        ]
+
+        # Learnable padding tokens
+        self.cap_pad_token = mx.zeros((1, cfg.dim))
+        self.x_pad_token = mx.zeros((1, cfg.dim))
+
+        # Image patch embedder — key uses "2-1" suffix for patch_size=2
+        # We store as a dict to match weight key `all_x_embedder.2-1.{weight,bias}`
+        patch_dim = cfg.in_channels * cfg.patch_size * cfg.patch_size  # 16 * 4 = 64
+        self.all_x_embedder = {"2-1": nn.Linear(patch_dim, cfg.dim)}
+
+        # Context refiner (no AdaLN)
+        self.context_refiner = [RefinerBlock(cfg) for _ in range(cfg.n_refiner_layers)]
+
+        # Main DiT layers (with AdaLN)
+        self.layers = [DiTBlock(cfg) for _ in range(cfg.n_layers)]
+
+        # Noise refiner (with AdaLN)
+        self.noise_refiner = [DiTBlock(cfg) for _ in range(cfg.n_refiner_layers)]
+
+        # Final layer — key uses "2-1" suffix
+        self.all_final_layer = {
+            "2-1": FinalLayer(cfg.dim, patch_dim, cfg.t_embed_dim)
+        }
+
+        # Precomputed RoPE frequency tables (matches diffusers RopeEmbedder)
+        self._rope = RopeEmbedder(cfg.axes_dims, cfg.axes_lens, cfg.rope_theta)
+
+    def _patchify(self, x: mx.array) -> mx.array:
+        """Convert image latents to patch sequence.
+
+        Matches diffusers: channels-last within each patch.
+        x: (B, C, H, W) → (B, H//p * W//p, p*p*C)
+
+        diffusers logic:
+            image.view(C, 1, 1, h, pH, w, pW)
+            image.permute(1, 3, 5, 2, 4, 6, 0)  # (1, h, w, 1, pH, pW, C)
+            reshape → (h*w, pH*pW*C)
+        """
+        B, C, H, W = x.shape
+        p = self.cfg.patch_size
+        pH, pW = H // p, W // p
+        # (B, C, pH, p, pW, p)
+        x = x.reshape(B, C, pH, p, pW, p)
+        # → (B, pH, pW, p, p, C)  — channels LAST per patch
+        x = x.transpose(0, 2, 4, 3, 5, 1)
+        # → (B, pH*pW, p*p*C)
+        x = x.reshape(B, pH * pW, p * p * C)
+        return x
+
+    def _unpatchify(self, x: mx.array, h: int, w: int) -> mx.array:
+        """Convert patch sequence back to image latents.
+
+        Matches diffusers: channels-last within each patch.
+        x: (B, pH*pW, p*p*C) → (B, C, H, W)
+
+        diffusers logic:
+            x.view(1, h, w, 1, pH, pW, C)
+            x.permute(6, 0, 3, 1, 4, 2, 5)  # (C, 1, 1, h, pH, w, pW)
+            reshape → (C, H, W)
+        """
+        B = x.shape[0]
+        p = self.cfg.patch_size
+        C = self.cfg.in_channels
+        pH, pW = h // p, w // p
+        # (B, pH, pW, p, p, C)
+        x = x.reshape(B, pH, pW, p, p, C)
+        # → (B, C, pH, p, pW, p)
+        x = x.transpose(0, 5, 1, 3, 2, 4)
+        # → (B, C, H, W)
+        x = x.reshape(B, C, h, w)
+        return x
+
+    def __call__(
+        self,
+        x: mx.array,
+        t: mx.array,
+        cap_feats: mx.array,
+        cap_mask: mx.array | None = None,
+    ) -> mx.array:
+        """Forward pass — matches diffusers ZImageTransformer2DModel.forward().
+
+        Correct execution order (from diffusers source):
+            1. t_embed
+            2. x_embed → noise_refiner  (image tokens with RoPE)
+            3. cap_embed → context_refiner  (text tokens with RoPE)
+            4. build unified [img, cap] sequence  (IMAGE FIRST in basic mode)
+            5. main layers (30 blocks with AdaLN + RoPE)
+            6. final_layer on FULL unified sequence
+            7. extract image tokens → unpatchify
+
+        Args:
+            x: (B, C, H, W) noisy latents
+            t: (B,) timesteps (1-sigma, scaled by pipeline)
+            cap_feats: (B, L_text, cap_feat_dim) text encoder hidden states
+            cap_mask: (B, L_text) boolean mask for padding
+
+        Returns:
+            noise_pred: (B, C, H, W) predicted noise
+        """
+        B, C, H, W = x.shape
+        cfg = self.cfg
+        p = cfg.patch_size
+        pH, pW = H // p, W // p
+
+        # 1. Timestep embedding → adaln_input
+        adaln_input = self.t_embedder(t * cfg.t_scale)  # (B, 256)
+
+        # 2. Patchify + embed image latents
+        img = self._patchify(x)  # (B, pH*pW, patch_dim=64)
+        img = self.all_x_embedder["2-1"](img)  # (B, pH*pW, dim=3840)
+
+        L_cap_orig = cap_feats.shape[1]
+        L_img = img.shape[1]
+
+        # Pad caption to SEQ_MULTI_OF=32 (matching diffusers _pad_with_ids)
+        SEQ_MULTI_OF = 32
+        pad_len = (-L_cap_orig) % SEQ_MULTI_OF
+        L_cap = L_cap_orig + pad_len
+
+        # Build position IDs matching diffusers (cap: t=1..L_cap_orig, img: t=L_cap_orig+1)
+        # NOTE: position IDs use original cap length (not padded), padding tokens get (0,0,0) IDs
+        img_pos_ids, cap_pos_ids = build_position_ids(L_cap_orig, pH, pW)
+
+        # Look up RoPE frequencies from precomputed tables
+        img_freqs = self._rope(img_pos_ids)  # (L_img, rope_half_dim)
+        cap_freqs_orig = self._rope(cap_pos_ids)  # (L_cap_orig, rope_half_dim)
+
+        # Pad cap RoPE freqs with zeros for padding positions (same as diffusers)
+        if pad_len > 0:
+            cap_freqs = mx.concatenate([
+                cap_freqs_orig,
+                mx.zeros((pad_len, cap_freqs_orig.shape[-1]))
+            ], axis=0)
+        else:
+            cap_freqs = cap_freqs_orig
+
+        # noise_refiner on image tokens (with AdaLN, with RoPE)
+        for block in self.noise_refiner:
+            img = block(img, adaln_input, img_freqs)
+
+        # 3. Caption embedding (cap_embedder is RMSNorm then Linear)
+        cap = self.cap_embedder[0](cap_feats)  # RMSNorm
+        cap = self.cap_embedder[1](cap)        # Linear → (B, L_cap_orig, dim=3840)
+
+        # Pad caption with cap_pad_token (matching diffusers _pad_with_ids).
+        # In diffusers, ALL tokens (real + pad) attend to each other fully —
+        # cap_pad_token is a learned vector, not masked out. The diffusers
+        # "attn_mask" is only for batch-level padding (all-True for BS=1).
+        if pad_len > 0:
+            pad_tok = mx.broadcast_to(self.cap_pad_token, (B, pad_len, cfg.dim))
+            cap = mx.concatenate([cap, pad_tok], axis=1)  # (B, L_cap, dim)
+
+        # context_refiner on text tokens (no AdaLN, WITH RoPE, no mask needed)
+        for block in self.context_refiner:
+            cap = block(cap, cap_freqs)
+
+        # 4. Build unified sequence [img, cap] — IMAGE FIRST (diffusers basic mode)
+        unified = mx.concatenate([img, cap], axis=1)  # (B, L_img + L_cap, dim)
+        unified_freqs = mx.concatenate([img_freqs, cap_freqs], axis=0)
+
+        # 5. Main DiT layers (30 blocks, with AdaLN conditioning + RoPE)
+        for block in self.layers:
+            unified = block(unified, adaln_input, unified_freqs)
+
+        # 6. Final layer on FULL unified sequence (as diffusers does)
+        unified = self.all_final_layer["2-1"](unified, adaln_input)
+
+        # 7. Extract image tokens (first L_img tokens) and unpatchify
+        img_out = unified[:, :L_img, :]  # (B, L_img, patch_dim=64)
+        out = self._unpatchify(img_out, H, W)
+        return out