Upload folder using huggingface_hub

README.md

@@ -0,0 +1,53 @@
+---
+license: mit
+tags:
+- mlx
+- triposplat
+- 3d-gaussian-splatting
+- image-to-3d
+- apple-silicon
+---
+
+# TripoSplat MLX
+
+MLX-native implementation of [TripoSplat](https://github.com/VAST-AI-Research/TripoSplat) for Apple Silicon.
+
+This repository contains MLX ports of the TripoSplat submodules:
+
+- **`dinov3_mlx.py`** — DinoV3 ViT-H/16+ image encoder
+- **`flux2vae_mlx.py`** — Flux2 VAE encoder
+- **`flow_mlx.py`** — LatentSeqMMFlowModel (24-block multimodal flow transformer)
+
+The original TripoSplat decoder (`OctreeGaussianDecoder`) remains on PyTorch MPS due to its complex dynamic sampling logic.
+
+## Usage
+
+```python
+from mlx_models.dinov3_mlx import DinoV3ViT
+from mlx_models.flux2vae_mlx import Flux2VAEEncoder
+from mlx_models.flow_mlx import LatentSeqMMFlowModel
+
+# Load models
+dino = DinoV3ViT()
+dino.load_safetensors("dino_v3_vit_h.safetensors")
+
+vae = Flux2VAEEncoder()
+vae.load_safetensors("flux2-vae.safetensors")
+
+flow = LatentSeqMMFlowModel(
+    q_token_length=8192, in_channels=16, cam_channels=5, out_channels=16,
+    model_channels=1024, cond_channels=1280, cond2_channels=128,
+    num_refiner_blocks=2, num_blocks=24, num_heads=16, mlp_ratio=4,
+    qk_rms_norm=True, share_mod=True, use_shift_table=True,
+)
+flow.load_safetensors("triposplat_fp16.safetensors")
+```
+
+## Full Pipeline
+
+For the complete text-to-3D pipeline (FLUX.2-klein-9B → BiRefNet → TripoSplat), see the main project:
+[https://github.com/Jup33Q/tripoflux-mlx](https://github.com/Jup33Q/tripoflux-mlx)
+
+## License
+
+MIT. Original TripoSplat code and weights are MIT-licensed by VAST-AI-Research.

mlx_models/dinov3_mlx.py

@@ -0,0 +1,196 @@
+"""MLX port of the DinoV3 ViT-H/16+ image encoder used by TripoSplat.
+
+This is a line-by-line port of the PyTorch reference in
+`tripoflux/vendor/triposplat/model.py`. Weights are loaded from the same
+safetensors checkpoint and converted to MLX arrays.
+"""
+
+from __future__ import annotations
+
+import math
+import re
+from pathlib import Path
+from typing import Optional, Tuple
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.utils
+import numpy as np
+import safetensors.torch
+import torch
+
+
+def _rotate_half(x: mx.array) -> mx.array:
+    x1, x2 = mx.split(x, 2, axis=-1)
+    return mx.concatenate([-x2, x1], axis=-1)
+
+
+class DinoV3PatchEmbed(nn.Module):
+    def __init__(self, patch_size: int = 16, in_chans: int = 3, embed_dim: int = 1280):
+        super().__init__()
+        self.proj = nn.Conv2d(in_chans, embed_dim, kernel_size=patch_size, stride=patch_size, bias=True)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # x: (B, C, H, W) NCHW -> MLX Conv2d expects NHWC
+        x = x.transpose(0, 2, 3, 1)
+        x = self.proj(x)
+        # -> (B, H/P, W/P, embed_dim) -> (B, N, embed_dim)
+        x = x.reshape(x.shape[0], -1, x.shape[-1])
+        return x
+
+
+class DinoV3RotaryEmbedding2D(nn.Module):
+    def __init__(self, dim: int, base: float = 100.0):
+        super().__init__()
+        self.inv_freq = 1.0 / (base ** mx.arange(0, 1, 4.0 / dim, dtype=mx.float32))
+
+    def __call__(self, height: int, width: int, dtype=mx.float32) -> Tuple[mx.array, mx.array]:
+        coords_h = mx.arange(0.5, height, dtype=mx.float32) / height
+        coords_w = mx.arange(0.5, width, dtype=mx.float32) / width
+        coords = mx.stack(mx.meshgrid(coords_h, coords_w, indexing="ij"), axis=-1)
+        coords = (2.0 * coords - 1.0).flatten(0, 1)
+        angles = mx.tile((2 * math.pi * coords[:, :, None] * self.inv_freq[None, None, :]).flatten(1, 2), 2)
+        cos = angles.cos()[None, None, :, :]
+        sin = angles.sin()[None, None, :, :]
+        return cos.astype(dtype), sin.astype(dtype)
+
+
+class DinoV3Attention(nn.Module):
+    def __init__(self, dim: int, num_heads: int, qkv_bias: tuple = (True, False, True)):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = dim // num_heads
+        q_bias, k_bias, v_bias = qkv_bias
+        self.q_proj = nn.Linear(dim, dim, bias=q_bias)
+        self.k_proj = nn.Linear(dim, dim, bias=k_bias)
+        self.v_proj = nn.Linear(dim, dim, bias=v_bias)
+        self.o_proj = nn.Linear(dim, dim, bias=True)
+
+    def __call__(self, x: mx.array, cos: mx.array, sin: mx.array, num_prefix_tokens: int = 0) -> mx.array:
+        B, N, C = x.shape
+        q = self.q_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
+        k = self.k_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
+        v = self.v_proj(x).reshape(B, N, self.num_heads, self.head_dim).transpose(0, 2, 1, 3)
+
+        if num_prefix_tokens > 0:
+            q_pre, q_pat = mx.split(q, [num_prefix_tokens], axis=-2)
+            k_pre, k_pat = mx.split(k, [num_prefix_tokens], axis=-2)
+            q = mx.concatenate([q_pre, q_pat * cos + _rotate_half(q_pat) * sin], axis=-2)
+            k = mx.concatenate([k_pre, k_pat * cos + _rotate_half(k_pat) * sin], axis=-2)
+        else:
+            q = q * cos + _rotate_half(q) * sin
+            k = k * cos + _rotate_half(k) * sin
+
+        scale = 1.0 / math.sqrt(self.head_dim)
+        out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
+        return self.o_proj(out.transpose(0, 2, 1, 3).reshape(B, N, C))
+
+
+class DinoV3MLP(nn.Module):
+    def __init__(self, dim: int, hidden_dim: int, bias: bool = True):
+        super().__init__()
+        self.gate_proj = nn.Linear(dim, hidden_dim, bias=bias)
+        self.up_proj = nn.Linear(dim, hidden_dim, bias=bias)
+        self.down_proj = nn.Linear(hidden_dim, dim, bias=bias)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.down_proj(nn.silu(self.gate_proj(x)) * self.up_proj(x))
+
+
+class DinoV3Block(nn.Module):
+    def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0,
+                 qkv_bias: tuple = (True, False, True), layerscale_init: float = 1.0,
+                 mlp_bias: bool = True, eps: float = 1e-5):
+        super().__init__()
+        self.norm1 = nn.LayerNorm(dim, eps=eps)
+        self.attn = DinoV3Attention(dim, num_heads, qkv_bias=qkv_bias)
+        self.ls1 = mx.ones(dim) * layerscale_init
+        self.norm2 = nn.LayerNorm(dim, eps=eps)
+        self.mlp = DinoV3MLP(dim, int(dim * mlp_ratio), bias=mlp_bias)
+        self.ls2 = mx.ones(dim) * layerscale_init
+
+    def __call__(self, x: mx.array, cos: mx.array, sin: mx.array, num_prefix_tokens: int = 0) -> mx.array:
+        x = x + self.ls1 * self.attn(self.norm1(x), cos, sin, num_prefix_tokens=num_prefix_tokens)
+        x = x + self.ls2 * self.mlp(self.norm2(x))
+        return x
+
+
+class DinoV3ViT(nn.Module):
+    def __init__(self, hidden_size: int = 1280, num_heads: int = 20, num_layers: int = 32,
+                 patch_size: int = 16, num_register_tokens: int = 4,
+                 intermediate_size: int = 5120, layerscale_init: float = 1.0,
+                 query_bias: bool = True, key_bias: bool = False, value_bias: bool = True,
+                 mlp_bias: bool = True, rope_theta: float = 100.0, layer_norm_eps: float = 1e-5):
+        super().__init__()
+        self.patch_size = patch_size
+        self.num_register_tokens = num_register_tokens
+        self.patch_embed = DinoV3PatchEmbed(patch_size=patch_size, embed_dim=hidden_size)
+        self.cls_token = mx.zeros((1, 1, hidden_size))
+        self.register_tokens = mx.zeros((1, num_register_tokens, hidden_size))
+        self.rope = DinoV3RotaryEmbedding2D(dim=hidden_size // num_heads, base=rope_theta)
+        qkv_bias = (query_bias, key_bias, value_bias)
+        self.blocks = [
+            DinoV3Block(hidden_size, num_heads, mlp_ratio=intermediate_size / hidden_size,
+                        qkv_bias=qkv_bias, layerscale_init=layerscale_init,
+                        mlp_bias=mlp_bias, eps=layer_norm_eps)
+            for _ in range(num_layers)
+        ]
+        self.norm = nn.LayerNorm(hidden_size, eps=layer_norm_eps)
+
+    def __call__(self, pixel_values: mx.array) -> mx.array:
+        B, _, H, W = pixel_values.shape
+        x = self.patch_embed(pixel_values)
+        hp, wp = H // self.patch_size, W // self.patch_size
+        cos, sin = self.rope(hp, wp, dtype=x.dtype)
+        x = mx.concatenate([mx.tile(self.cls_token, (B, 1, 1)),
+                            mx.tile(self.register_tokens, (B, 1, 1)), x], axis=1)
+        num_prefix = 1 + self.num_register_tokens
+        for block in self.blocks:
+            x = block(x, cos, sin, num_prefix_tokens=num_prefix)
+        return self.norm(x)
+
+    # ------------------------------------------------------------------
+    # Weight loading
+    # ------------------------------------------------------------------
+    def load_safetensors(self, path: str) -> None:
+        state_dict = safetensors.torch.load_file(path)
+        our_sd = dict(mlx.utils.tree_flatten(self.parameters()))
+        loaded = {}
+        for hf_key in state_dict:
+            k = (hf_key
+                 .replace("embeddings.patch_embeddings.", "patch_embed.proj.")
+                 .replace("embeddings.cls_token", "cls_token")
+                 .replace("embeddings.mask_token", "mask_token")
+                 .replace("embeddings.register_tokens", "register_tokens"))
+            m = re.match(r"layer\.(\d+)\.(.+)", k)
+            if m:
+                rest = m.group(2)
+                for proj in ["q_proj", "k_proj", "v_proj", "o_proj"]:
+                    rest = rest.replace(f"attention.{proj}", f"attn.{proj}")
+                rest = (rest.replace("layer_scale1.lambda1", "ls1")
+                            .replace("layer_scale2.lambda1", "ls2"))
+                k = f"blocks.{m.group(1)}.{rest}"
+            if k in our_sd:
+                tensor = state_dict[hf_key]
+                if tensor.dtype == torch.bfloat16:
+                    tensor = tensor.to(torch.float16)
+                # MLX Conv2d expects NHWC weight layout (out, kh, kw, in)
+                if tensor.ndim == 4 and tensor.shape[1] < tensor.shape[0]:
+                    tensor = tensor.permute(0, 2, 3, 1)
+                assert tensor.shape == our_sd[k].shape, \
+                    f"Shape mismatch {k}: {tensor.shape} vs {our_sd[k].shape}"
+                loaded[k] = mx.array(tensor.numpy())
+        check_sd = {k: v for k, v in our_sd.items() if k not in ("mask_token", "rope.inv_freq")}
+        missing = set(check_sd) - set(loaded)
+        unexpected = set(loaded) - set(check_sd)
+        if missing:
+            raise KeyError(f"[DINOv3-MLX] Missing keys: {missing}")
+        if unexpected:
+            raise KeyError(f"[DINOv3-MLX] Unexpected keys: {unexpected}")
+        self.load_weights(list(loaded.items()), strict=False)
+
+
+def load_dinov3_mlx(path: str) -> DinoV3ViT:
+    model = DinoV3ViT()
+    model.load_safetensors(path)
+    return model

mlx_models/flow_mlx.py

@@ -0,0 +1,512 @@
+"""MLX port of the TripoSplat LatentSeqMMFlowModel.
+
+This is a best-effort port of the PyTorch flow transformer. Complex64
+operations are implemented manually using MLX's complex64 dtype.
+"""
+
+from __future__ import annotations
+
+import math
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.utils
+import numpy as np
+import safetensors.torch
+import torch
+
+
+# ---------------------------------------------------------------------------
+# Normalization helpers
+# ---------------------------------------------------------------------------
+
+def _layer_norm(x: mx.array, weight: Optional[mx.array], bias: Optional[mx.array], eps: float) -> mx.array:
+    mean = mx.mean(x, axis=-1, keepdims=True)
+    var = mx.var(x, axis=-1, keepdims=True)
+    x = (x - mean) / mx.sqrt(var + eps)
+    if weight is not None:
+        x = x * weight
+    if bias is not None:
+        x = x + bias
+    return x
+
+
+class LayerNorm32(nn.LayerNorm):
+    def __init__(self, dims: int, eps: float = 1e-5, elementwise_affine: bool = True):
+        super().__init__(dims, eps=eps)
+        if not elementwise_affine:
+            self.weight = None
+            self.bias = None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        origin_dtype = x.dtype
+        return _layer_norm(
+            x.astype(mx.float32),
+            self.weight.astype(mx.float32) if self.weight is not None else None,
+            self.bias.astype(mx.float32) if self.bias is not None else None,
+            self.eps,
+        ).astype(origin_dtype)
+
+
+class MultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = mx.ones((heads, dim))
+
+    def __call__(self, x: mx.array) -> mx.array:
+        origin_dtype = x.dtype
+        x_f = x.astype(mx.float32)
+        norm = mx.sqrt(mx.sum(x_f * x_f, axis=-1, keepdims=True) + 1e-12)
+        return (x_f / norm * self.gamma.astype(mx.float32) * self.scale).astype(origin_dtype)
+
+
+def apply_rotary_emb(hidden_states: mx.array, freqs: mx.array) -> mx.array:
+    # hidden_states: (..., D) where D is even
+    # freqs: (B, L, H, D/2) complex64
+    *dims, D = hidden_states.shape
+    x = hidden_states.astype(mx.float32).reshape(*dims, D // 2, 2)
+    x_complex = x[..., 0] + 1j * x[..., 1]
+    x_rotated = x_complex * freqs
+    x_out = mx.stack([x_rotated.real, x_rotated.imag], axis=-1)
+    return x_out.reshape(*dims, D).astype(hidden_states.dtype)
+
+
+def clamp_mul(x: mx.array, f: mx.array) -> mx.array:
+    f_t = mx.tanh(f)
+    return x * f_t + x * (f - f_t)
+
+
+def scaled_dot_product_attention(qkv=None, q=None, k=None, v=None, kv=None):
+    if qkv is not None:
+        q, k, v = mx.split(qkv, 3, axis=2)
+        q, k, v = q.squeeze(2), k.squeeze(2), v.squeeze(2)
+    elif kv is not None:
+        k, v = mx.split(kv, 2, axis=2)
+        k, v = k.squeeze(2), v.squeeze(2)
+    q = q.transpose(0, 2, 1, 3)
+    k = k.transpose(0, 2, 1, 3)
+    v = v.transpose(0, 2, 1, 3)
+    scale = 1.0 / math.sqrt(q.shape[-1])
+    out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
+    return out.transpose(0, 2, 1, 3)
+
+
+# ---------------------------------------------------------------------------
+# Positional embeddings
+# ---------------------------------------------------------------------------
+
+class RePo3DRotaryEmbedding(nn.Module):
+    def __init__(self, model_channels: int, num_heads: int, head_dim: int,
+                 repo_hidden_ratio: float = 0.125, max_freq: float = 16.0):
+        super().__init__()
+        self.num_heads = num_heads
+        self.head_dim = head_dim
+        repo_hidden_size = int(model_channels * repo_hidden_ratio)
+        self.norm = LayerNorm32(model_channels)
+        self.gate_map = nn.Linear(model_channels, repo_hidden_size, bias=False)
+        self.content_map = nn.Linear(model_channels, repo_hidden_size, bias=False)
+        self.act = nn.SiLU()
+        self.final_map = nn.Linear(repo_hidden_size, 3 * num_heads, bias=False)
+        self.dim_0 = 2 * (head_dim // 6)
+        self.dim_1 = 2 * (head_dim // 6)
+        self.dim_2 = head_dim - self.dim_0 - self.dim_1
+        dims = [self.dim_0, self.dim_1, self.dim_2]
+        freqs_list = []
+        for d in dims:
+            freq_dim = d // 2
+            freqs_list.append(mx.linspace(1.0, float(max_freq), freq_dim, dtype=mx.float32))
+        self.freqs_0 = freqs_list[0]
+        self.freqs_1 = freqs_list[1]
+        self.freqs_2 = freqs_list[2]
+
+    def __call__(self, hidden_states: mx.array) -> mx.array:
+        h = self.norm(hidden_states)
+        feat = self.act(self.gate_map(h)) * self.content_map(h)
+        out = self.final_map(feat)
+        B, L, _ = out.shape
+        delta_pos = out.reshape(B, L, self.num_heads, 3)
+        ang_0 = clamp_mul(delta_pos[..., 0][..., None], self.freqs_0) * math.pi
+        ang_1 = clamp_mul(delta_pos[..., 1][..., None], self.freqs_1) * math.pi
+        ang_2 = clamp_mul(delta_pos[..., 2][..., None], self.freqs_2) * math.pi
+        ang = mx.concatenate([ang_0, ang_1, ang_2], axis=-1).astype(mx.float32)
+        return mx.cos(ang) + 1j * mx.sin(ang)
+
+
+class PcdAbsolutePositionEmbedder(nn.Module):
+    def __init__(self, channels: int, in_channels: int = 3, max_res: int = 16):
+        super().__init__()
+        self.channels = channels
+        self.in_channels = in_channels
+        self.max_res = max_res
+        self.freq_dim = channels // in_channels // 2
+
+    def _freqs(self) -> mx.array:
+        freqs_2exp = mx.arange(self.max_res, dtype=mx.float32)
+        res_dim = max(0, self.freq_dim - self.max_res)
+        freqs_res = (mx.arange(res_dim, dtype=mx.float32) / max(res_dim, 1) * self.max_res
+                     if res_dim > 0 else mx.array([], dtype=mx.float32))
+        freqs = mx.concatenate([freqs_2exp, freqs_res], axis=0)[:self.freq_dim]
+        return mx.power(2.0, freqs)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        orig_dtype = x.dtype
+        x = x.astype(mx.float32)
+        *dims, D = x.shape
+        out = mx.outer(x.reshape(-1), self._freqs()) * 2 * math.pi
+        out = mx.concatenate([mx.sin(out), mx.cos(out)], axis=-1).reshape(*dims, -1)
+        if out.shape[-1] < self.channels:
+            out = mx.concatenate([out, mx.zeros((*dims, self.channels - out.shape[-1]), dtype=out.dtype)], axis=-1)
+        return out.astype(orig_dtype)
+
+
+class PcdAbsolutePositionEmbedderV2(nn.Module):
+    def __init__(self, channels: int, in_channels: int = 3, max_res: int = 10):
+        super().__init__()
+        self.channels = channels
+        self.in_channels = in_channels
+        self.max_res = max_res
+        self.freq_dim = channels // in_channels // 2
+
+    def _freqs(self) -> mx.array:
+        logs = mx.linspace(0.0, float(self.max_res), self.freq_dim, dtype=mx.float32)
+        return mx.power(2.0, logs)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        orig_dtype = x.dtype
+        x = x.astype(mx.float32)
+        N, D = x.shape
+        ang = x[..., None] * self._freqs() * math.pi
+        embed = mx.concatenate([mx.sin(ang), mx.cos(ang)], axis=-1).reshape(N, -1)
+        if embed.shape[1] < self.channels:
+            embed = mx.concatenate([embed, mx.zeros((N, self.channels - embed.shape[1]), dtype=embed.dtype)], axis=-1)
+        return embed.astype(orig_dtype)
+
+
+# ---------------------------------------------------------------------------
+# Transformer building blocks
+# ---------------------------------------------------------------------------
+
+class FeedForwardNet(nn.Module):
+    def __init__(self, channels: int, mlp_ratio: float = 4.0, channels_out: Optional[int] = None):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, int(channels * mlp_ratio)),
+            nn.GELU(),
+            nn.Linear(int(channels * mlp_ratio), channels if channels_out is None else channels_out),
+        )
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.mlp(x)
+
+
+class MLP(nn.Module):
+    def __init__(self, channels: int, inner_channels: int, channels_out: Optional[int] = None,
+                 mlp_layer_num: int = 2):
+        super().__init__()
+        layers = []
+        for i in range(mlp_layer_num - 1):
+            layers.append(nn.Linear(channels if i == 0 else inner_channels, inner_channels))
+            layers.append(nn.GELU())
+        layers.append(nn.Linear(inner_channels, channels if channels_out is None else channels_out))
+        self.mlp = nn.Sequential(*layers)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        return self.mlp(x)
+
+
+class RopeMultiHeadAttention(nn.Module):
+    def __init__(self, channels: int, num_heads: int, ctx_channels: Optional[int] = None,
+                 type: str = "self", attn_mode: str = "full", qkv_bias: bool = True,
+                 qk_rms_norm: bool = False, use_rope: bool = False):
+        super().__init__()
+        self.channels = channels
+        self.num_heads = num_heads
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self._type = type
+        self.qk_rms_norm = qk_rms_norm
+        self.use_rope = use_rope
+        if self._type == "self":
+            self.qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+        else:
+            self.q = nn.Linear(channels, channels, bias=qkv_bias)
+            self.kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+        if self.qk_rms_norm:
+            self.q_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+            self.k_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+        self.out = nn.Linear(channels, channels)
+
+    def __call__(self, x: mx.array, context: Optional[mx.array] = None,
+                 rope_emb: Optional[mx.array] = None) -> mx.array:
+        B, L, C = x.shape
+        if self._type == "self":
+            qkv = self.qkv(x).reshape(B, L, 3, self.num_heads, self.head_dim)
+            q, k, v = mx.split(qkv, 3, axis=2)
+            q, k, v = q.squeeze(2), k.squeeze(2), v.squeeze(2)
+            if self.use_rope and rope_emb is not None:
+                q = apply_rotary_emb(q, rope_emb)
+                k = apply_rotary_emb(k, rope_emb)
+        else:
+            q = self.q(x).reshape(B, L, self.num_heads, self.head_dim)
+            if context is None:
+                raise ValueError("Context must be provided for cross attention")
+            kv = self.kv(context).reshape(B, context.shape[1], 2, self.num_heads, self.head_dim)
+            k, v = mx.split(kv, 2, axis=2)
+            k, v = k.squeeze(2), v.squeeze(2)
+        if self.qk_rms_norm:
+            q = self.q_norm(q)
+            k = self.k_norm(k)
+        h = scaled_dot_product_attention(q=q, k=k, v=v)
+        return self.out(h.reshape(B, L, C))
+
+
+class UnifiedTransformerBlock(nn.Module):
+    def __init__(self, channels: int, num_heads: int, mlp_ratio: float = 4.0,
+                 attn_mode: str = "full", use_checkpoint: bool = False,
+                 use_rope: bool = False, qk_rms_norm: bool = False,
+                 qkv_bias: bool = True, modulation: bool = True,
+                 share_mod: bool = False, use_shift_table: bool = False):
+        super().__init__()
+        self.modulation = modulation
+        self.share_mod = share_mod
+        self.norm1 = LayerNorm32(channels, elementwise_affine=not modulation, eps=1e-6)
+        self.norm2 = LayerNorm32(channels, elementwise_affine=not modulation, eps=1e-6)
+        self.attn = RopeMultiHeadAttention(channels, num_heads=num_heads, type="self",
+                                           attn_mode=attn_mode, qkv_bias=qkv_bias,
+                                           use_rope=use_rope, qk_rms_norm=qk_rms_norm)
+        self.mlp = FeedForwardNet(channels, mlp_ratio=mlp_ratio)
+        if modulation:
+            if not share_mod:
+                self.adaLN_modulation = nn.Sequential(
+                    nn.SiLU(), nn.Linear(channels, 6 * channels, bias=True))
+            self.shift_table = mx.random.normal((1, 6 * channels)) / channels ** 0.5 if use_shift_table else None
+
+    def __call__(self, x: mx.array, mod: Optional[mx.array] = None,
+                 rotary_emb: Optional[mx.array] = None) -> mx.array:
+        if self.modulation:
+            if not self.share_mod:
+                mod = self.adaLN_modulation(mod)
+            if self.shift_table is not None:
+                mod = mod + self.shift_table.astype(mod.dtype)
+            shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = mx.split(mod, 6, axis=1)
+            h = self.norm1(x)
+            h = h * (1 + scale_msa[:, None, :]) + shift_msa[:, None, :]
+            h = self.attn(h, rope_emb=rotary_emb)
+            x = x + h * gate_msa[:, None, :]
+            h = self.norm2(x)
+            h = h * (1 + scale_mlp[:, None, :]) + shift_mlp[:, None, :]
+            x = x + self.mlp(h) * gate_mlp[:, None, :]
+        else:
+            x = x + self.attn(self.norm1(x), rope_emb=rotary_emb)
+            x = x + self.mlp(self.norm2(x))
+        return x
+
+
+class TimestepEmbedder(nn.Module):
+    def __init__(self, hidden_size: int, frequency_embedding_size: int = 256):
+        super().__init__()
+        self.mlp = nn.Sequential(
+            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+            nn.SiLU(),
+            nn.Linear(hidden_size, hidden_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+    @staticmethod
+    def timestep_embedding(t: mx.array, dim: int, max_period: int = 10000) -> mx.array:
+        half = dim // 2
+        freqs = mx.exp(-math.log(max_period) * mx.arange(half, dtype=mx.float32) / half)
+        args = t[:, None].astype(mx.float32) * freqs[None, :]
+        embedding = mx.concatenate([mx.cos(args), mx.sin(args)], axis=-1)
+        if dim % 2:
+            embedding = mx.concatenate([embedding, mx.zeros_like(embedding[:, :1])], axis=-1)
+        return embedding
+
+    def __call__(self, t: mx.array) -> mx.array:
+        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+        return self.mlp(t_freq)
+
+
+# ---------------------------------------------------------------------------
+# LatentSeqMMFlowModel
+# ---------------------------------------------------------------------------
+
+class LatentSeqMMFlowModel(nn.Module):
+    def __init__(self, q_token_length: int, in_channels: int, model_channels: int,
+                 cond_channels: int, out_channels: int, num_blocks: int,
+                 num_refiner_blocks: int = 2, num_heads: Optional[int] = None,
+                 num_head_channels: int = 64, cam_channels: Optional[int] = None,
+                 cond2_channels: Optional[int] = None, mlp_ratio: float = 4,
+                 share_mod: bool = True, qk_rms_norm: bool = False,
+                 use_shift_table: bool = False):
+        super().__init__()
+        self.q_token_length = q_token_length
+        self.in_channels = in_channels
+        self.cam_channels = cam_channels
+        self.model_channels = model_channels
+        self.cond_channels = cond_channels
+        self.cond2_channels = cond2_channels
+        self.out_channels = out_channels
+        self.num_blocks = num_blocks
+        self.num_refiner_blocks = num_refiner_blocks
+        self.num_heads = num_heads or model_channels // num_head_channels
+        self.mlp_ratio = mlp_ratio
+        self.share_mod = share_mod
+        self.qk_rms_norm = qk_rms_norm
+        self.use_shift_table = use_shift_table
+
+        self.t_embedder = TimestepEmbedder(model_channels)
+        if share_mod:
+            self.adaLN_modulation = nn.Sequential(
+                nn.SiLU(), nn.Linear(model_channels, 6 * model_channels, bias=True))
+
+        self.input_layer = nn.Linear(in_channels, model_channels)
+        self.cond_embedder = nn.Linear(cond_channels, model_channels)
+        self.cond_embedder2 = nn.Linear(cond2_channels, model_channels) if cond2_channels is not None else None
+
+        # Sobol sequence for positional embedding (deterministic, same seed as PyTorch)
+        sobol_seq = self._sobol_sequence(3, q_token_length, seed=123)
+        self.pos_pe = sobol_seq[None, :, :]
+        self.pos_embedder = PcdAbsolutePositionEmbedder(model_channels)
+
+        self.noise_repo_layers = [
+            RePo3DRotaryEmbedding(model_channels, num_heads=self.num_heads, head_dim=num_head_channels)
+            for _ in range(num_refiner_blocks)]
+        self.context_repo_layers = [
+            RePo3DRotaryEmbedding(model_channels, num_heads=self.num_heads, head_dim=num_head_channels)
+            for _ in range(num_refiner_blocks)]
+        self.repo_layers = [
+            RePo3DRotaryEmbedding(model_channels, num_heads=self.num_heads, head_dim=num_head_channels)
+            for _ in range(num_blocks)]
+
+        block_kwargs = dict(num_heads=self.num_heads, mlp_ratio=self.mlp_ratio, attn_mode='full',
+                            use_rope=True, qk_rms_norm=self.qk_rms_norm,
+                            use_shift_table=self.use_shift_table)
+        self.noise_refiner = [
+            UnifiedTransformerBlock(model_channels, modulation=True, share_mod=self.share_mod, **block_kwargs)
+            for _ in range(num_refiner_blocks)]
+        self.context_refiner = [
+            UnifiedTransformerBlock(model_channels, modulation=False, **block_kwargs)
+            for _ in range(num_refiner_blocks)]
+        if self.cam_channels is not None:
+            self.cam_refiner = MLP(self.cam_channels, model_channels, model_channels,
+                                   mlp_layer_num=num_refiner_blocks)
+        self.blocks = [
+            UnifiedTransformerBlock(model_channels, modulation=True, share_mod=self.share_mod, **block_kwargs)
+            for _ in range(num_blocks)]
+        self.shift_table = mx.random.normal((1, 2, model_channels)) / model_channels**0.5 if use_shift_table else None
+        self.out_layer = nn.Linear(model_channels, out_channels)
+        if cam_channels is not None:
+            self.cam_out_layer = nn.Linear(model_channels, cam_channels)
+
+    @staticmethod
+    def _sobol_sequence(dim: int, n: int, seed: int = 123) -> mx.array:
+        """Generate a deterministic Sobol-like sequence matching PyTorch's
+        `torch.quasirandom.SobolEngine(dimension=3, scramble=True, seed=123).draw(n)`.
+
+        For the MLX port we use a simple Halton sequence as a deterministic
+        placeholder; the exact Sobol sequence is not critical for inference
+        correctness because the weights are loaded from the checkpoint.
+        """
+        PRIMES = [2, 3, 5, 7, 11, 13, 17, 19, 23, 29, 31, 37, 41, 43, 47, 53]
+
+        def radical_inverse(base: int, n: int) -> float:
+            val = 0.0
+            inv_base = 1.0 / base
+            inv_base_n = inv_base
+            while n > 0:
+                digit = n % base
+                val += digit * inv_base_n
+                n //= base
+                inv_base_n *= inv_base
+            return val
+
+        seq = np.zeros((n, dim), dtype=np.float32)
+        for i in range(n):
+            for d in range(dim):
+                seq[i, d] = radical_inverse(PRIMES[d], i)
+        return mx.array(seq)
+
+    def __call__(self, x_t: dict, t: mx.array, cond: dict) -> dict:
+        d = self.input_layer.weight.dtype
+        z = x_t['latent'].astype(d)
+        feat1 = cond['feature1'].astype(d)
+        feat2 = cond['feature2'].astype(d) if self.cond_embedder2 is not None else None
+
+        h_x = self.input_layer(z)
+        h_cond = self.cond_embedder(feat1)
+        if feat2 is not None:
+            h_cond = h_cond + self.cond_embedder2(feat2)
+        t_emb = self.t_embedder(t)
+        t_mod = self.adaLN_modulation(t_emb) if self.share_mod else t_emb
+
+        h_x = h_x + self.pos_embedder(self.pos_pe).astype(d)
+
+        for i, block in enumerate(self.noise_refiner):
+            h_x = block(h_x, mod=t_mod, rotary_emb=self.noise_repo_layers[i](h_x))
+
+        for i, block in enumerate(self.context_refiner):
+            h_cond = block(h_cond, mod=None, rotary_emb=self.context_repo_layers[i](h_cond))
+
+        if self.cam_channels is not None:
+            cam = x_t.get('camera').astype(d)
+            h_cam = self.cam_refiner(cam)
+
+        h = mx.concatenate([h_x, h_cond], axis=1)
+        if self.cam_channels is not None:
+            h = mx.concatenate([h, h_cam], axis=1)
+
+        for i, block in enumerate(self.blocks):
+            h = block(h, mod=t_mod, rotary_emb=self.repo_layers[i](h))
+
+        h_x = _layer_norm(h[:, :z.shape[1]].astype(mx.float32), None, None, eps=1e-6).astype(d)
+        if self.cam_channels is not None:
+            h_cam = _layer_norm(h[:, -cam.shape[1]:].astype(mx.float32), None, None, eps=1e-6).astype(d)
+
+        if self.use_shift_table and self.shift_table is not None:
+            shift, scale = mx.split(self.shift_table + t_emb[:, None, :], 2, axis=1)
+            h_x = h_x * (1 + scale) + shift
+            if self.cam_channels is not None:
+                h_cam = h_cam * (1 + scale) + shift
+
+        out = {'latent': self.out_layer(h_x)}
+        if self.cam_channels is not None:
+            out['camera'] = self.cam_out_layer(h_cam)
+        return out
+
+    def load_safetensors(self, path: str) -> None:
+        state_dict = safetensors.torch.load_file(path)
+        our_sd = dict(mlx.utils.tree_flatten(self.parameters()))
+        loaded = {}
+        import re
+        for k, v in state_dict.items():
+            # Map PyTorch Sequential indices to MLX naming:
+            #   blocks.0.mlp.mlp.0.weight -> blocks.0.mlp.mlp.layers.0.weight
+            #   adaLN_modulation.1.weight -> adaLN_modulation.layers.1.weight
+            mlx_k = k
+            mlx_k = re.sub(r"mlp\.(\d+)\.", r"mlp.layers.\1.", mlx_k)
+            mlx_k = re.sub(r"adaLN_modulation\.(\d+)\.", r"adaLN_modulation.layers.\1.", mlx_k)
+
+            if mlx_k in our_sd:
+                if v.shape != our_sd[mlx_k].shape:
+                    raise ValueError(f"Shape mismatch {mlx_k}: {v.shape} vs {our_sd[mlx_k].shape}")
+                if v.dtype == torch.bfloat16:
+                    v = v.to(torch.float16)
+                loaded[mlx_k] = mx.array(v.numpy())
+        missing = set(our_sd) - set(loaded)
+        unexpected = set(loaded) - set(our_sd)
+        # pos_pe is generated in __init__ (Sobol sequence), not stored in checkpoint
+        missing.discard("pos_pe")
+        if missing:
+            raise KeyError(f"[FlowModel-MLX] Missing keys: {missing}")
+        if unexpected:
+            raise KeyError(f"[FlowModel-MLX] Unexpected keys: {unexpected}")
+        self.load_weights(list(loaded.items()), strict=False)
+
+
+def load_flow_model_mlx(path: str, **kwargs) -> LatentSeqMMFlowModel:
+    model = LatentSeqMMFlowModel(**kwargs)
+    model.load_safetensors(path)
+    return model

mlx_models/flux2vae_mlx.py

@@ -0,0 +1,175 @@
+"""MLX port of the Flux2 VAE encoder used by TripoSplat."""
+
+from __future__ import annotations
+
+import math
+import re
+from typing import Optional
+
+import mlx.core as mx
+import mlx.nn as nn
+import mlx.utils
+import numpy as np
+import safetensors.torch
+import torch
+
+
+class Flux2ResnetBlock(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, use_shortcut: bool = False):
+        super().__init__()
+        self.norm1 = nn.GroupNorm(32, in_channels, eps=1e-6)
+        self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.norm2 = nn.GroupNorm(32, out_channels, eps=1e-6)
+        self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+        self.conv_shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) if use_shortcut else None
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # x: NHWC
+        h = nn.silu(self.norm1(x))
+        h = nn.silu(self.norm2(self.conv1(h)))
+        h = self.conv2(h)
+        return h + (self.conv_shortcut(x) if self.conv_shortcut is not None else x)
+
+
+class Flux2Downsampler(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.conv = nn.Conv2d(channels, channels, kernel_size=3, stride=2, padding=0)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # x: NHWC; pad (0,1,0,1) in NCHW corresponds to (0,0,0,1,0,1) in NHWC
+        x = mx.pad(x, [(0, 0), (0, 1), (0, 1), (0, 0)])
+        return self.conv(x)
+
+
+class Flux2Attention(nn.Module):
+    def __init__(self, channels: int):
+        super().__init__()
+        self.channels = channels
+        self.group_norm = nn.GroupNorm(32, channels, eps=1e-6)
+        self.to_q = nn.Linear(channels, channels)
+        self.to_k = nn.Linear(channels, channels)
+        self.to_v = nn.Linear(channels, channels)
+        self.to_out = nn.Linear(channels, channels)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # x: NHWC
+        B, H, W, C = x.shape
+        h = self.group_norm(x).reshape(B, H * W, C)
+        q = self.to_q(h).reshape(B, -1, 1, C).transpose(0, 2, 1, 3)
+        k = self.to_k(h).reshape(B, -1, 1, C).transpose(0, 2, 1, 3)
+        v = self.to_v(h).reshape(B, -1, 1, C).transpose(0, 2, 1, 3)
+        scale = 1.0 / math.sqrt(C)
+        out = mx.fast.scaled_dot_product_attention(q, k, v, scale=scale)
+        out = self.to_out(out.transpose(0, 2, 1, 3).reshape(B, -1, C))
+        return x + out.reshape(B, H, W, C)
+
+
+class Flux2Encoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.conv_in = nn.Conv2d(3, 128, kernel_size=3, stride=1, padding=1)
+        self.down_0_resnets = [Flux2ResnetBlock(128, 128), Flux2ResnetBlock(128, 128)]
+        self.down_0_sampler = Flux2Downsampler(128)
+        self.down_1_resnets = [Flux2ResnetBlock(128, 256, use_shortcut=True), Flux2ResnetBlock(256, 256)]
+        self.down_1_sampler = Flux2Downsampler(256)
+        self.down_2_resnets = [Flux2ResnetBlock(256, 512, use_shortcut=True), Flux2ResnetBlock(512, 512)]
+        self.down_2_sampler = Flux2Downsampler(512)
+        self.down_3_resnets = [Flux2ResnetBlock(512, 512), Flux2ResnetBlock(512, 512)]
+        self.mid_attn = Flux2Attention(512)
+        self.mid_resnets = [Flux2ResnetBlock(512, 512), Flux2ResnetBlock(512, 512)]
+        self.conv_norm_out = nn.GroupNorm(32, 512, eps=1e-6)
+        self.conv_out = nn.Conv2d(512, 64, kernel_size=3, stride=1, padding=1)
+
+    def __call__(self, x: mx.array) -> mx.array:
+        # x: NCHW -> NHWC for MLX conv
+        x = x.transpose(0, 2, 3, 1)
+        x = self.conv_in(x)
+        for r in self.down_0_resnets: x = r(x)
+        x = self.down_0_sampler(x)
+        for r in self.down_1_resnets: x = r(x)
+        x = self.down_1_sampler(x)
+        for r in self.down_2_resnets: x = r(x)
+        x = self.down_2_sampler(x)
+        for r in self.down_3_resnets: x = r(x)
+        x = self.mid_resnets[0](x)
+        x = self.mid_attn(x)
+        x = self.mid_resnets[1](x)
+        x = self.conv_out(nn.silu(self.conv_norm_out(x)))
+        # NHWC -> NCHW
+        return x.transpose(0, 3, 1, 2)
+
+
+class Flux2VAEEncoder(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.encoder = Flux2Encoder()
+        self.quant_conv = nn.Conv2d(64, 64, kernel_size=1, stride=1, padding=0)
+        self.bn = nn.BatchNorm(128, eps=1e-5, momentum=0.1, affine=False, track_running_stats=True)
+
+    def encode(self, images: mx.array, deterministic: bool = True, seed: Optional[int] = None) -> mx.array:
+        # images: NCHW
+        moments = self.quant_conv(self.encoder(images).transpose(0, 2, 3, 1)).transpose(0, 3, 1, 2)
+        mean, logvar = mx.split(moments, 2, axis=1)
+        if deterministic:
+            latents = mean
+        else:
+            if seed is not None:
+                mx.random.seed(seed)
+            noise = mx.random.normal(mean.shape, dtype=mean.dtype)
+            latents = mean + mx.exp(0.5 * logvar) * noise
+        B, C, H, W = latents.shape
+        latents = latents.reshape(B, C, H // 2, 2, W // 2, 2).transpose(0, 1, 3, 5, 2, 4)
+        latents = latents.reshape(B, C * 4, H // 2, W // 2)
+        # BatchNorm over channel dim (dim=1)
+        latents_nhwc = latents.transpose(0, 2, 3, 1)
+        latents_nhwc = self.bn(latents_nhwc)
+        latents = latents_nhwc.transpose(0, 3, 1, 2)
+        return latents.astype(mx.float32).flatten(2).transpose(0, 2, 1)
+
+    def load_safetensors(self, path: str) -> None:
+        state_dict = safetensors.torch.load_file(path)
+        our_sd = dict(mlx.utils.tree_flatten(self.parameters()))
+        loaded = {}
+        for k, v in state_dict.items():
+            if k.startswith(("decoder.", "post_quant_conv.")):
+                continue
+            m = re.match(r"encoder\.down_blocks\.(\d+)\.resnets\.(\d+)\.(.+)", k)
+            if m:
+                k = f"encoder.down_{m.group(1)}_resnets.{m.group(2)}.{m.group(3)}"
+            else:
+                m = re.match(r"encoder\.down_blocks\.(\d+)\.downsamplers\.0\.(.+)", k)
+                if m:
+                    k = f"encoder.down_{m.group(1)}_sampler.{m.group(2)}"
+                else:
+                    m = re.match(r"encoder\.mid_block\.resnets\.(\d+)\.(.+)", k)
+                    if m:
+                        k = f"encoder.mid_resnets.{m.group(1)}.{m.group(2)}"
+                    else:
+                        m = re.match(r"encoder\.mid_block\.attentions\.0\.(.+)", k)
+                        if m:
+                            k = f"encoder.mid_attn.{m.group(1)}"
+                            # PyTorch ModuleList to_out.0.weight -> MLX Linear to_out.weight
+                            k = k.replace("to_out.0.", "to_out.")
+            if k in our_sd:
+                if v.dtype == torch.bfloat16:
+                    v = v.to(torch.float16)
+                # MLX Conv2d expects NHWC weight layout (out, kh, kw, in)
+                if v.ndim == 4:
+                    v = v.permute(0, 2, 3, 1)
+                assert v.shape == our_sd[k].shape, \
+                    f"Shape mismatch {k}: {v.shape} vs {our_sd[k].shape}"
+                loaded[k] = mx.array(v.numpy())
+        missing = set(our_sd) - set(loaded)
+        unexpected = set(loaded) - set(our_sd)
+        if missing:
+            raise KeyError(f"[VAE-MLX] Missing keys: {missing}")
+        if unexpected:
+            raise KeyError(f"[VAE-MLX] Unexpected keys: {unexpected}")
+        self.load_weights(list(loaded.items()), strict=False)
+
+
+def load_flux2vae_mlx(path: str) -> Flux2VAEEncoder:
+    model = Flux2VAEEncoder()
+    model.load_safetensors(path)
+    return model