Release v0.4.0: selective MLX model downloads

3da88e3 verified 2 days ago

26.2 kB

from __future__ import annotations

import gc
import hashlib
import importlib.metadata
import json
import logging
import math
import os
import time
import uuid
from collections import OrderedDict
from pathlib import Path

import mlx.core as mx
import numpy as np

from .config import DiTConfig
from .layout import image_positions, rope_cache
from .ops import attention, gated_add, modulate, rms_norm, rotate, silu, swiglu
from .weights import Linear, WeightStore

log = logging.getLogger(__name__)


class DiTBlock:
    def __init__(
        self, store, prefix, config, *, bits=None, group_size=64, packed_cache=None
    ):
        self.config = config
        self.n1 = store.get(prefix + ".norm1.weight")
        self.n2 = store.get(prefix + ".norm2.weight")
        self.qn = store.get(prefix + ".attn.q_norm.weight")
        self.kn = store.get(prefix + ".attn.k_norm.weight")
        if packed_cache is not None:
            from .packed import PackedBF16, PackedLinear

            cache_path = Path(packed_cache) / f"{prefix}.safetensors"
            names = ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2")
            attrs = ("qkv", "out", "fc1", "fc2")
            if cache_path.exists():
                data, metadata = mx.load(str(cache_path), return_metadata=True)
                if metadata.get("format") != "h3-palette12-v1":
                    raise ValueError("Unsupported packed checkpoint cache")
                shapes = json.loads(metadata["shapes"])
                for attr, name in zip(attrs, names, strict=True):
                    packed = PackedBF16.from_tensors(data, name, shapes[name])
                    setattr(self, attr, PackedLinear(packed, data.get(name + ".bias")))
                mx.eval(data)
                return
        grouped = (config.num_attention_heads, config.attention_head_dim)
        self.qkv = Linear.load(
            store,
            prefix + ".attn.qkv_proj",
            bits=bits,
            group_size=group_size,
            grouped_heads=grouped,
        )
        self.out = Linear.load(
            store, prefix + ".attn.out_proj", bits=bits, group_size=group_size
        )
        self.fc1 = Linear.load(
            store, prefix + ".mlp.fc1", bits=bits, group_size=group_size
        )
        self.fc2 = Linear.load(
            store, prefix + ".mlp.fc2", bits=bits, group_size=group_size
        )
        if packed_cache is not None:
            data, shapes = {}, {}
            for attr, name in zip(attrs, names, strict=True):
                linear = getattr(self, attr)
                packed = PackedLinear.from_linear(linear)
                setattr(self, attr, packed)
                shapes[name] = packed.weight.shape
                data.update(packed.weight.tensors(name))
                if packed.bias is not None:
                    data[name + ".bias"] = packed.bias
            cache_path.parent.mkdir(parents=True, exist_ok=True)
            temporary = cache_path.with_name(
                f".{cache_path.stem}-{uuid.uuid4().hex}.safetensors"
            )
            try:
                mx.save_safetensors(
                    str(temporary),
                    data,
                    metadata={
                        "format": "h3-palette12-v1",
                        "shapes": json.dumps(shapes),
                    },
                )
                temporary.replace(cache_path)
            finally:
                temporary.unlink(missing_ok=True)

    @property
    def nbytes(self):
        return sum(
            x.nbytes
            for x in (
                self.n1,
                self.n2,
                self.qn,
                self.kn,
                self.qkv,
                self.out,
                self.fc1,
                self.fc2,
            )
        )

    @staticmethod
    def project(linear, x, state):
        return linear(x) if state is None else linear(x, state=state)

    def attn(self, x, rope, query_chunk, kernel="native", projection_weights=None):
        c = self.config
        states = (None, None) if projection_weights is None else projection_weights[:2]
        qkv = self.project(self.qkv, x, states[0]).reshape(
            1, x.shape[0], 3, c.num_attention_heads, c.attention_head_dim
        )
        q, k, v = (qkv[:, :, i] for i in range(3))
        exact_qk = False
        if kernel in ("metal-exact", "metal-qk-exact") and rope is not None:
            from .tuning import qk_norm_rope_exact, supports_qk

            if supports_qk(qkv[0]):
                q, k = qk_norm_rope_exact(
                    qkv[0], self.qn, self.kn, *rope, c.qk_norm_eps
                )
                exact_qk = True
        if kernel == "metal-qk" and rope is not None:
            from .kernels import qk_norm_rope

            q, k = qk_norm_rope(qkv[0], self.qn, self.kn, *rope, c.qk_norm_eps)
        elif not exact_qk:
            q = rms_norm(q, self.qn, c.qk_norm_eps)
            k = rms_norm(k, self.kn, c.qk_norm_eps)
            if rope is not None:
                q, k = rotate(q, *rope), rotate(k, *rope)
        a = attention(q, k, v, query_chunk=query_chunk).reshape(x.shape[0], -1)
        return self.project(self.out, a, states[1])

    def mlp(self, x, chunk, projection_weights=None):
        states = (None, None) if projection_weights is None else projection_weights[2:]
        if not chunk or x.shape[0] <= chunk:
            return self.project(
                self.fc2, swiglu(self.project(self.fc1, x, states[0])), states[1]
            )
        result = []
        for start in range(0, x.shape[0], chunk):
            out = self.project(
                self.fc2,
                swiglu(self.project(self.fc1, x[start : start + chunk], states[0])),
                states[1],
            )
            mx.eval(out)
            result.append(out)
        return mx.concatenate(result, axis=0)

    def __call__(
        self,
        x,
        *,
        params=None,
        indices=None,
        rope=None,
        mlp_chunk=1024,
        query_chunk=0,
        kernel="native",
        threads=256,
        projection_weights=None,
    ):
        c = self.config
        if params is None:
            x = x + self.attn(rms_norm(x, self.n1, c.norm_eps), None, query_chunk)
            return x + self.mlp(rms_norm(x, self.n2, c.norm_eps), mlp_chunk)
        shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = params
        h = modulate(
            x,
            self.n1,
            shift_a,
            scale_a,
            indices,
            c.norm_eps,
            kernel=kernel,
            threads=threads,
        )
        x = gated_add(
            x,
            self.attn(h, rope, query_chunk, kernel, projection_weights),
            gate_a,
            indices,
        )
        h = modulate(
            x,
            self.n2,
            shift_m,
            scale_m,
            indices,
            c.norm_eps,
            kernel=kernel,
            threads=threads,
        )
        return gated_add(x, self.mlp(h, mlp_chunk, projection_weights), gate_m, indices)


class DiT:
    def __init__(
        self,
        model_path,
        *,
        cache_dir="cache",
        cache_gb=38,
        bits=None,
        group_size=64,
        mlp_chunk=0,
        query_chunk=0,
        kernel="native",
        threads=256,
        quantized_path=None,
        compile_blocks=False,
        eval_every=1,
        keep_refiners=False,
        weight_storage="bf16",
        async_blocks=False,
        cache_refined_text=True,
        prefetch_blocks=0,
        tune_matmul=False,
    ):
        self.path = Path(model_path)
        self.config = DiTConfig.load(self.path / "transformer/config.json")
        self.store = WeightStore(self.path / "transformer")
        self.core_store = self.store
        self.release = None
        release_path = self.path / "release.json"
        if release_path.exists():
            from .hub import load_release

            self.release = load_release(self.path)
            if quantized_path is not None or bits is not None:
                raise ValueError(
                    "Published variants already select their weight precision"
                )
            core_format = self.release["core_format"]
            if core_format == "quantized":
                meta = json.loads((self.path / "core/quantization.json").read_text())
                expected = hashlib.sha256(
                    (self.path / "transformer/config.json").read_bytes()
                ).hexdigest()
                if (
                    meta.get("format") != "h3-mlx-core-v1"
                    or meta.get("source_config_sha256") != expected
                    or meta.get("bits") != self.release["bits"]
                    or meta.get("group_size") != self.release["group_size"]
                ):
                    raise ValueError("Published quantized core/configuration mismatch")
                bits, group_size = meta["bits"], meta["group_size"]
                self.core_store = WeightStore(self.path / "core")
            elif core_format == "packed12":
                weight_storage = "packed12"
            if core_format != "packed12" and weight_storage == "packed12":
                raise ValueError("Select the packed12 download to use packed storage")
        if quantized_path is not None:
            meta = json.loads((Path(quantized_path) / "quantization.json").read_text())
            if meta.get("format") != "h3-mlx-core-v1":
                raise ValueError("Unsupported quantized checkpoint format")
            if meta["source_fingerprint"] != self.store.fingerprint():
                raise ValueError(
                    "Quantized core does not match this original checkpoint"
                )
            if bits is not None and bits != meta["bits"]:
                raise ValueError(
                    "--quantize conflicts with the converted core precision"
                )
            bits, group_size = meta["bits"], meta["group_size"]
            self.core_store = WeightStore(quantized_path)
        self.heads = WeightStore(self.path / "pdd_fused_heads.safetensors")
        self.cache_dir = Path(cache_dir)
        safe_limit = max(
            0, mx.device_info()["max_recommended_working_set_size"] - 8 * 1024**3
        )
        self.block_limit = min(int(cache_gb * 1024**3), safe_limit)
        self.blocks = OrderedDict()
        self.block_bytes = 0
        self.bits, self.group_size = bits, group_size
        self.mlp_chunk, self.query_chunk, self.kernel, self.threads = (
            mlp_chunk,
            query_chunk,
            kernel,
            threads,
        )
        self.metrics = {
            "block_load_s": 0.0,
            "block_compute_s": 0.0,
            "block_loads": 0,
            "block_wall_s": 0.0,
        }
        self.constants = {}
        if eval_every < 0:
            raise ValueError("eval_every must be nonnegative")
        self.compile_blocks = compile_blocks
        self.eval_every = eval_every
        self.keep_refiners = keep_refiners
        self.compiled_blocks = {}
        self.async_blocks = async_blocks
        self.cache_refined_text = cache_refined_text
        self.refiner_cache_hit = False
        if prefetch_blocks < 0:
            raise ValueError("prefetch_blocks must be nonnegative")
        self.prefetch_blocks = prefetch_blocks
        self.tune_matmul = tune_matmul
        if weight_storage not in ("bf16", "packed12"):
            raise ValueError("weight_storage must be bf16 or packed12")
        if weight_storage == "packed12" and (
            bits is not None or quantized_path is not None
        ):
            raise ValueError("Lossless packing cannot be combined with quantization")
        self.weight_storage = weight_storage
        self.packed_cache = None
        if weight_storage == "packed12":
            key = hashlib.sha256(
                (self.store.fingerprint() + "palette12-v1").encode()
            ).hexdigest()[:24]
            self.packed_cache = self.cache_dir / f"packed12-{key}"
            if self.release is not None:
                self.packed_cache = self.path / "core"
                for i in range(self.config.num_layers):
                    if not (self.packed_cache / f"blocks.{i}.safetensors").is_file():
                        raise FileNotFoundError(f"Missing published packed block {i}")

    def _constant(self, key):
        if key not in self.constants:
            self.constants[key] = self.store.get(key)
        return self.constants[key]

    def _linear(self, prefix):
        # All constants here are deliberately unquantized, including FP32 islands.
        if prefix not in self.constants:
            self.constants[prefix] = Linear.load(self.store, prefix)
        return self.constants[prefix]

    def _block(self, prefix):
        if prefix in self.blocks:
            self.blocks.move_to_end(prefix)
            return self.blocks[prefix]
        # Pin a fixed resident subset. A rotating LRU would evict every block
        # before its next denoise pass when the whole model does not fit.
        before = time.perf_counter()
        block = DiTBlock(
            self.core_store,
            prefix,
            self.config,
            bits=self.bits,
            group_size=self.group_size,
            packed_cache=self.packed_cache if prefix.startswith("blocks.") else None,
        )
        if self.tune_matmul and prefix.startswith("blocks."):
            from .tuning import TunedFC2

            block.fc2 = TunedFC2(block.fc2)
        self.metrics["block_load_s"] += time.perf_counter() - before
        self.metrics["block_loads"] += 1
        if self.block_bytes + block.nbytes <= self.block_limit:
            self.blocks[prefix] = block
            self.block_bytes += block.nbytes
        return block

    def prepare_modulation(self, sigmas, *, edit=False):
        """Read each 520 MB AdaLN weight once, then discard it permanently.
        A schedule-specific 6x3x6xH table replaces 26 GB of original BF16
        weights. No timestep interpolation or low-rank approximation is used.
        """
        c = self.config
        times = [1 - s for s in sigmas[:-1]] + ([0.999] if edit else [])
        if self.release is not None:
            from .hub import validate_modulation

            validate_modulation(self.release, sigmas)
            key = "editing" if edit else "generation"
            target = self.path / "modulation" / f"{key}.safetensors"
            if getattr(self, "modulation_key", None) == str(target):
                return
            table, metadata = mx.load(str(target), return_metadata=True)
            identity = json.loads(metadata.get("identity", "{}"))
            if (
                identity.get("times") != times
                or identity.get("config") != vars(c)
                or identity.get("method") != "bf16-linear-fp32-time"
            ):
                raise ValueError("Published modulation table identity mismatch")
            for i in range(c.num_layers):
                if table[f"block.{i}"].shape != (len(times), 3, 6, c.hidden_size):
                    raise ValueError("Published modulation table dimensions mismatch")
            if table["final"].shape != (len(times), 2, c.hidden_size):
                raise ValueError("Published final modulation dimensions mismatch")
            mx.eval(table)
            self.modulation = table
            self.modulation_key = str(target)
            return
        identity = {
            "version": 5,
            "source": self.store.fingerprint(),
            "times": times,
            "config": self.config.__dict__,
            "mlx": importlib.metadata.version("mlx"),
            "tf32": os.environ.get("MLX_ENABLE_TF32"),
            "method": "bf16-linear-fp32-time",
        }
        key = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest()[
            :24
        ]
        if getattr(self, "modulation_key", None) == key:
            return
        target = self.cache_dir / f"adaln-{key}.safetensors"
        if target.exists():
            log.info("Loading cached native AdaLN schedule %s", target)
            self.modulation = mx.load(str(target))
            mx.eval(self.modulation)
            self.modulation_key = key
            return
        log.info("Precomputing AdaLN: releasing ~26 GB of schedule-only weights")
        half = c.timestep_input_dim // 2
        freq = mx.exp(-math.log(10000.0) * mx.arange(half, dtype=mx.float32) / half)
        args = mx.array(times, mx.float32)[:, None] * freq[None]
        t = mx.concatenate((mx.cos(args), mx.sin(args)), axis=-1)
        t = self._linear("time_embedder.proj_out")(
            silu(self._linear("time_embedder.proj_in")(t))
        )
        t = silu(t).astype(mx.bfloat16)
        mx.eval(t)
        table = {}
        for i in range(c.num_layers):
            prefix = f"blocks.{i}.adaln_proj.linear"
            layer = Linear.load(self.store, prefix)
            table[f"block.{i}"] = layer(t).reshape(len(times), 3, 6, c.hidden_size)
            mx.eval(table[f"block.{i}"])
            del layer
            if (i + 1) % 10 == 0:
                log.info("AdaLN %d/%d", i + 1, c.num_layers)
                mx.clear_cache()
        layer = Linear.load(self.store, "final_layer.adaln_proj.linear")
        table["final"] = layer(t).reshape(len(times), 2, c.hidden_size)
        mx.eval(table)
        self.cache_dir.mkdir(parents=True, exist_ok=True)
        tmp = target.with_suffix(".tmp.safetensors")
        mx.save_safetensors(
            str(tmp), table, metadata={"identity": json.dumps(identity)}
        )
        tmp.replace(target)
        self.modulation = table
        self.modulation_key = key
        mx.clear_cache()

    def refine(self, text):
        self.refiner_cache_hit = False
        if self.cache_refined_text:
            identity = {
                "format": "h3-refined-text-v1",
                "source": self.store.fingerprint(),
                "core_source": self.core_store.fingerprint(),
                "weight_storage": self.weight_storage,
                "config": vars(self.config),
                "bits": self.bits,
                "group_size": self.group_size,
                "mlp_chunk": self.mlp_chunk,
                "query_chunk": self.query_chunk,
                "mlx": importlib.metadata.version("mlx"),
                "tf32": os.environ.get("MLX_ENABLE_TF32"),
                "text_shape": text.shape,
                "text_sha256": hashlib.sha256(
                    np.asarray(text.astype(mx.bfloat16).view(mx.uint16)).tobytes()
                ).hexdigest(),
            }
            key = hashlib.sha256(
                json.dumps(identity, sort_keys=True).encode()
            ).hexdigest()[:24]
            target = self.cache_dir / f"refined-{key}.safetensors"
            if target.exists():
                x = mx.load(str(target))["hidden"]
                mx.eval(x)
                self.refiner_cache_hit = True
                return x
        x = self._linear("condition_proj")(text.astype(mx.bfloat16))
        for i in range(self.config.token_refiner_num_layers):
            block = self._block(f"token_refiner.blocks.{i}")
            x = block(x, mlp_chunk=self.mlp_chunk, query_chunk=self.query_chunk)
            mx.eval(x)
        x = rms_norm(
            x,
            self._constant("token_refiner.final_norm.weight"),
            self.config.final_norm_eps,
        )
        mx.eval(x)
        if self.cache_refined_text:
            target.parent.mkdir(parents=True, exist_ok=True)
            temporary = target.with_name(
                f".{target.stem}-{uuid.uuid4().hex}.safetensors"
            )
            try:
                mx.save_safetensors(
                    str(temporary),
                    {"hidden": x},
                    metadata={"identity": json.dumps(identity)},
                )
                temporary.replace(target)
            finally:
                temporary.unlink(missing_ok=True)
        # Refiners run only once. Free their ~1.5 GB for the fifty main blocks.
        for key in list(self.blocks):
            if key.startswith("token_refiner.") and not self.keep_refiners:
                old = self.blocks.pop(key)
                self.block_bytes -= old.nbytes
                del old
        mx.clear_cache()
        return x

    def predict(
        self, rows, text, step, h, w, *, references=(), text_tags=None, on_layer=None
    ):
        ntext = text.shape[0]
        nref = sum(rh * rw // 4 for rh, rw in references)
        x = mx.concatenate(
            (text, self._linear("video_patch_proj")(rows).astype(mx.bfloat16)), axis=0
        )
        positions = image_positions(ntext, h, w, references)
        rope = rope_cache(positions, self._constant("rope.inv_freq"))
        tags = mx.concatenate(
            (
                mx.ones(ntext, mx.int32) if text_tags is None else text_tags,
                mx.zeros(rows.shape[0], mx.int32),
            )
        )
        # Table has current timestep for every modality; editing adds the fixed .999 anchor.
        time_ids = mx.concatenate(
            (
                mx.zeros(ntext, mx.int32),
                mx.full(nref, len(self.modulation["final"]) - 1, mx.int32),
                mx.zeros(rows.shape[0] - nref, mx.int32),
            )
        )
        loop_start = time.perf_counter()
        loads_before = self.metrics["block_load_s"]
        prefetch = self.prefetch_blocks and any(
            f"blocks.{i}" not in self.blocks for i in range(self.config.num_layers)
        )
        sync_every = self.prefetch_blocks if prefetch else self.eval_every
        for i in range(self.config.num_layers):
            block = self._block(f"blocks.{i}")
            table = self.modulation[f"block.{i}"]
            # Select current and condition timestep as a compact [6,6,H] table.
            current = table[step]
            if references:
                compact = mx.concatenate((current, table[-1]), axis=0)
                indices = mx.where(time_ids > 0, tags + 3, tags)
            else:
                compact, indices = current, tags
            params = tuple(compact[:, j] for j in range(6))
            if self.compile_blocks:
                # One compiled callable per resident block; parameters, positions
                # and input remain dynamic. Compiled closures must never pin
                # streamed weights beyond the configured residency budget.
                key = f"blocks.{i}"
                runner = self.compiled_blocks.get(key)
                if runner is None:

                    def forward(
                        x, params, indices, rope, projection_weights=None, block=block
                    ):
                        return block(
                            x,
                            params=params,
                            indices=indices,
                            rope=rope,
                            mlp_chunk=self.mlp_chunk,
                            query_chunk=self.query_chunk,
                            kernel=self.kernel,
                            threads=self.threads,
                            projection_weights=projection_weights,
                        )

                    runner = mx.compile(forward)
                    if key in self.blocks:
                        self.compiled_blocks[key] = runner
                if self.weight_storage == "packed12":
                    # Explicit dynamic inputs prevent constant-folding the
                    # reconstruction into a retained dense weight matrix.
                    state = tuple(
                        p.weight.state()
                        for p in (block.qkv, block.out, block.fc1, block.fc2)
                    )
                    x = runner(x, params, indices, rope, state)
                else:
                    x = runner(x, params, indices, rope)
            else:
                x = block(
                    x,
                    params=params,
                    indices=indices,
                    rope=rope,
                    mlp_chunk=self.mlp_chunk,
                    query_chunk=self.query_chunk,
                    kernel=self.kernel,
                    threads=self.threads,
                )
            if (self.async_blocks or prefetch) and not on_layer:
                mx.async_eval(x)
            if on_layer or (sync_every and (i + 1) % sync_every == 0):
                mx.eval(x)
            if on_layer:
                on_layer(step, i, x)
        mx.eval(x)
        self.metrics["block_wall_s"] += time.perf_counter() - loop_start
        self.metrics["block_compute_s"] += (
            time.perf_counter()
            - loop_start
            - (self.metrics["block_load_s"] - loads_before)
        )
        final = self.modulation["final"][step]
        if references:
            final_table = mx.stack((final, self.modulation["final"][-1]))
            final_indices = (time_ids > 0).astype(mx.int32)
        else:
            final_table = final[None]
            final_indices = mx.zeros(x.shape[0], mx.int32)
        hidden = modulate(
            x,
            self._constant("final_layer.norm.weight"),
            final_table[:, 0],
            final_table[:, 1],
            final_indices,
            self.config.final_norm_eps,
            kernel=self.kernel,
            threads=self.threads,
        )
        # Native per-step PDD output head and FP32 projection/Euler state.
        weight = self.heads.get("video_out.weight", rows=slice(step, step + 1))[0]
        bias = self.heads.get("video_out.bias", rows=slice(step, step + 1))[0]
        velocity = mx.addmm(bias, hidden.astype(mx.float32), weight.T)
        mx.eval(velocity)
        return velocity[ntext:]

    def clear(self):
        self.compiled_blocks.clear()
        self.blocks.clear()
        self.constants.clear()
        self.block_bytes = 0
        self.store.clear()
        if self.core_store is not self.store:
            self.core_store.clear()
        self.heads.clear()
        gc.collect()
        mx.clear_cache()

	from __future__ import annotations

	import gc
	import hashlib
	import importlib.metadata
	import json
	import logging
	import math
	import os
	import time
	import uuid
	from collections import OrderedDict
	from pathlib import Path

	import mlx.core as mx
	import numpy as np

	from .config import DiTConfig
	from .layout import image_positions, rope_cache
	from .ops import attention, gated_add, modulate, rms_norm, rotate, silu, swiglu
	from .weights import Linear, WeightStore

	log = logging.getLogger(__name__)


	class DiTBlock:
	def __init__(
	self, store, prefix, config, *, bits=None, group_size=64, packed_cache=None
	):
	self.config = config
	self.n1 = store.get(prefix + ".norm1.weight")
	self.n2 = store.get(prefix + ".norm2.weight")
	self.qn = store.get(prefix + ".attn.q_norm.weight")
	self.kn = store.get(prefix + ".attn.k_norm.weight")
	if packed_cache is not None:
	from .packed import PackedBF16, PackedLinear

	cache_path = Path(packed_cache) / f"{prefix}.safetensors"
	names = ("attn.qkv_proj", "attn.out_proj", "mlp.fc1", "mlp.fc2")
	attrs = ("qkv", "out", "fc1", "fc2")
	if cache_path.exists():
	data, metadata = mx.load(str(cache_path), return_metadata=True)
	if metadata.get("format") != "h3-palette12-v1":
	raise ValueError("Unsupported packed checkpoint cache")
	shapes = json.loads(metadata["shapes"])
	for attr, name in zip(attrs, names, strict=True):
	packed = PackedBF16.from_tensors(data, name, shapes[name])
	setattr(self, attr, PackedLinear(packed, data.get(name + ".bias")))
	mx.eval(data)
	return
	grouped = (config.num_attention_heads, config.attention_head_dim)
	self.qkv = Linear.load(
	store,
	prefix + ".attn.qkv_proj",
	bits=bits,
	group_size=group_size,
	grouped_heads=grouped,
	)
	self.out = Linear.load(
	store, prefix + ".attn.out_proj", bits=bits, group_size=group_size
	)
	self.fc1 = Linear.load(
	store, prefix + ".mlp.fc1", bits=bits, group_size=group_size
	)
	self.fc2 = Linear.load(
	store, prefix + ".mlp.fc2", bits=bits, group_size=group_size
	)
	if packed_cache is not None:
	data, shapes = {}, {}
	for attr, name in zip(attrs, names, strict=True):
	linear = getattr(self, attr)
	packed = PackedLinear.from_linear(linear)
	setattr(self, attr, packed)
	shapes[name] = packed.weight.shape
	data.update(packed.weight.tensors(name))
	if packed.bias is not None:
	data[name + ".bias"] = packed.bias
	cache_path.parent.mkdir(parents=True, exist_ok=True)
	temporary = cache_path.with_name(
	f".{cache_path.stem}-{uuid.uuid4().hex}.safetensors"
	)
	try:
	mx.save_safetensors(
	str(temporary),
	data,
	metadata={
	"format": "h3-palette12-v1",
	"shapes": json.dumps(shapes),
	},
	)
	temporary.replace(cache_path)
	finally:
	temporary.unlink(missing_ok=True)

	@property
	def nbytes(self):
	return sum(
	x.nbytes
	for x in (
	self.n1,
	self.n2,
	self.qn,
	self.kn,
	self.qkv,
	self.out,
	self.fc1,
	self.fc2,
	)
	)

	@staticmethod
	def project(linear, x, state):
	return linear(x) if state is None else linear(x, state=state)

	def attn(self, x, rope, query_chunk, kernel="native", projection_weights=None):
	c = self.config
	states = (None, None) if projection_weights is None else projection_weights[:2]
	qkv = self.project(self.qkv, x, states[0]).reshape(
	1, x.shape[0], 3, c.num_attention_heads, c.attention_head_dim
	)
	q, k, v = (qkv[:, :, i] for i in range(3))
	exact_qk = False
	if kernel in ("metal-exact", "metal-qk-exact") and rope is not None:
	from .tuning import qk_norm_rope_exact, supports_qk

	if supports_qk(qkv[0]):
	q, k = qk_norm_rope_exact(
	qkv[0], self.qn, self.kn, *rope, c.qk_norm_eps
	)
	exact_qk = True
	if kernel == "metal-qk" and rope is not None:
	from .kernels import qk_norm_rope

	q, k = qk_norm_rope(qkv[0], self.qn, self.kn, *rope, c.qk_norm_eps)
	elif not exact_qk:
	q = rms_norm(q, self.qn, c.qk_norm_eps)
	k = rms_norm(k, self.kn, c.qk_norm_eps)
	if rope is not None:
	q, k = rotate(q, rope), rotate(k, rope)
	a = attention(q, k, v, query_chunk=query_chunk).reshape(x.shape[0], -1)
	return self.project(self.out, a, states[1])

	def mlp(self, x, chunk, projection_weights=None):
	states = (None, None) if projection_weights is None else projection_weights[2:]
	if not chunk or x.shape[0] <= chunk:
	return self.project(
	self.fc2, swiglu(self.project(self.fc1, x, states[0])), states[1]
	)
	result = []
	for start in range(0, x.shape[0], chunk):
	out = self.project(
	self.fc2,
	swiglu(self.project(self.fc1, x[start : start + chunk], states[0])),
	states[1],
	)
	mx.eval(out)
	result.append(out)
	return mx.concatenate(result, axis=0)

	def __call__(
	self,
	x,
	*,
	params=None,
	indices=None,
	rope=None,
	mlp_chunk=1024,
	query_chunk=0,
	kernel="native",
	threads=256,
	projection_weights=None,
	):
	c = self.config
	if params is None:
	x = x + self.attn(rms_norm(x, self.n1, c.norm_eps), None, query_chunk)
	return x + self.mlp(rms_norm(x, self.n2, c.norm_eps), mlp_chunk)
	shift_a, scale_a, gate_a, shift_m, scale_m, gate_m = params
	h = modulate(
	x,
	self.n1,
	shift_a,
	scale_a,
	indices,
	c.norm_eps,
	kernel=kernel,
	threads=threads,
	)
	x = gated_add(
	x,
	self.attn(h, rope, query_chunk, kernel, projection_weights),
	gate_a,
	indices,
	)
	h = modulate(
	x,
	self.n2,
	shift_m,
	scale_m,
	indices,
	c.norm_eps,
	kernel=kernel,
	threads=threads,
	)
	return gated_add(x, self.mlp(h, mlp_chunk, projection_weights), gate_m, indices)


	class DiT:
	def __init__(
	self,
	model_path,
	*,
	cache_dir="cache",
	cache_gb=38,
	bits=None,
	group_size=64,
	mlp_chunk=0,
	query_chunk=0,
	kernel="native",
	threads=256,
	quantized_path=None,
	compile_blocks=False,
	eval_every=1,
	keep_refiners=False,
	weight_storage="bf16",
	async_blocks=False,
	cache_refined_text=True,
	prefetch_blocks=0,
	tune_matmul=False,
	):
	self.path = Path(model_path)
	self.config = DiTConfig.load(self.path / "transformer/config.json")
	self.store = WeightStore(self.path / "transformer")
	self.core_store = self.store
	self.release = None
	release_path = self.path / "release.json"
	if release_path.exists():
	from .hub import load_release

	self.release = load_release(self.path)
	if quantized_path is not None or bits is not None:
	raise ValueError(
	"Published variants already select their weight precision"
	)
	core_format = self.release["core_format"]
	if core_format == "quantized":
	meta = json.loads((self.path / "core/quantization.json").read_text())
	expected = hashlib.sha256(
	(self.path / "transformer/config.json").read_bytes()
	).hexdigest()
	if (
	meta.get("format") != "h3-mlx-core-v1"
	or meta.get("source_config_sha256") != expected
	or meta.get("bits") != self.release["bits"]
	or meta.get("group_size") != self.release["group_size"]
	):
	raise ValueError("Published quantized core/configuration mismatch")
	bits, group_size = meta["bits"], meta["group_size"]
	self.core_store = WeightStore(self.path / "core")
	elif core_format == "packed12":
	weight_storage = "packed12"
	if core_format != "packed12" and weight_storage == "packed12":
	raise ValueError("Select the packed12 download to use packed storage")
	if quantized_path is not None:
	meta = json.loads((Path(quantized_path) / "quantization.json").read_text())
	if meta.get("format") != "h3-mlx-core-v1":
	raise ValueError("Unsupported quantized checkpoint format")
	if meta["source_fingerprint"] != self.store.fingerprint():
	raise ValueError(
	"Quantized core does not match this original checkpoint"
	)
	if bits is not None and bits != meta["bits"]:
	raise ValueError(
	"--quantize conflicts with the converted core precision"
	)
	bits, group_size = meta["bits"], meta["group_size"]
	self.core_store = WeightStore(quantized_path)
	self.heads = WeightStore(self.path / "pdd_fused_heads.safetensors")
	self.cache_dir = Path(cache_dir)
	safe_limit = max(
	0, mx.device_info()["max_recommended_working_set_size"] - 8 * 1024**3
	)
	self.block_limit = min(int(cache_gb * 1024**3), safe_limit)
	self.blocks = OrderedDict()
	self.block_bytes = 0
	self.bits, self.group_size = bits, group_size
	self.mlp_chunk, self.query_chunk, self.kernel, self.threads = (
	mlp_chunk,
	query_chunk,
	kernel,
	threads,
	)
	self.metrics = {
	"block_load_s": 0.0,
	"block_compute_s": 0.0,
	"block_loads": 0,
	"block_wall_s": 0.0,
	}
	self.constants = {}
	if eval_every < 0:
	raise ValueError("eval_every must be nonnegative")
	self.compile_blocks = compile_blocks
	self.eval_every = eval_every
	self.keep_refiners = keep_refiners
	self.compiled_blocks = {}
	self.async_blocks = async_blocks
	self.cache_refined_text = cache_refined_text
	self.refiner_cache_hit = False
	if prefetch_blocks < 0:
	raise ValueError("prefetch_blocks must be nonnegative")
	self.prefetch_blocks = prefetch_blocks
	self.tune_matmul = tune_matmul
	if weight_storage not in ("bf16", "packed12"):
	raise ValueError("weight_storage must be bf16 or packed12")
	if weight_storage == "packed12" and (
	bits is not None or quantized_path is not None
	):
	raise ValueError("Lossless packing cannot be combined with quantization")
	self.weight_storage = weight_storage
	self.packed_cache = None
	if weight_storage == "packed12":
	key = hashlib.sha256(
	(self.store.fingerprint() + "palette12-v1").encode()
	).hexdigest()[:24]
	self.packed_cache = self.cache_dir / f"packed12-{key}"
	if self.release is not None:
	self.packed_cache = self.path / "core"
	for i in range(self.config.num_layers):
	if not (self.packed_cache / f"blocks.{i}.safetensors").is_file():
	raise FileNotFoundError(f"Missing published packed block {i}")

	def _constant(self, key):
	if key not in self.constants:
	self.constants[key] = self.store.get(key)
	return self.constants[key]

	def _linear(self, prefix):
	# All constants here are deliberately unquantized, including FP32 islands.
	if prefix not in self.constants:
	self.constants[prefix] = Linear.load(self.store, prefix)
	return self.constants[prefix]

	def _block(self, prefix):
	if prefix in self.blocks:
	self.blocks.move_to_end(prefix)
	return self.blocks[prefix]
	# Pin a fixed resident subset. A rotating LRU would evict every block
	# before its next denoise pass when the whole model does not fit.
	before = time.perf_counter()
	block = DiTBlock(
	self.core_store,
	prefix,
	self.config,
	bits=self.bits,
	group_size=self.group_size,
	packed_cache=self.packed_cache if prefix.startswith("blocks.") else None,
	)
	if self.tune_matmul and prefix.startswith("blocks."):
	from .tuning import TunedFC2

	block.fc2 = TunedFC2(block.fc2)
	self.metrics["block_load_s"] += time.perf_counter() - before
	self.metrics["block_loads"] += 1
	if self.block_bytes + block.nbytes <= self.block_limit:
	self.blocks[prefix] = block
	self.block_bytes += block.nbytes
	return block

	def prepare_modulation(self, sigmas, *, edit=False):
	"""Read each 520 MB AdaLN weight once, then discard it permanently.

	A schedule-specific 6x3x6xH table replaces 26 GB of original BF16
	weights. No timestep interpolation or low-rank approximation is used.
	"""
	c = self.config
	times = [1 - s for s in sigmas[:-1]] + ([0.999] if edit else [])
	if self.release is not None:
	from .hub import validate_modulation

	validate_modulation(self.release, sigmas)
	key = "editing" if edit else "generation"
	target = self.path / "modulation" / f"{key}.safetensors"
	if getattr(self, "modulation_key", None) == str(target):
	return
	table, metadata = mx.load(str(target), return_metadata=True)
	identity = json.loads(metadata.get("identity", "{}"))
	if (
	identity.get("times") != times
	or identity.get("config") != vars(c)
	or identity.get("method") != "bf16-linear-fp32-time"
	):
	raise ValueError("Published modulation table identity mismatch")
	for i in range(c.num_layers):
	if table[f"block.{i}"].shape != (len(times), 3, 6, c.hidden_size):
	raise ValueError("Published modulation table dimensions mismatch")
	if table["final"].shape != (len(times), 2, c.hidden_size):
	raise ValueError("Published final modulation dimensions mismatch")
	mx.eval(table)
	self.modulation = table
	self.modulation_key = str(target)
	return
	identity = {
	"version": 5,
	"source": self.store.fingerprint(),
	"times": times,
	"config": self.config.__dict__,
	"mlx": importlib.metadata.version("mlx"),
	"tf32": os.environ.get("MLX_ENABLE_TF32"),
	"method": "bf16-linear-fp32-time",
	}
	key = hashlib.sha256(json.dumps(identity, sort_keys=True).encode()).hexdigest()[
	:24
	]
	if getattr(self, "modulation_key", None) == key:
	return
	target = self.cache_dir / f"adaln-{key}.safetensors"
	if target.exists():
	log.info("Loading cached native AdaLN schedule %s", target)
	self.modulation = mx.load(str(target))
	mx.eval(self.modulation)
	self.modulation_key = key
	return
	log.info("Precomputing AdaLN: releasing ~26 GB of schedule-only weights")
	half = c.timestep_input_dim // 2
	freq = mx.exp(-math.log(10000.0) * mx.arange(half, dtype=mx.float32) / half)
	args = mx.array(times, mx.float32)[:, None] * freq[None]
	t = mx.concatenate((mx.cos(args), mx.sin(args)), axis=-1)
	t = self._linear("time_embedder.proj_out")(
	silu(self._linear("time_embedder.proj_in")(t))
	)
	t = silu(t).astype(mx.bfloat16)
	mx.eval(t)
	table = {}
	for i in range(c.num_layers):
	prefix = f"blocks.{i}.adaln_proj.linear"
	layer = Linear.load(self.store, prefix)
	table[f"block.{i}"] = layer(t).reshape(len(times), 3, 6, c.hidden_size)
	mx.eval(table[f"block.{i}"])
	del layer
	if (i + 1) % 10 == 0:
	log.info("AdaLN %d/%d", i + 1, c.num_layers)
	mx.clear_cache()
	layer = Linear.load(self.store, "final_layer.adaln_proj.linear")
	table["final"] = layer(t).reshape(len(times), 2, c.hidden_size)
	mx.eval(table)
	self.cache_dir.mkdir(parents=True, exist_ok=True)
	tmp = target.with_suffix(".tmp.safetensors")
	mx.save_safetensors(
	str(tmp), table, metadata={"identity": json.dumps(identity)}
	)
	tmp.replace(target)
	self.modulation = table
	self.modulation_key = key
	mx.clear_cache()

	def refine(self, text):
	self.refiner_cache_hit = False
	if self.cache_refined_text:
	identity = {
	"format": "h3-refined-text-v1",
	"source": self.store.fingerprint(),
	"core_source": self.core_store.fingerprint(),
	"weight_storage": self.weight_storage,
	"config": vars(self.config),
	"bits": self.bits,
	"group_size": self.group_size,
	"mlp_chunk": self.mlp_chunk,
	"query_chunk": self.query_chunk,
	"mlx": importlib.metadata.version("mlx"),
	"tf32": os.environ.get("MLX_ENABLE_TF32"),
	"text_shape": text.shape,
	"text_sha256": hashlib.sha256(
	np.asarray(text.astype(mx.bfloat16).view(mx.uint16)).tobytes()
	).hexdigest(),
	}
	key = hashlib.sha256(
	json.dumps(identity, sort_keys=True).encode()
	).hexdigest()[:24]
	target = self.cache_dir / f"refined-{key}.safetensors"
	if target.exists():
	x = mx.load(str(target))["hidden"]
	mx.eval(x)
	self.refiner_cache_hit = True
	return x
	x = self._linear("condition_proj")(text.astype(mx.bfloat16))
	for i in range(self.config.token_refiner_num_layers):
	block = self._block(f"token_refiner.blocks.{i}")
	x = block(x, mlp_chunk=self.mlp_chunk, query_chunk=self.query_chunk)
	mx.eval(x)
	x = rms_norm(
	x,
	self._constant("token_refiner.final_norm.weight"),
	self.config.final_norm_eps,
	)
	mx.eval(x)
	if self.cache_refined_text:
	target.parent.mkdir(parents=True, exist_ok=True)
	temporary = target.with_name(
	f".{target.stem}-{uuid.uuid4().hex}.safetensors"
	)
	try:
	mx.save_safetensors(
	str(temporary),
	{"hidden": x},
	metadata={"identity": json.dumps(identity)},
	)
	temporary.replace(target)
	finally:
	temporary.unlink(missing_ok=True)
	# Refiners run only once. Free their ~1.5 GB for the fifty main blocks.
	for key in list(self.blocks):
	if key.startswith("token_refiner.") and not self.keep_refiners:
	old = self.blocks.pop(key)
	self.block_bytes -= old.nbytes
	del old
	mx.clear_cache()
	return x

	def predict(
	self, rows, text, step, h, w, *, references=(), text_tags=None, on_layer=None
	):
	ntext = text.shape[0]
	nref = sum(rh * rw // 4 for rh, rw in references)
	x = mx.concatenate(
	(text, self._linear("video_patch_proj")(rows).astype(mx.bfloat16)), axis=0
	)
	positions = image_positions(ntext, h, w, references)
	rope = rope_cache(positions, self._constant("rope.inv_freq"))
	tags = mx.concatenate(
	(
	mx.ones(ntext, mx.int32) if text_tags is None else text_tags,
	mx.zeros(rows.shape[0], mx.int32),
	)
	)
	# Table has current timestep for every modality; editing adds the fixed .999 anchor.
	time_ids = mx.concatenate(
	(
	mx.zeros(ntext, mx.int32),
	mx.full(nref, len(self.modulation["final"]) - 1, mx.int32),
	mx.zeros(rows.shape[0] - nref, mx.int32),
	)
	)
	loop_start = time.perf_counter()
	loads_before = self.metrics["block_load_s"]
	prefetch = self.prefetch_blocks and any(
	f"blocks.{i}" not in self.blocks for i in range(self.config.num_layers)
	)
	sync_every = self.prefetch_blocks if prefetch else self.eval_every
	for i in range(self.config.num_layers):
	block = self._block(f"blocks.{i}")
	table = self.modulation[f"block.{i}"]
	# Select current and condition timestep as a compact [6,6,H] table.
	current = table[step]
	if references:
	compact = mx.concatenate((current, table[-1]), axis=0)
	indices = mx.where(time_ids > 0, tags + 3, tags)
	else:
	compact, indices = current, tags
	params = tuple(compact[:, j] for j in range(6))
	if self.compile_blocks:
	# One compiled callable per resident block; parameters, positions
	# and input remain dynamic. Compiled closures must never pin
	# streamed weights beyond the configured residency budget.
	key = f"blocks.{i}"
	runner = self.compiled_blocks.get(key)
	if runner is None:

	def forward(
	x, params, indices, rope, projection_weights=None, block=block
	):
	return block(
	x,
	params=params,
	indices=indices,
	rope=rope,
	mlp_chunk=self.mlp_chunk,
	query_chunk=self.query_chunk,
	kernel=self.kernel,
	threads=self.threads,
	projection_weights=projection_weights,
	)

	runner = mx.compile(forward)
	if key in self.blocks:
	self.compiled_blocks[key] = runner
	if self.weight_storage == "packed12":
	# Explicit dynamic inputs prevent constant-folding the
	# reconstruction into a retained dense weight matrix.
	state = tuple(
	p.weight.state()
	for p in (block.qkv, block.out, block.fc1, block.fc2)
	)
	x = runner(x, params, indices, rope, state)
	else:
	x = runner(x, params, indices, rope)
	else:
	x = block(
	x,
	params=params,
	indices=indices,
	rope=rope,
	mlp_chunk=self.mlp_chunk,
	query_chunk=self.query_chunk,
	kernel=self.kernel,
	threads=self.threads,
	)
	if (self.async_blocks or prefetch) and not on_layer:
	mx.async_eval(x)
	if on_layer or (sync_every and (i + 1) % sync_every == 0):
	mx.eval(x)
	if on_layer:
	on_layer(step, i, x)
	mx.eval(x)
	self.metrics["block_wall_s"] += time.perf_counter() - loop_start
	self.metrics["block_compute_s"] += (
	time.perf_counter()
	- loop_start
	- (self.metrics["block_load_s"] - loads_before)
	)
	final = self.modulation["final"][step]
	if references:
	final_table = mx.stack((final, self.modulation["final"][-1]))
	final_indices = (time_ids > 0).astype(mx.int32)
	else:
	final_table = final[None]
	final_indices = mx.zeros(x.shape[0], mx.int32)
	hidden = modulate(
	x,
	self._constant("final_layer.norm.weight"),
	final_table[:, 0],
	final_table[:, 1],
	final_indices,
	self.config.final_norm_eps,
	kernel=self.kernel,
	threads=self.threads,
	)
	# Native per-step PDD output head and FP32 projection/Euler state.
	weight = self.heads.get("video_out.weight", rows=slice(step, step + 1))[0]
	bias = self.heads.get("video_out.bias", rows=slice(step, step + 1))[0]
	velocity = mx.addmm(bias, hidden.astype(mx.float32), weight.T)
	mx.eval(velocity)
	return velocity[ntext:]

	def clear(self):
	self.compiled_blocks.clear()
	self.blocks.clear()
	self.constants.clear()
	self.block_bytes = 0
	self.store.clear()
	if self.core_store is not self.store:
	self.core_store.clear()
	self.heads.clear()
	gc.collect()
	mx.clear_cache()