"""DiT blocks for Anima MLX.""" from __future__ import annotations from pathlib import Path from typing import Any, Mapping from .attention import scaled_dot_product_attention from .primitives import gelu, layer_norm, linear, rms_norm, silu def apply_cosmos_rotary_pos_emb(x: Any, freqs: Any) -> Any: import mlx.core as mx x_pairs = x.reshape(*x.shape[:-1], 2, -1) x_pairs = mx.swapaxes(x_pairs, -2, -1) x_pairs = mx.expand_dims(x_pairs.astype(mx.float32), axis=-2) out = freqs[..., 0] * x_pairs[..., 0] + freqs[..., 1] * x_pairs[..., 1] return mx.swapaxes(out, -1, -2).reshape(*x.shape).astype(x.dtype) def video_rope3d(shape: tuple[int, int, int, int, int], head_dim: int = 128) -> Any: import mlx.core as mx _, t_len, h_len, w_len, _ = shape dim_h = head_dim // 6 * 2 dim_w = dim_h dim_t = head_dim - 2 * dim_h dim_spatial_range = mx.arange(0, dim_h, 2).astype(mx.float32)[: dim_h // 2] / dim_h dim_temporal_range = mx.arange(0, dim_t, 2).astype(mx.float32)[: dim_t // 2] / dim_t h_ntk_factor = 4.0 ** (dim_h / (dim_h - 2)) w_ntk_factor = 4.0 ** (dim_w / (dim_w - 2)) t_ntk_factor = 1.0 ** (dim_t / (dim_t - 2)) h_freqs = 1.0 / ((10000.0 * h_ntk_factor) ** dim_spatial_range) w_freqs = 1.0 / ((10000.0 * w_ntk_factor) ** dim_spatial_range) t_freqs = 1.0 / ((10000.0 * t_ntk_factor) ** dim_temporal_range) seq = mx.arange(max(h_len, w_len, t_len)).astype(mx.float32) half_h = mx.outer(seq[:h_len], h_freqs) half_w = mx.outer(seq[:w_len], w_freqs) half_t = mx.outer(seq[:t_len], t_freqs) def stack_freqs(x: Any) -> Any: return mx.stack([mx.cos(x), -mx.sin(x), mx.sin(x), mx.cos(x)], axis=-1) half_h = stack_freqs(half_h) half_w = stack_freqs(half_w) half_t = stack_freqs(half_t) emb_t = mx.broadcast_to(half_t[:, None, None, :, :], (t_len, h_len, w_len, dim_t // 2, 4)) emb_h = mx.broadcast_to(half_h[None, :, None, :, :], (t_len, h_len, w_len, dim_h // 2, 4)) emb_w = mx.broadcast_to(half_w[None, None, :, :, :], (t_len, h_len, w_len, dim_w // 2, 4)) emb = mx.concatenate([emb_t, emb_h, emb_w], axis=-2) return emb.reshape(t_len * h_len * w_len, head_dim // 2, 2, 2).astype(mx.float32) class DiTAttention: def __init__(self, weights: Mapping[str, Any], prefix: str, *, n_heads: int = 16, head_dim: int = 128) -> None: self.weights = weights self.prefix = prefix self.n_heads = n_heads self.head_dim = head_dim def _weight(self, name: str) -> Any: return self.weights[f"{self.prefix}.{name}"] def _linear(self, name: str, x: Any) -> Any: return linear(x, self._weight(f"{name}.weight")) def __call__( self, x: Any, *, context: Any | None = None, rope_emb: Any | None = None, trace: dict[str, Any] | None = None, trace_prefix: str = "", ) -> Any: import mlx.core as mx is_selfattn = context is None context = x if context is None else context input_shape = x.shape[:-1] context_shape = context.shape[:-1] q = self._linear("q_proj", x).reshape(*input_shape, self.n_heads, self.head_dim) k = self._linear("k_proj", context).reshape(*context_shape, self.n_heads, self.head_dim) v = self._linear("v_proj", context).reshape(*context_shape, self.n_heads, self.head_dim) if trace is not None: trace[f"{trace_prefix}q_proj"] = q trace[f"{trace_prefix}k_proj"] = k trace[f"{trace_prefix}v_proj"] = v q = rms_norm(q, self._weight("q_norm.weight")) k = rms_norm(k, self._weight("k_norm.weight")) if trace is not None: trace[f"{trace_prefix}q_norm"] = q trace[f"{trace_prefix}k_norm"] = k if is_selfattn and rope_emb is not None: q = apply_cosmos_rotary_pos_emb(q, rope_emb) k = apply_cosmos_rotary_pos_emb(k, rope_emb) if trace is not None: trace[f"{trace_prefix}q_rope"] = q trace[f"{trace_prefix}k_rope"] = k q = mx.swapaxes(q, 1, 2) k = mx.swapaxes(k, 1, 2) v = mx.swapaxes(v, 1, 2) out = scaled_dot_product_attention(q, k, v) out = mx.swapaxes(out, 1, 2).reshape(*input_shape, self.n_heads * self.head_dim) if trace is not None: trace[f"{trace_prefix}attention_kernel_output"] = out output = self._linear("output_proj", out) if trace is not None: trace[f"{trace_prefix}output_proj"] = output trace[f"{trace_prefix}output"] = output return output class DiTBlock: def __init__(self, weights: Mapping[str, Any], prefix: str = "blocks.0") -> None: self.weights = self._normalize_keys(weights) self.prefix = prefix self.self_attn = DiTAttention(self.weights, f"{prefix}.self_attn") self.cross_attn = DiTAttention(self.weights, f"{prefix}.cross_attn") @classmethod def from_safetensors(cls, path: str | Path, *, block_index: int = 0, dtype: str = "float32") -> "DiTBlock": from anima_mlx.utils.weights import load_mlx_safetensors_subset path = _resolve_diffusion_block_path(path, block_index) prefix = f"net.blocks.{block_index}." weights = load_mlx_safetensors_subset(path, prefix=prefix, strip_prefix="net.", dtype=dtype) return cls(weights, prefix=f"blocks.{block_index}") @staticmethod def _normalize_keys(weights: Mapping[str, Any]) -> dict[str, Any]: normalized: dict[str, Any] = {} for key, value in weights.items(): if key.startswith("net."): key = key.removeprefix("net.") normalized[key] = value return normalized def _weight(self, name: str) -> Any: return self.weights[f"{self.prefix}.{name}"] def _adaln(self, name: str, emb: Any, adaln_lora: Any) -> tuple[Any, Any, Any]: hidden = silu(emb) hidden = linear(hidden, self._weight(f"adaln_modulation_{name}.1.weight")) hidden = linear(hidden, self._weight(f"adaln_modulation_{name}.2.weight")) hidden = hidden + adaln_lora size = hidden.shape[-1] // 3 return hidden[..., :size], hidden[..., size : 2 * size], hidden[..., 2 * size :] @staticmethod def _modulation_shape(x: Any, value: Any) -> Any: return value.reshape(value.shape[0], value.shape[1], *([1] * (len(x.shape) - 3)), value.shape[-1]) def __call__( self, x: Any, emb: Any, context: Any, rope_emb: Any, adaln_lora: Any, *, trace: dict[str, Any] | None = None, trace_prefix: str = "", ) -> Any: residual_dtype = x.dtype b, t_len, h_len, w_len, dim = x.shape shift, scale, gate = self._adaln("self_attn", emb, adaln_lora) if trace is not None: trace[f"{trace_prefix}self_attn_shift"] = shift trace[f"{trace_prefix}self_attn_scale"] = scale trace[f"{trace_prefix}self_attn_gate"] = gate shift = self._modulation_shape(x, shift) scale = self._modulation_shape(x, scale) gate = self._modulation_shape(x, gate) normed = layer_norm(x) * (1 + scale) + shift if trace is not None: trace[f"{trace_prefix}self_attn_normed"] = normed self_attn_input = normed.reshape(b, t_len * h_len * w_len, dim) if trace is not None: trace[f"{trace_prefix}self_attn_input"] = self_attn_input result = self.self_attn( self_attn_input, rope_emb=rope_emb, trace=trace, trace_prefix=f"{trace_prefix}self_attn_", ) if trace is not None: trace[f"{trace_prefix}self_attn_output_flat"] = result result = result.reshape(b, t_len, h_len, w_len, dim) if trace is not None: trace[f"{trace_prefix}self_attn_output_reshaped"] = result x = x + gate.astype(residual_dtype) * result.astype(residual_dtype) if trace is not None: trace[f"{trace_prefix}self_attn_residual_output"] = x shift, scale, gate = self._adaln("cross_attn", emb, adaln_lora) if trace is not None: trace[f"{trace_prefix}cross_attn_shift"] = shift trace[f"{trace_prefix}cross_attn_scale"] = scale trace[f"{trace_prefix}cross_attn_gate"] = gate shift = self._modulation_shape(x, shift) scale = self._modulation_shape(x, scale) gate = self._modulation_shape(x, gate) normed = layer_norm(x) * (1 + scale) + shift if trace is not None: trace[f"{trace_prefix}cross_attn_normed"] = normed cross_attn_input = normed.reshape(b, t_len * h_len * w_len, dim) if trace is not None: trace[f"{trace_prefix}cross_attn_input"] = cross_attn_input result = self.cross_attn( cross_attn_input, context=context, trace=trace, trace_prefix=f"{trace_prefix}cross_attn_", ) if trace is not None: trace[f"{trace_prefix}cross_attn_output_flat"] = result result = result.reshape(b, t_len, h_len, w_len, dim) if trace is not None: trace[f"{trace_prefix}cross_attn_output_reshaped"] = result x = x + gate.astype(residual_dtype) * result.astype(residual_dtype) if trace is not None: trace[f"{trace_prefix}cross_attn_residual_output"] = x shift, scale, gate = self._adaln("mlp", emb, adaln_lora) if trace is not None: trace[f"{trace_prefix}mlp_shift"] = shift trace[f"{trace_prefix}mlp_scale"] = scale trace[f"{trace_prefix}mlp_gate"] = gate shift = self._modulation_shape(x, shift) scale = self._modulation_shape(x, scale) gate = self._modulation_shape(x, gate) normed = layer_norm(x) * (1 + scale) + shift if trace is not None: trace[f"{trace_prefix}mlp_normed"] = normed hidden = linear(normed, self._weight("mlp.layer1.weight")) if trace is not None: trace[f"{trace_prefix}mlp_layer1_output"] = hidden hidden = gelu(hidden) if trace is not None: trace[f"{trace_prefix}mlp_gelu_output"] = hidden result = linear(hidden, self._weight("mlp.layer2.weight")) if trace is not None: trace[f"{trace_prefix}mlp_layer2_output"] = result residual_product = gate.astype(residual_dtype) * result.astype(residual_dtype) output = x + residual_product if trace is not None: trace[f"{trace_prefix}mlp_residual_product"] = residual_product trace[f"{trace_prefix}mlp_residual_output"] = output return output def dit_timesteps(timesteps_b_t: Any, num_channels: int = 2048) -> Any: import math import mlx.core as mx timesteps = timesteps_b_t.reshape(-1).astype(mx.float32) half_dim = num_channels // 2 exponent = -math.log(10000) * mx.arange(half_dim).astype(mx.float32) exponent = exponent / (half_dim - 0.0) emb = mx.exp(exponent) emb = timesteps[:, None] * emb[None, :] emb = mx.concatenate([mx.cos(emb), mx.sin(emb)], axis=-1) return emb.reshape(timesteps_b_t.shape[0], timesteps_b_t.shape[1], num_channels) def patchify_latent(latent: Any) -> Any: import mlx.core as mx b, _, t_len, h_len, w_len = latent.shape pad_h = h_len % 2 pad_w = w_len % 2 if pad_h or pad_w: latent = mx.pad(latent, [(0, 0), (0, 0), (0, 0), (0, pad_h), (0, pad_w)]) h_len = latent.shape[-2] w_len = latent.shape[-1] padding_mask = mx.zeros((b, 1, h_len, w_len), dtype=latent.dtype) if pad_h or pad_w: valid_h = h_len - pad_h valid_w = w_len - pad_w valid = mx.zeros((b, 1, valid_h, valid_w), dtype=latent.dtype) if pad_w: valid = mx.concatenate([valid, mx.ones((b, 1, valid_h, pad_w), dtype=latent.dtype)], axis=-1) if pad_h: valid = mx.concatenate([valid, mx.ones((b, 1, pad_h, w_len), dtype=latent.dtype)], axis=-2) padding_mask = valid patch_input = mx.concatenate([latent, mx.repeat(padding_mask[:, :, None, :, :], t_len, axis=2)], axis=1) b, channels, t_len, h_len, w_len = patch_input.shape return ( patch_input.reshape(b, channels, t_len, h_len // 2, 2, w_len // 2, 2) .transpose(0, 2, 3, 5, 1, 4, 6) .reshape(b, t_len, h_len // 2, w_len // 2, 68) ) def unpatchify(x: Any) -> Any: b, t_len, h_len, w_len, _ = x.shape return ( x.reshape(b, t_len, h_len, w_len, 2, 2, 1, 16) .transpose(0, 7, 1, 6, 2, 4, 3, 5) .reshape(b, 16, t_len, h_len * 2, w_len * 2) ) class DiT: def __init__( self, weights: Mapping[str, Any], *, block_count: int = 28, weights_path: str | Path | None = None, dtype: str = "float32", eval_each_block: bool = True, eval_interval: int | None = None, ) -> None: self.weights = DiTBlock._normalize_keys(weights) self.block_count = block_count self.weights_path = Path(weights_path) if weights_path is not None else None self.dtype = dtype self.eval_interval = 1 if eval_interval is None and eval_each_block else eval_interval or 0 self.blocks = None if self.weights_path is not None else [ DiTBlock(self.weights, prefix=f"blocks.{index}") for index in range(block_count) ] @classmethod def from_safetensors( cls, path: str | Path, *, block_count: int = 28, dtype: str = "float32", lazy_blocks: bool = True, eval_interval: int = 1, ) -> "DiT": from anima_mlx.utils.weights import load_mlx_safetensors_subset if eval_interval < 0: raise ValueError("eval_interval must be non-negative") path = _resolve_diffusion_path(path) core_prefixes = ( "net.x_embedder.", "net.t_embedder.", "net.t_embedding_norm.", "net.final_layer.", ) weights = load_mlx_safetensors_subset( path, strip_prefix="net.", key_filter=lambda key: key.startswith(("net.blocks.", *core_prefixes) if not lazy_blocks else core_prefixes), dtype=dtype, ) if not lazy_blocks and not any(key.startswith("blocks.") for key in weights): for block_index in range(block_count): weights.update(_load_diffusion_block_weights(path, block_index, dtype=dtype)) return cls( weights, block_count=block_count, weights_path=path if lazy_blocks else None, dtype=dtype, eval_interval=eval_interval, ) def _weight(self, name: str) -> Any: return self.weights[name] def prepare_inputs_trace(self, latent: Any, timestep: Any) -> tuple[Any, Any, Any, dict[str, Any]]: patch_input = patchify_latent(latent) x = linear(patch_input, self._weight("x_embedder.proj.1.weight")) t_emb = dit_timesteps(timestep, 2048) t_emb_for_embedder = t_emb.astype(x.dtype) hidden = linear(t_emb_for_embedder, self._weight("t_embedder.1.linear_1.weight")) hidden_activated = silu(hidden) adaln_lora = linear(hidden_activated, self._weight("t_embedder.1.linear_2.weight")) t_embedding = rms_norm(t_emb_for_embedder, self._weight("t_embedding_norm.weight")) return x, t_embedding, adaln_lora, { "patch_input": patch_input, "x_embedder_output": x, "timestep_embedding_raw": t_emb, "t_embedder_linear_1_output": hidden, "t_embedder_silu_output": hidden_activated, "t_embedding": t_embedding, "adaln_lora": adaln_lora, } def prepare_inputs(self, latent: Any, timestep: Any) -> tuple[Any, Any, Any]: x, t_embedding, adaln_lora, _ = self.prepare_inputs_trace(latent, timestep) return x, t_embedding, adaln_lora def final_layer(self, x: Any, emb: Any, adaln_lora: Any, *, trace: dict[str, Any] | None = None) -> Any: hidden = silu(emb) if trace is not None: trace["final_adaln_silu_output"] = hidden hidden = linear(hidden, self._weight("final_layer.adaln_modulation.1.weight")) if trace is not None: trace["final_adaln_linear_1_output"] = hidden hidden = linear(hidden, self._weight("final_layer.adaln_modulation.2.weight")) if trace is not None: trace["final_adaln_linear_2_output"] = hidden hidden = hidden + adaln_lora[:, :, :4096] shift, scale = hidden[..., :2048], hidden[..., 2048:] if trace is not None: trace["final_shift"] = shift trace["final_scale"] = scale shift = DiTBlock._modulation_shape(x, shift) scale = DiTBlock._modulation_shape(x, scale) normed = layer_norm(x) * (1 + scale) + shift if trace is not None: trace["final_normed"] = normed output = linear(normed, self._weight("final_layer.linear.weight")) if trace is not None: trace["final_linear_output"] = output return output def _block(self, index: int) -> DiTBlock: if self.blocks is not None: return self.blocks[index] if self.weights_path is None: raise RuntimeError("DiT block weights are not available") return DiTBlock.from_safetensors(self.weights_path, block_index=index, dtype=self.dtype) def __call__( self, latent: Any, timestep: Any, context: Any, *, checkpoint_indices: set[int] | None = None, trace: bool = False, ) -> Any | tuple[Any, dict[str, Any], Any] | tuple[Any, dict[str, Any]]: checkpoint_indices = checkpoint_indices or set() x, t_embedding, adaln_lora, trace_data = self.prepare_inputs_trace(latent, timestep) rope = video_rope3d(tuple(x.shape), head_dim=128) import mlx.core as mx if trace: trace_data["rope"] = rope rope = mx.expand_dims(mx.expand_dims(rope, axis=1), axis=0) if trace: trace_data["rope_expanded"] = rope checkpoints: dict[str, Any] = {} for index in range(self.block_count): block = self._block(index) if trace and index in {0, 1, 13, 27}: trace_data[f"block_{index}_input"] = x block_trace = trace_data if trace and index in {0, 1} else None block_x = x.astype(mx.float32) if x.dtype == mx.float16 else x block_t_embedding = t_embedding.astype(mx.float32) if t_embedding.dtype == mx.float16 else t_embedding block_adaln_lora = adaln_lora.astype(mx.float32) if adaln_lora.dtype == mx.float16 else adaln_lora x = block( block_x, block_t_embedding, context, rope, block_adaln_lora, trace=block_trace, trace_prefix=f"block_{index}_", ) if self.eval_interval and ((index + 1) % self.eval_interval == 0 or index == self.block_count - 1): mx.eval(x) if trace and index in {0, 1, 13, 27}: trace_data[f"block_{index}_output"] = x if index in checkpoint_indices: checkpoints[f"block_{index}_output"] = x if trace: trace_data["final_layer_input"] = x final_x = x.astype(mx.float32) if x.dtype == mx.float16 else x final_t_embedding = t_embedding.astype(mx.float32) if t_embedding.dtype == mx.float16 else t_embedding final_adaln_lora = adaln_lora.astype(mx.float32) if adaln_lora.dtype == mx.float16 else adaln_lora patch_output = self.final_layer(final_x, final_t_embedding, final_adaln_lora, trace=trace_data if trace else None) unpatchified = unpatchify(patch_output) denoised = unpatchified[:, :, : latent.shape[-3], : latent.shape[-2], : latent.shape[-1]] if trace: trace_data["patch_output"] = patch_output trace_data["unpatchified"] = unpatchified trace_data["denoised"] = denoised return denoised, trace_data if checkpoint_indices: return denoised, checkpoints, patch_output return denoised def _resolve_diffusion_path(path: str | Path) -> Path: resolved = Path(path) if resolved.is_dir(): core_path = resolved / "diffusion_core.safetensors" if core_path.exists(): return core_path return resolved / "diffusion.safetensors" return resolved def _resolve_diffusion_block_path(path: str | Path, block_index: int) -> Path: resolved = Path(path) if resolved.is_dir(): block_path = resolved / "diffusion_blocks" / f"block_{block_index:02d}.safetensors" if block_path.exists(): return block_path return resolved / "diffusion.safetensors" if resolved.name == "diffusion_core.safetensors": block_path = resolved.parent / "diffusion_blocks" / f"block_{block_index:02d}.safetensors" if block_path.exists(): return block_path fallback = resolved.parent / "diffusion.safetensors" if fallback.exists(): return fallback return resolved def _load_diffusion_block_weights(path: str | Path, block_index: int, *, dtype: str) -> dict[str, Any]: from anima_mlx.utils.weights import load_mlx_safetensors_subset block_path = _resolve_diffusion_block_path(path, block_index) prefix = f"net.blocks.{block_index}." return load_mlx_safetensors_subset(block_path, prefix=prefix, strip_prefix="net.", dtype=dtype)