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()