moss-image-mlx 0.1.0__py3-none-any.whl

moss_image_mlx/__init__.py

@@ -0,0 +1,6 @@
+"""MOSS-Image on Apple Silicon (MLX)."""
+
+from .pipeline import MOSSImagePipeline, Request, save
+
+__version__ = "0.1.0"
+__all__ = ["MOSSImagePipeline", "Request", "save", "__version__"]

moss_image_mlx/cache.py

@@ -0,0 +1,57 @@
+"""On-disk cache of prompt embeddings.
+
+Re-rendering a prompt (new seed, steps, size or guidance) then skips the 9B text encoder entirely: no ~16 GB (bf16)
+load and no encode. Entries are keyed by the encoder's identity (config, weight files, quantization, the
+implementation's numerics version), the token limit and the exact prompt text.
+"""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import os
+from pathlib import Path
+
+import mlx.core as mx
+
+# Bump when the encoder numerics change, so stale embeddings are never reused.
+ENCODER_NUMERICS_VERSION = "transformers-faithful-1"
+
+
+def default_cache_dir() -> Path:
+    return Path(os.environ.get("MOSS_MLX_CACHE", Path.home() / ".cache" / "moss-image-mlx")) / "prompt-embeds"
+
+
+def encoder_id(text_encoder_dir: str | Path) -> str:
+    d = Path(text_encoder_dir)
+    h = hashlib.sha256(ENCODER_NUMERICS_VERSION.encode())
+    h.update((d / "config.json").read_bytes())
+    for f in sorted(d.glob("*.safetensors")):
+        h.update(f"{f.name}:{f.stat().st_size}".encode())
+    return h.hexdigest()[:16]
+
+
+class PromptCache:
+    def __init__(self, text_encoder_dir: str | Path, root: str | Path | None = None, tag: str = "bf16"):
+        """tag distinguishes encoder variants built from the same files (e.g. quantized at load)."""
+        self.dir = Path(root or default_cache_dir()) / f"{encoder_id(text_encoder_dir)}-{tag}"
+
+    def _path(self, prompt: str, max_tokens: int) -> Path:
+        key = hashlib.sha256(json.dumps([prompt, max_tokens], ensure_ascii=False).encode()).hexdigest()
+        return self.dir / f"{key}.safetensors"
+
+    def get(self, prompt: str, max_tokens: int) -> mx.array | None:
+        path = self._path(prompt, max_tokens)
+        if not path.exists():
+            return None
+        try:
+            return mx.load(str(path))["hidden"]
+        except Exception:
+            return None
+
+    def put(self, prompt: str, max_tokens: int, hidden: mx.array):
+        self.dir.mkdir(parents=True, exist_ok=True)
+        path = self._path(prompt, max_tokens)
+        tmp = path.with_suffix(".tmp.safetensors")
+        mx.save_safetensors(str(tmp), {"hidden": hidden})
+        tmp.replace(path)

moss_image_mlx/cli.py

@@ -0,0 +1,156 @@
+"""moss-mlx: MOSS-Image text-to-image on Apple Silicon (MLX).
+
+Examples:
+  moss-mlx "A red ceramic teapot on a wooden table" --model ~/models/MOSS-Image-SFT -o teapot.png
+  moss-mlx "A poster titled HELLO" --model ... --seeds 1,2,3,4 -o poster.png      # one load, four images
+  moss-mlx --parameters assets/images/005/parameters.json --model ... --transformer .../MOSS-Image-Base/transformer
+  moss-mlx --batch jobs.jsonl --model ...      # one JSON object per line (Request fields, or {"parameters": path})
+"""
+
+import argparse
+import json
+import sys
+import time
+from dataclasses import replace
+from pathlib import Path
+
+from .pipeline import MOSSImagePipeline, Request, save
+
+TIERS = {"8bit": "VisualInference/MOSS-Image-SFT-MLX-8bit", "compact": "VisualInference/MOSS-Image-SFT-MLX-compact"}
+
+
+def default_model() -> str:
+    """8bit on Macs with >= 24 GB of memory (9.2 GB peak at 1024^2), compact below (7.7 GB peak)."""
+    import subprocess
+
+    try:
+        mem = int(subprocess.run(["sysctl", "-n", "hw.memsize"], capture_output=True, text=True).stdout)
+    except Exception:
+        mem = 0
+    return TIERS["8bit"] if mem >= 24 * 1024**3 else TIERS["compact"]
+
+
+def resolve_model(model: str | None) -> str:
+    """A local directory, a tier name (8bit, compact), or a Hugging Face repo id (downloaded and cached)."""
+    model = model or default_model()
+    model = TIERS.get(model, model)
+    if Path(model).expanduser().is_dir():
+        return str(Path(model).expanduser())
+    import subprocess
+
+    print(f"fetching {model} (cached after the first run)", file=sys.stderr)
+    # Download in a child process: the Hub client keeps several GB of transfer buffers that this process would
+    # otherwise still hold when the text encoder loads (measured +2.7 GB on the first run's peak).
+    code = "import sys; from huggingface_hub import snapshot_download; print(snapshot_download(sys.argv[1]))"
+    done = subprocess.run([sys.executable, "-c", code, model], stdout=subprocess.PIPE, text=True)
+    if done.returncode:
+        sys.exit(f"could not download {model}")
+    return done.stdout.strip().splitlines()[-1]
+
+
+def _layers(text: str) -> list[int]:
+    return [int(x) for x in str(text).replace(",", " ").split()] if str(text).strip() else []
+
+
+def quant_spec(text: str | None):
+    """'8' -> 8-bit g64 affine; '4g32' -> 4-bit g32; 'nvfp4' / 'mxfp8'; or a JSON spec."""
+    if not text or text in ("bf16", "none"):
+        return None
+    if text.startswith("{"):
+        return json.loads(text)
+    if text == "nvfp4":
+        return {"bits": 4, "group_size": 16, "mode": "nvfp4"}
+    if text == "mxfp8":
+        return {"bits": 8, "group_size": 32, "mode": "mxfp8"}
+    bits, _, group = text.partition("g")
+    return {"bits": int(bits), "group_size": int(group or 64)}
+
+
+def _from_parameters(path) -> dict:
+    params = json.load(open(path))["parameters"]
+    return {"prompt": params["prompt"], "negative_prompt": params.get("negative_prompt", ""),
+            "width": params["width"], "height": params["height"], "steps": params["steps"],
+            "guidance": float(params["cfg"]), "shift": float(params["shift"]), "seed": params["seed"],
+            "uncond_skip_layers": _layers(params.get("uncond_skip_layers", "22"))}
+
+
+def main(argv=None):
+    p = argparse.ArgumentParser(prog="moss-mlx", description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("prompt", nargs="?", help="Prompt text or layout JSON")
+    p.add_argument("--model", help="Tier (8bit, compact), Hugging Face repo id or local directory. Default: 8bit on "
+                                   ">= 24 GB Macs, compact below")
+    p.add_argument("--transformer", help="Override the transformer directory (e.g. Base with the shared encoder)")
+    p.add_argument("--parameters", help="Upstream parameters.json (official example format); CLI flags override it")
+    p.add_argument("--batch", help="JSONL file of jobs; rendered with each model component loaded once")
+    p.add_argument("--negative-prompt")
+    p.add_argument("--width", type=int)
+    p.add_argument("--height", type=int)
+    p.add_argument("--steps", type=int)
+    p.add_argument("--guidance", type=float)
+    p.add_argument("--shift", type=float)
+    p.add_argument("--seed", type=int)
+    p.add_argument("--seeds", help="Comma-separated seeds: one image per seed, loaded once")
+    p.add_argument("--uncond-skip-layers", help="e.g. '22'; empty string disables")
+    p.add_argument("-o", "--output", default="moss.png", help="Output PNG (several images get -<n> suffixes)")
+    p.add_argument("--dit-dtype", choices=("bfloat16", "float32"), help="Transformer precision (default: as stored)")
+    p.add_argument("--dit-quant", help="Quantize the transformer at load: 8, 6, 4g32, nvfp4, mxfp8 or JSON")
+    p.add_argument("--te-quant", help="Quantize the text encoder at load (same forms)")
+    p.add_argument("--precise", action="store_true", help="fp32 residual stream and elementwise ops (bf16 matmuls)")
+    p.add_argument("--no-cache", action="store_true", help="Do not read or write the prompt-embedding cache")
+    p.add_argument("--keep-resident", action="store_true", help="Do not free components between stages")
+    args = p.parse_args(argv)
+    args.model = resolve_model(args.model)
+
+    base = _from_parameters(args.parameters) if args.parameters else {}
+    for name in ("prompt", "negative_prompt", "width", "height", "steps", "guidance", "shift", "seed"):
+        if getattr(args, name) is not None:
+            base[name] = getattr(args, name)
+    if args.uncond_skip_layers is not None:
+        base["uncond_skip_layers"] = _layers(args.uncond_skip_layers)
+
+    jobs = []
+    if args.batch:
+        for line in Path(args.batch).read_text().splitlines():
+            if line.strip():
+                job = json.loads(line)
+                fields = _from_parameters(job.pop("parameters")) if "parameters" in job else {}
+                out = job.pop("output", None)
+                jobs.append((Request(**{**base, **fields, **job}), out))
+    else:
+        if "prompt" not in base:
+            p.error("a prompt, --parameters or --batch is required")
+        req = Request(**base)
+        seeds = [int(s) for s in args.seeds.split(",")] if args.seeds else [req.seed]
+        jobs = [(replace(req, seed=s), None) for s in seeds]
+    from .layout import check
+
+    for r, _ in jobs:
+        if r.prompt.lstrip().startswith("{"):
+            for problem in check(r.prompt, r.width, r.height):
+                print(f"layout warning: {problem}", file=sys.stderr)
+    out = Path(args.output)
+    jobs = [(r, o or (str(out) if len(jobs) == 1 else str(out.with_name(f"{out.stem}-{i}{out.suffix}"))))
+            for i, (r, o) in enumerate(jobs)]
+
+    import mlx.core as mx
+
+    dtype = getattr(mx, args.dit_dtype) if args.dit_dtype else None
+    pipe = MOSSImagePipeline(args.model, args.transformer, low_memory=not args.keep_resident, dit_dtype=dtype,
+                             dit_quant=quant_spec(args.dit_quant), te_quant=quant_spec(args.te_quant),
+                             cache_dir=False if args.no_cache else None, precise=args.precise)
+    start = time.perf_counter()
+
+    def progress(i, n):
+        print(f"\rstep {i + 1}/{n}  {time.perf_counter() - start:7.1f} s", end="", file=sys.stderr, flush=True)
+
+    paths = dict((id(r), o) for r, o in jobs)
+    for req, image in pipe.generate_many([r for r, _ in jobs], on_step=progress):
+        print(file=sys.stderr)
+        save(image, req, paths[id(req)], {"timings": dict(pipe.timings), "precise": args.precise,
+                                          "dit_quant": args.dit_quant or "bf16", "te_quant": args.te_quant or "bf16"})
+        print(paths[id(req)])
+    print(json.dumps(pipe.timings), file=sys.stderr)
+
+
+if __name__ == "__main__":
+    main()

moss_image_mlx/dit.py

@@ -0,0 +1,293 @@
+"""MOSS-Image diffusion transformer in MLX.
+
+Mirrors MOSSImageTransformer in the official pipeline_moss_image.py (OpenMOSS, Apache-2.0) operation by operation:
+bf16 weights and residual stream, RMSNorm reduced in fp32 then cast back before the weight multiply, RoPE phases in
+fp32 (computed on the CPU, cast to bf16), the timestep sinusoid in float64, and a single shared adaLN modulation
+for every block. Parameter names equal the reference state-dict keys (except time_embedding, see sanitize).
+"""
+
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import mlx.core as mx
+import mlx.nn as nn
+import numpy as np
+
+
+class RMSNorm(nn.Module):
+    """torch: (x.float() * rsqrt(mean(x.float()**2) + 1e-6)).to(x.dtype) * weight"""
+
+    def __init__(self, dim: int):
+        super().__init__()
+        self.weight = mx.ones((dim,))
+
+    def __call__(self, x):
+        return mx.fast.rms_norm(x.astype(mx.float32), None, 1e-6).astype(x.dtype) * self.weight
+
+
+@mx.compile
+def situ(hidden, gate):
+    return (4.0 * mx.tanh(gate / 4.0) * mx.sigmoid(gate)) * (25.0 * mx.tanh(hidden / 25.0))
+
+
+@mx.compile
+def _norm_modulate(x, weight, scale, shift):
+    """RMSNorm (fp32 reduction, cast, weight) then adaLN modulation, each op rounding like the reference."""
+    return (mx.fast.rms_norm(x.astype(mx.float32), None, 1e-6).astype(x.dtype) * weight) * (1 + scale) + shift
+
+
+@mx.compile
+def _gated_add(x, y, gate):
+    return x + y * gate
+
+
+def silu_like_torch(x):
+    """torch's fused silu on bf16 computes in fp32 and rounds once."""
+    x32 = x.astype(mx.float32)
+    return (x32 * mx.sigmoid(x32)).astype(x.dtype)
+
+
+@mx.compile
+def rotate(x, sin, cos):
+    half = x.shape[-1] // 2
+    return x * cos + mx.concatenate([-x[..., half:], x[..., :half]], axis=-1) * sin
+
+
+def attention(q, k, v, rope, heads: int):
+    """q, k, v: (B, N, dim); q and k are already per-head RMS-normalised. rope: (sin, cos), each (1, N, 1, hd)."""
+    b, n, dim = q.shape
+    hd = dim // heads
+    sin, cos = rope
+    q = rotate(q.reshape(b, n, heads, hd), sin, cos).transpose(0, 2, 1, 3)
+    k = rotate(k.reshape(b, n, heads, hd), sin, cos).transpose(0, 2, 1, 3)
+    v = v.reshape(b, n, heads, hd).transpose(0, 2, 1, 3)
+    out = mx.fast.scaled_dot_product_attention(q, k, v, scale=hd**-0.5)
+    return out.transpose(0, 2, 1, 3).reshape(b, n, dim)
+
+
+def _head_norm(norm: RMSNorm, x, heads: int):
+    b, n, dim = x.shape
+    return norm(x.reshape(b, n, heads, dim // heads)).reshape(b, n, dim)
+
+
+# "Precise" mode: matmuls and attention keep bf16 inputs (speed), everything elementwise runs in fp32 and the
+# residual stream stays fp32, so bf16 rounding enters only at matmul/attention inputs.
+F32, BF16 = mx.float32, mx.bfloat16
+
+
+def _norm32(norm: RMSNorm, x):
+    return mx.fast.rms_norm(x.astype(F32), norm.weight.astype(F32), 1e-6)
+
+
+def _head_norm32(norm: RMSNorm, x, heads: int):
+    b, n, dim = x.shape
+    return _norm32(norm, x.reshape(b, n, heads, dim // heads)).reshape(b, n, dim)
+
+
+def _attention32(q, k, v, rope, heads: int):
+    """q, k fp32 (normalised), v bf16; RoPE in fp32, attention in bf16."""
+    b, n, dim = q.shape
+    hd = dim // heads
+    sin, cos = rope
+    q = rotate(q.reshape(b, n, heads, hd), sin, cos).astype(BF16).transpose(0, 2, 1, 3)
+    k = rotate(k.reshape(b, n, heads, hd), sin, cos).astype(BF16).transpose(0, 2, 1, 3)
+    v = v.reshape(b, n, heads, hd).transpose(0, 2, 1, 3)
+    out = mx.fast.scaled_dot_product_attention(q, k, v, scale=hd**-0.5)
+    return out.transpose(0, 2, 1, 3).reshape(b, n, dim)
+
+
+class TextBlock(nn.Module):
+    def __init__(self, dim: int, ffn_dim: int, heads: int):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = RMSNorm(dim), RMSNorm(dim)
+        self.q_norm, self.k_norm = RMSNorm(dim // heads), RMSNorm(dim // heads)
+        self.qkv = nn.Linear(dim, 3 * dim)
+        self.proj = nn.Linear(dim, dim)
+        self.ffn_gate, self.ffn_up = nn.Linear(dim, ffn_dim), nn.Linear(dim, ffn_dim)
+        self.ffn_down = nn.Linear(ffn_dim, dim)
+
+    def __call__(self, x, rope, precise: bool = False):
+        if precise:
+            q, k, v = mx.split(self.qkv(_norm32(self.norm1, x).astype(BF16)), 3, axis=-1)
+            q, k = _head_norm32(self.q_norm, q, self.heads), _head_norm32(self.k_norm, k, self.heads)
+            x = x + self.proj(_attention32(q, k, v, rope, self.heads)).astype(F32)
+            y = _norm32(self.norm2, x).astype(BF16)
+            h = situ(self.ffn_up(y).astype(F32), self.ffn_gate(y).astype(F32))
+            return x + self.ffn_down(h.astype(BF16)).astype(F32)
+        q, k, v = mx.split(self.qkv(self.norm1(x)), 3, axis=-1)
+        q, k = _head_norm(self.q_norm, q, self.heads), _head_norm(self.k_norm, k, self.heads)
+        x = x + self.proj(attention(q, k, v, rope, self.heads))
+        y = self.norm2(x)
+        return x + self.ffn_down(situ(self.ffn_up(y), self.ffn_gate(y)))
+
+
+class TextProjection(nn.Module):
+    def __init__(self, text_dim: int, dim: int, ffn_dim: int, heads: int):
+        super().__init__()
+        self.proj = nn.Linear(text_dim, dim)
+        self.blocks = [TextBlock(dim, ffn_dim, heads) for _ in range(2)]
+
+    def __call__(self, x, rope, precise: bool = False):
+        x = self.proj(x.astype(BF16 if precise else self.proj.weight.dtype))
+        if precise:
+            x = x.astype(F32)
+        for block in self.blocks:
+            x = block(x, rope, precise)
+        return x
+
+
+class Block(nn.Module):
+    def __init__(self, dim: int, ffn_dim: int, heads: int):
+        super().__init__()
+        self.heads = heads
+        self.norm1, self.norm2 = RMSNorm(dim), RMSNorm(dim)
+        self.q_norm, self.k_norm = RMSNorm(dim // heads), RMSNorm(dim // heads)
+        self.proj = {name: nn.Linear(dim, dim, bias=False) for name in ("q", "k", "v", "o")}
+        self.proj["ffn_up_gate"] = nn.Linear(dim, 2 * ffn_dim, bias=False)
+        self.proj["ffn_down"] = nn.Linear(ffn_dim, dim, bias=False)
+
+    def __call__(self, x, rope, mod, precise: bool = False):
+        shift_a, scale_a, gate_a, shift_f, scale_f, gate_f = (mod[:, i : i + 1] for i in range(6))
+        p = self.proj
+        if precise:  # x and mod are fp32 here
+            y = (_norm32(self.norm1, x) * (1 + scale_a) + shift_a).astype(BF16)
+            q = _head_norm32(self.q_norm, p["q"](y), self.heads)
+            k = _head_norm32(self.k_norm, p["k"](y), self.heads)
+            x = x + p["o"](_attention32(q, k, p["v"](y), rope, self.heads)).astype(F32) * gate_a
+            y = (_norm32(self.norm2, x) * (1 + scale_f) + shift_f).astype(BF16)
+            hidden, gate = mx.split(p["ffn_up_gate"](y).astype(F32), 2, axis=-1)
+            return x + p["ffn_down"](situ(hidden, gate).astype(BF16)).astype(F32) * gate_f
+        y = _norm_modulate(x, self.norm1.weight, scale_a, shift_a)
+        q = _head_norm(self.q_norm, p["q"](y), self.heads)
+        k = _head_norm(self.k_norm, p["k"](y), self.heads)
+        x = _gated_add(x, p["o"](attention(q, k, p["v"](y), rope, self.heads)), gate_a)
+        y = _norm_modulate(x, self.norm2.weight, scale_f, shift_f)
+        hidden, gate = mx.split(p["ffn_up_gate"](y), 2, axis=-1)
+        return _gated_add(x, p["ffn_down"](situ(hidden, gate)), gate_f)
+
+
+def rope_table(height: int, width: int, text_length: int, head_dim: int, base: float, dtype) -> tuple:
+    """(sin, cos), each (1, N, 1, head_dim), for image tokens then text tokens; coordinates (r, t, h, w)."""
+    bands_n = head_dim // 8
+    exponent = np.arange(bands_n, dtype=np.float32) / np.float32(bands_n)
+    bands = (np.float32(1) / (np.float32(base) ** exponent)).astype(np.float32)
+    h, w = np.meshgrid(np.arange(height), np.arange(width), indexing="ij")
+    image = np.stack([np.zeros_like(h), np.full_like(h, text_length), h, w], -1).reshape(-1, 4).astype(np.float32)
+    image[:, 2] -= np.float32((height - 1) / 2)
+    image[:, 3] -= np.float32((width - 1) / 2)
+    text = np.zeros((text_length, 4), dtype=np.float32)
+    text[:, 1] = np.arange(text_length, dtype=np.float32)
+    phase = np.concatenate([image, text], 0)[..., None] * bands  # (N, 4, bands)
+    n = phase.shape[0]
+    sin = np.sin(phase).reshape(n, -1)  # (N, head_dim / 2): axes r, t, h, w
+    cos = np.cos(phase).reshape(n, -1)
+    sin = mx.array(np.concatenate([sin, sin], -1)).astype(dtype)[None, :, None, :]
+    cos = mx.array(np.concatenate([cos, cos], -1)).astype(dtype)[None, :, None, :]
+    return sin, cos
+
+
+def timestep_embedding(timestep: float, dim: int) -> np.ndarray:
+    """float64 sinusoid as in the reference (cos first), returned as float32."""
+    t = np.float64(np.float32(timestep))
+    half = dim // 2
+    freq = 10000.0 ** (-np.arange(half, dtype=np.float64) / half)
+    phase = t * freq
+    return np.concatenate([np.cos(phase), np.sin(phase)]).astype(np.float32)[None]
+
+
+class MOSSImageTransformer(nn.Module):
+    def __init__(self, dim=3072, ffn_dim=8192, num_heads=24, num_layers=42, text_dim=4096, freq_dim=256,
+                 latent_channels=128, rope_base=10000.0):
+        super().__init__()
+        self.config = dict(dim=dim, ffn_dim=ffn_dim, num_heads=num_heads, num_layers=num_layers, text_dim=text_dim,
+                           freq_dim=freq_dim, latent_channels=latent_channels, rope_base=rope_base)
+        self.video_embedder = nn.Linear(latent_channels, dim)
+        self.text_embedder = TextProjection(text_dim, dim, ffn_dim, num_heads)
+        self.time_embedding = [nn.Linear(freq_dim, dim), nn.Linear(dim, dim)]
+        self.shared_block_modulation = nn.Linear(dim, 6 * dim)
+        self.layers = [Block(dim, ffn_dim, num_heads) for _ in range(num_layers)]
+        self.final_norm_video = RMSNorm(dim)
+        self.final_linear_video = nn.Linear(dim, latent_channels, bias=False)
+        self.video_head_modulation = nn.Linear(dim, 2 * dim)
+        self._rope_cache = {}
+
+    def modulation(self, timestep: float):
+        """Shared block modulation (1, 6, dim) and head modulation (1, 2, dim), both in the weight dtype."""
+        dtype = self.time_embedding[0].weight.dtype
+        emb = mx.array(timestep_embedding(timestep, self.config["freq_dim"])).astype(dtype)
+        time = self.time_embedding[1](silu_like_torch(self.time_embedding[0](emb)))
+        dim = self.config["dim"]
+        mod = self.shared_block_modulation(time).reshape(1, 6, dim)
+        head = self.video_head_modulation(time).reshape(1, 2, dim)
+        return mod, head
+
+    def rope(self, height: int, width: int, text_length: int, dtype):
+        key = (height, width, text_length, dtype)
+        if key not in self._rope_cache:
+            self._rope_cache = {key: rope_table(height, width, text_length, self.config["dim"] // self.config["num_heads"],
+                                                self.config["rope_base"], dtype)}
+        return self._rope_cache[key]
+
+    def __call__(self, latents, timestep: float, context, skip_layers=()):
+        """latents (1, 128, H, W) fp32; context (1, L, text_dim) with no padding. Returns velocity (1, 128, H, W)."""
+        b, c, height, width = latents.shape
+        if b != 1 or context.shape[0] != 1:
+            raise ValueError("The reference evaluates each sample separately; pass batch 1.")
+        dtype = self.video_embedder.weight.dtype
+        precise = bool(getattr(self, "precise", False)) and dtype != F32
+        mod, head = self.modulation(timestep)
+        video = self.video_embedder(latents.reshape(1, c, height * width).transpose(0, 2, 1).astype(dtype))
+        act = F32 if precise else dtype
+        if precise:
+            mod, head, video = mod.astype(F32), head.astype(F32), video.astype(F32)
+        sin, cos = self.rope(height, width, context.shape[1], act)
+        n_img = height * width
+        text = self.text_embedder(context, (sin[:, n_img:], cos[:, n_img:]), precise)
+        x = mx.concatenate([video, text], axis=1)
+        skip = set(skip_layers or ())
+        for index, layer in enumerate(self.layers):
+            if index not in skip:
+                x = layer(x, (sin, cos), mod, precise)
+        if precise:
+            x = _norm32(self.final_norm_video, x[:, :n_img]) * (1 + head[:, 1:2]) + head[:, 0:1]
+            x = self.final_linear_video(x.astype(BF16))
+        else:
+            x = self.final_norm_video(x[:, :n_img])
+            x = x * (1 + head[:, 1:2]) + head[:, 0:1]
+            x = self.final_linear_video(x)
+        return x.transpose(0, 2, 1).reshape(1, c, height, width)
+
+    @staticmethod
+    def sanitize(weights: dict) -> dict:
+        out = {}
+        for key, value in weights.items():
+            if key == "rope.bands":  # derived constant, recomputed in fp32
+                continue
+            out[key.replace("time_embedding.2.", "time_embedding.1.")] = value
+        return out
+
+
+def load_transformer(path: str | Path, dtype=None) -> MOSSImageTransformer:
+    """Load a `transformer/` directory: the original diffusers layout, or an exported MLX one (possibly quantized,
+    marked by "moss_image_mlx" in config.json)."""
+    path = Path(path)
+    config = json.loads((path / "config.json").read_text())
+    keys = ("dim", "ffn_dim", "num_heads", "num_layers", "text_dim", "freq_dim", "latent_channels", "rope_base")
+    model = MOSSImageTransformer(**{k: config[k] for k in keys if k in config})
+    exported = config.get("moss_image_mlx")
+    if exported and exported.get("quantization"):
+        from .quantize import quantize_dit
+
+        quantize_dit(model, exported["quantization"])  # build QuantizedLinear modules before loading their weights
+    weights = {}
+    for file in sorted(path.glob("*.safetensors")):
+        weights.update(mx.load(str(file)))
+    if not exported:
+        weights = model.sanitize(weights)
+    if dtype is not None:
+        weights = {k: v.astype(dtype) if mx.issubdtype(v.dtype, mx.floating) else v for k, v in weights.items()}
+    model.load_weights(list(weights.items()), strict=True)
+    return model

moss_image_mlx/export.py

@@ -0,0 +1,135 @@
+"""Export a MOSS-Image checkpoint to the MLX release layout, optionally quantized.
+
+  moss-mlx-export --model ~/hf/MOSS-Image-SFT --dit-quant 8 --te-quant 4 --out ~/out/MOSS-Image-SFT-MLX-compact
+
+Layout (each component config.json carries {"moss_image_mlx": {"format": 1, "quantization": spec | null}}):
+  config.json            top-level summary (also what the Hub counts as a download for library "mlx")
+  model_index.json       diffusers-style index of the components
+  transformer/           MLX parameter names, sharded safetensors
+  text_encoder/          Qwen3.5 language backbone only (vision tower dropped), norms stored as fp32 (1 + w)
+  tokenizer/, scheduler/ copied
+  vae/                   decoder only, MLX conv layout, fp32
+Quantization runs on the CPU stream and is evaluated module by module, so the bf16 source never sits in memory whole.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import shutil
+from pathlib import Path
+
+import mlx.core as mx
+from mlx.utils import tree_flatten
+
+from .cli import quant_spec
+from .dit import load_transformer
+from .quantize import describe, quantize_dit, quantize_text_encoder
+from .text_encoder import TextEncoder
+from .vae import load_vae_decoder
+
+FORMAT = 1
+SHARD_BYTES = 4 << 30
+
+
+def _eval_by_module(model):
+    for _, module in model.named_modules():
+        params = [v for _, v in tree_flatten(module.parameters())]
+        if params:
+            mx.eval(params)
+
+
+def _save_sharded(params: dict, out: Path, prefix: str = "model"):
+    out.mkdir(parents=True, exist_ok=True)
+    shards, cur, size = [], {}, 0
+    for k, v in params.items():
+        if cur and size + v.nbytes > SHARD_BYTES:
+            shards.append(cur)
+            cur, size = {}, 0
+        cur[k] = v
+        size += v.nbytes
+    shards.append(cur)
+    index = {}
+    for i, shard in enumerate(shards):
+        name = f"{prefix}-{i + 1:05d}-of-{len(shards):05d}.safetensors" if len(shards) > 1 else f"{prefix}.safetensors"
+        mx.save_safetensors(str(out / name), shard, metadata={"format": "mlx"})
+        index.update({k: name for k in shard})
+    if len(shards) > 1:
+        total = sum(v.nbytes for v in params.values())
+        (out / f"{prefix}.safetensors.index.json").write_text(json.dumps({"metadata": {"total_size": total},
+                                                                          "weight_map": index}, indent=1))
+
+
+def _component_config(src: Path, spec, drop=()):
+    config = json.loads(src.read_text())
+    for key in drop:
+        config.pop(key, None)
+    config["moss_image_mlx"] = {"format": FORMAT, "quantization": spec}
+    return config
+
+
+def export(model_dir, out, transformer_dir=None, dit_spec=None, te_spec=None, source_id=None, tier=None):
+    model_dir, out = Path(model_dir), Path(out)
+    transformer_dir = Path(transformer_dir) if transformer_dir else model_dir / "transformer"
+    out.mkdir(parents=True, exist_ok=True)
+    mx.set_default_device(mx.cpu)
+
+    dit = load_transformer(transformer_dir)
+    quantize_dit(dit, dit_spec)
+    _eval_by_module(dit)
+    _save_sharded(dict(tree_flatten(dit.parameters())), out / "transformer")
+    cfg = _component_config(transformer_dir / "config.json", dit_spec)
+    (out / "transformer" / "config.json").write_text(json.dumps(cfg, indent=2) + "\n")
+    del dit
+    mx.clear_cache()
+
+    te = TextEncoder(model_dir / "text_encoder", model_dir / "tokenizer")
+    quantize_text_encoder(te.model, te_spec)
+    _eval_by_module(te.model)
+    _save_sharded(dict(tree_flatten(te.model.parameters())), out / "text_encoder")
+    cfg = _component_config(model_dir / "text_encoder" / "config.json", te_spec, drop=("vision_config",))
+    cfg["architectures"] = ["Qwen3_5TextModel"]
+    (out / "text_encoder" / "config.json").write_text(json.dumps(cfg, indent=2) + "\n")
+    del te
+    mx.clear_cache()
+
+    vae = load_vae_decoder(model_dir / "vae")
+    _save_sharded(dict(tree_flatten(vae.parameters())), out / "vae", prefix="decoder")
+    cfg = _component_config(model_dir / "vae" / "config.json", None)
+    cfg["include_encoder"] = False
+    (out / "vae" / "config.json").write_text(json.dumps(cfg, indent=2) + "\n")
+
+    for name in ("tokenizer", "scheduler"):
+        (out / name).mkdir(exist_ok=True)
+        for f in (model_dir / name).iterdir():
+            if f.is_file() and f.suffix in (".json", ".jinja", ".txt", ".model"):
+                shutil.copy2(f, out / name / f.name)
+    index = json.loads((model_dir / "model_index.json").read_text())
+    index["moss_image_mlx"] = {"format": FORMAT}
+    (out / "model_index.json").write_text(json.dumps(index, indent=2) + "\n")
+    if (model_dir / "LICENSE").exists():
+        shutil.copy2(model_dir / "LICENSE", out / "LICENSE")
+    summary = {"model_type": "moss_image", "library_name": "moss-image-mlx", "format": FORMAT, "tier": tier,
+               "source_model": source_id, "transformer_quantization": describe(dit_spec),
+               "text_encoder_quantization": describe(te_spec),
+               "components": ["transformer", "text_encoder", "tokenizer", "vae", "scheduler"]}
+    (out / "config.json").write_text(json.dumps(summary, indent=2) + "\n")
+    return out
+
+
+def main(argv=None):
+    p = argparse.ArgumentParser(prog="moss-mlx-export", description=__doc__,
+                                formatter_class=argparse.RawDescriptionHelpFormatter)
+    p.add_argument("--model", required=True)
+    p.add_argument("--transformer")
+    p.add_argument("--dit-quant")
+    p.add_argument("--te-quant")
+    p.add_argument("--source-id")
+    p.add_argument("--tier")
+    p.add_argument("--out", required=True)
+    a = p.parse_args(argv)
+    print(export(a.model, a.out, a.transformer, quant_spec(a.dit_quant), quant_spec(a.te_quant), a.source_id, a.tier))
+
+
+if __name__ == "__main__":
+    main()