diff --git a/convert_to_mlx.py b/convert_to_mlx.py new file mode 100644 index 00000000..94a96a03 --- /dev/null +++ b/convert_to_mlx.py @@ -0,0 +1,289 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +"""Convert PyTorch Depth Anything 3 weights to MLX format. + +Usage: + python convert_to_mlx.py --model da3-small --input weights.safetensors --output mlx_weights.safetensors + +This script: +1. Loads PyTorch weights (safetensors or .pth). +2. Maps parameter names from PyTorch model structure to MLX model structure. +3. Transposes Conv2d weights from PyTorch (OIHW) to MLX (OHWI) format. +4. Saves the result as a safetensors file loadable by mlx. +""" + +from __future__ import annotations + +import argparse +import re +from pathlib import Path + +import mlx.core as mx +import numpy as np + + +def load_pytorch_weights(path: str) -> dict[str, np.ndarray]: + """Load weights from a PyTorch safetensors or .pth file.""" + if path.endswith(".safetensors"): + from safetensors import safe_open + weights = {} + with safe_open(path, framework="numpy") as f: + for key in f.keys(): + new_key = key + if key.startswith("model."): + new_key = key[6:] + weights[new_key] = f.get_tensor(key) + return weights + else: + import torch + state_dict = torch.load(path, map_location="cpu", weights_only=False) + if "model" in state_dict: + state_dict = state_dict["model"] + + # Also handle cases where keys in state_dict["model"] still have "model." prefix + new_dict = {} + for k, v in state_dict.items(): + new_key = k + if k.startswith("model."): + new_key = k[6:] + new_dict[new_key] = v.numpy() if hasattr(v, 'numpy') else v + return new_dict + + +def map_key(pt_key: str) -> str | None: + """Map a PyTorch state_dict key to the corresponding MLX model key. + + Returns None if the key should be skipped (e.g., GS-related weights). + """ + k = pt_key + + # Skip GS-related and Aux weights (not used in MLX forward) + if any(p in k for p in ["gs_head", "gs_adapter", "aux", "refinenet", "layer1_rn", "layer2_rn", "layer3_rn", "layer4_rn", "output_conv"]): + # Wait, some of these might be used in the main head! + # refinenet1..4 are used. refinenetX_aux are NOT. + # output_conv1, output_conv2_a/b are used. output_convX_aux are NOT. + if "_aux" in k: + return None + if "gs_head" in k or "gs_adapter" in k: + return None + + # ---- Backbone (DinoV2) ---- + # backbone.pretrained.X -> backbone.X + k = k.replace("backbone.pretrained.", "backbone.") + + # Bulk rename for Swift compatibility (snake_case -> camelCase) + for snake, camel in [ + ("cls_token", "clsToken"), ("pos_embed", "posEmbed"), + ("patch_embed", "patchEmbed"), ("camera_token", "cameraToken"), + ("layer1_rn", "layer1Rn"), ("layer2_rn", "layer2Rn"), + ("layer3_rn", "layer3Rn"), ("layer4_rn", "layer4Rn"), + ("output_conv1", "outputConv1"), ("output_conv2_a", "outputConv2a"), + ("output_conv2_b", "outputConv2b"), + # New additions for missing MLX Swift layers: + ("out_conv", "outConv"), + (".mlp.", ".mlpLayer."), + (".swiglu.", ".swiGluLayer."), + ("q_norm", "qNorm"), + ("k_norm", "kNorm"), + ("cam_enc", "camEnc"), + ("pose_branch", "poseBranch"), + ("token_norm", "tokenNorm"), + ("trunk_norm", "trunkNorm"), + ("cam_dec", "camDec"), + ("backbone_fc", "backboneFc"), + ("fc_t", "fcT"), + ("fc_qvec", "fcQvec"), + ("fc_fov_linear", "fcFovLinear"), + ]: + k = k.replace(snake, camel) + + # ---- Head (DualDPT) ---- + k = re.sub(r"head\.resize_layers\.0\.", "head.resize0.", k) + k = re.sub(r"head\.resize_layers\.1\.", "head.resize1.", k) + k = re.sub(r"head\.resize_layers\.3\.", "head.resize3.", k) + + # head.scratch.output_conv -> head.output_conv + k = k.replace("head.scratch.", "head.") + + # head.refinenet1.resConfUnit1.conv1 -> head.refinenet1.resConfUnit1.conv1 + # head.output_conv1 -> head.output_conv1 + + # head.output_conv2.0 -> head.outputConv2a + # head.output_conv2.2 -> head.outputConv2b + k = re.sub(r"head\.output_conv2\.0\.", "head.outputConv2a.", k) + k = re.sub(r"head\.output_conv2\.2\.", "head.outputConv2b.", k) + # Skip output_conv2.1 (ReLU, no params) + if re.search(r"head\.output_conv2\.\d+\.", k) and "outputConv2a" not in k and "outputConv2b" not in k: + return None + + # Aux head: output_conv1_aux.N.M -> head.output_conv1_aux.N.layers.M + k = re.sub( + r"head\.output_conv1_aux\.(\d+)\.(\d+)\.", + r"head.output_conv1_aux.\1.layers.\2.", + k, + ) + + # head.output_conv2_aux.N.0 -> head.output_conv2_aux.N.conv1 + # head.output_conv2_aux.N.1 -> (Permute, skip) + # head.output_conv2_aux.N.2 -> head.output_conv2_aux.N.ln + # head.output_conv2_aux.N.3 -> (ReLU, skip) + # head.output_conv2_aux.N.4 -> (Identity, skip) + # head.output_conv2_aux.N.5 -> head.output_conv2_aux.N.conv2 + m = re.match(r"head\.output_conv2_aux\.(\d+)\.0\.(.*)", k) + if m: + k = f"head.output_conv2_aux.{m.group(1)}.conv1.{m.group(2)}" + m = re.match(r"head\.output_conv2_aux\.(\d+)\.2\.(.*)", k) + if m: + k = f"head.output_conv2_aux.{m.group(1)}.ln.{m.group(2)}" + m = re.match(r"head\.output_conv2_aux\.(\d+)\.[45]\.(.*)", k) + if m: + # Both 4 and 5 in PyTorch often map to the final conv in different configs + k = f"head.output_conv2_aux.{m.group(1)}.conv2.{m.group(2)}" + + # Skip Permute, ReLU, and optional Identity (indices 1, 3, 4) + if re.match(r"head\.output_conv2_aux\.\d+\.[134]\.", k) and "conv2" not in k: + return None + + # ---- Camera Encoder ---- + # cam_enc.trunk.N -> cam_enc.trunk.N + # cam_enc.pose_branch.fc1 -> cam_enc.pose_branch.fc1 + # cam_enc.token_norm -> cam_enc.token_norm + # cam_enc.trunk_norm -> cam_enc.trunk_norm + + # ---- Camera Decoder ---- + # camDec.backbone.0 -> camDec.backboneFc1 + # camDec.backbone.2 -> camDec.backboneFc2 + k = re.sub(r"camDec\.backbone\.0\.", "camDec.backboneFc1.", k) + k = re.sub(r"camDec\.backbone\.2\.", "camDec.backboneFc2.", k) + # Skip indices 1,3 (ReLU, no params) + if re.match(r"camDec\.backbone\.\d+\.", k): + return None + + # camDec.fc_fov.0 -> camDec.fcFovLinear + k = re.sub(r"camDec\.fc_fov\.0\.", "camDec.fcFovLinear.", k) + # Skip fc_fov.1 (ReLU) + if re.match(r"camDec\.fc_fov\.\d+\.", k): + return None + + # Handle skip_add (FloatFunctional) - no parameters + if "skip_add" in k: + return None + + return k + + +def is_conv_weight(key: str, shape: tuple) -> bool: + """Check if a weight tensor is a Conv2d/ConvTranspose2d weight that needs transposing. + + PyTorch conv weights: (out_ch, in_ch, kH, kW) -> 4D with spatial dims + MLX conv weights: (out_ch, kH, kW, in_ch) + """ + if len(shape) != 4: + return False + # Check if it's a weight (not bias) in a conv-like layer + if key.endswith(".weight"): + # Identify conv layers by name patterns (using Swift mapped names) + conv_patterns = [ + "patchEmbed.proj.weight", + "conv1.weight", "conv2.weight", + "outConv.weight", + "layer1Rn.weight", "layer2Rn.weight", + "layer3Rn.weight", "layer4Rn.weight", + "outputConv", + "resize0.weight", "resize1.weight", "resize3.weight", + "projects.", + ] + return any(p in key for p in conv_patterns) + return False + + +def transpose_conv_weight(weight: np.ndarray) -> np.ndarray: + """Transpose conv weight from PyTorch (OIHW) to MLX (OHWI).""" + return np.ascontiguousarray(weight.transpose(0, 2, 3, 1)) + + +def is_conv_transpose_weight(key: str) -> bool: + """Check if this is a ConvTranspose2d weight.""" + return "resize0.weight" in key or "resize1.weight" in key + + +def transpose_conv_transpose_weight(weight: np.ndarray) -> np.ndarray: + """Transpose ConvTranspose2d weight from PyTorch (I,O,kH,kW) to MLX (O,kH,kW,I).""" + return np.ascontiguousarray(weight.transpose(1, 2, 3, 0)) + + +def convert_weights( + pt_weights: dict[str, np.ndarray], +) -> dict[str, mx.array]: + """Convert PyTorch weights to MLX format. + + Returns: + dict mapping MLX key paths to mx.array values. + """ + mlx_weights = {} + skipped = [] + for pt_key, value in pt_weights.items(): + mlx_key = map_key(pt_key) + if mlx_key is None: + skipped.append(pt_key) + continue + + arr = value + if is_conv_transpose_weight(mlx_key): + arr = transpose_conv_transpose_weight(arr) + elif is_conv_weight(mlx_key, arr.shape): + arr = transpose_conv_weight(arr) + + mlx_weights[mlx_key] = mx.array(arr) + + if skipped: + print(f"Skipped {len(skipped)} keys (GS/unused):") + for s in skipped[:10]: + print(f" {s}") + if len(skipped) > 10: + print(f" ... and {len(skipped) - 10} more") + + return mlx_weights + + +def main(): + parser = argparse.ArgumentParser(description="Convert DA3 PyTorch weights to MLX") + parser.add_argument("--model", type=str, default="da3-small", + choices=["da3-small", "da3-base", "da3-large", "da3-giant"], + help="Model configuration name") + parser.add_argument("--input", type=str, required=True, + help="Path to PyTorch weights (.safetensors or .pth)") + parser.add_argument("--output", type=str, required=True, + help="Output path for MLX weights (.safetensors or .npz)") + parser.add_argument("--dtype", type=str, default="float16", + choices=["float32", "float16", "bfloat16"], + help="Output data type for MLX weights") + args = parser.parse_args() + + print(f"Loading PyTorch weights from {args.input}...") + pt_weights = load_pytorch_weights(args.input) + print(f" Loaded {len(pt_weights)} parameters") + + print(f"Converting weights to {args.dtype}...") + mlx_weights = convert_weights(pt_weights) + + # Cast weights to requested dtype + dtype = getattr(mx, args.dtype) + mlx_weights = {k: v.astype(dtype) for k, v in mlx_weights.items()} + + print(f" Converted {len(mlx_weights)} parameters") + + print(f"Saving MLX weights to {args.output}...") + if args.output.endswith(".safetensors"): + mx.save_safetensors(args.output, mlx_weights) + elif args.output.endswith(".npz"): + mx.savez(args.output, **mlx_weights) + else: + raise ValueError(f"Unsupported output format: {args.output}") + print("Done!") + + +if __name__ == "__main__": + main() diff --git a/inspect_npy.py b/inspect_npy.py new file mode 100644 index 00000000..110577c0 --- /dev/null +++ b/inspect_npy.py @@ -0,0 +1,59 @@ +import os +import numpy as np +import matplotlib.pyplot as plt + +def inspect_data(): + dir_path = "/Users/aweise/dev/sunscape2/sunscape2/DepthAnything3MLX/Tests/DepthAnything3MLXTests/" + input_path = os.path.join(dir_path, "input_images.npy") + depth_path = os.path.join(dir_path, "expected_depth.npy") + + if not os.path.exists(input_path) or not os.path.exists(depth_path): + print("Files not found. Please ensure they exist.") + return + + images = np.load(input_path) + depths = np.load(depth_path) + + print(f"Loaded input_images with shape: {images.shape}") + print(f"Loaded expected_depth with shape: {depths.shape}") + + # input_images should be (1, N, 518, 518, 3) + # expected_depth should be (N, 518, 518) + if len(images.shape) == 5: + images = images[0] + + N = images.shape[0] + + # De-normalize image for plotting + mean = np.array([0.485, 0.456, 0.406]) + std = np.array([0.229, 0.224, 0.225]) + + fig, axes = plt.subplots(N, 2, figsize=(10, 4 * N)) + + # Handle the case where N=1 so axes behavior is consistent + if N == 1: + axes = [axes] + + for i in range(N): + img = images[i] + # De-normalize and clip + img_visual = np.clip(img * std + mean, 0, 1) + depth = depths[i] + + ax_img = axes[i][0] + ax_depth = axes[i][1] + + ax_img.imshow(img_visual) + ax_img.set_title(f"Input Image {i+1}") + ax_img.axis('off') + + depth_plot = ax_depth.imshow(depth, cmap='inferno') + ax_depth.set_title(f"Expected Depth {i+1}") + ax_depth.axis('off') + fig.colorbar(depth_plot, ax=ax_depth, fraction=0.046, pad=0.04) + + plt.tight_layout() + plt.show() + +if __name__ == "__main__": + inspect_data() diff --git a/mlx_depth_anything_3/__init__.py b/mlx_depth_anything_3/__init__.py new file mode 100644 index 00000000..bc374f8b --- /dev/null +++ b/mlx_depth_anything_3/__init__.py @@ -0,0 +1,9 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +"""MLX implementation of Depth Anything 3 for Apple Silicon.""" + +from mlx_depth_anything_3.model import DepthAnything3Net +from mlx_depth_anything_3.inference import DepthAnything3 + +__all__ = ["DepthAnything3Net", "DepthAnything3"] diff --git a/mlx_depth_anything_3/backbone.py b/mlx_depth_anything_3/backbone.py new file mode 100644 index 00000000..22dfd12e --- /dev/null +++ b/mlx_depth_anything_3/backbone.py @@ -0,0 +1,382 @@ +# Copyright (c) Meta Platforms, Inc. and affiliates. +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +"""DinoV2 Vision Transformer backbone for MLX.""" + +from __future__ import annotations + +import math + +import mlx.core as mx +import mlx.nn as nn + +from mlx_depth_anything_3.layers import ( + Attention, + Block, + Mlp, + ModuleList, + PatchEmbed, + PositionGetter, + RotaryPositionEmbedding2D, + SwiGLUFFN, +) +from mlx_depth_anything_3.reference_view import ( + THRESH_FOR_REF_SELECTION, + reorder_by_reference, + restore_original_order, + select_reference_view, +) + + +class DinoVisionTransformer(nn.Module): + """DinoV2 Vision Transformer with alternating local/global attention.""" + + PATCH_SIZE = 14 + + def __init__( + self, + img_size: int = 518, + patch_size: int = 14, + in_chans: int = 3, + embed_dim: int = 768, + depth: int = 12, + num_heads: int = 12, + mlp_ratio: float = 4.0, + qkv_bias: bool = True, + ffn_bias: bool = True, + proj_bias: bool = True, + init_values: float = 1.0, + ffn_layer: str = "mlp", + num_register_tokens: int = 0, + interpolate_offset: float = 0.1, + alt_start: int = -1, + qknorm_start: int = -1, + rope_start: int = -1, + rope_freq: float = 100.0, + cat_token: bool = True, + ): + super().__init__() + self.patch_start_idx = 1 + self.embed_dim = embed_dim + self.alt_start = alt_start + self.qknorm_start = qknorm_start + self.rope_start = rope_start + self.cat_token = cat_token + self.num_tokens = 1 + self.n_blocks = depth + self.num_heads = num_heads + self.patch_size = patch_size + self.num_register_tokens = num_register_tokens + self.interpolate_offset = interpolate_offset + + self.patch_embed = PatchEmbed( + img_size=img_size, patch_size=patch_size, + in_chans=in_chans, embed_dim=embed_dim, + ) + num_patches = self.patch_embed.num_patches + self.cls_token = mx.zeros((1, 1, embed_dim)) + if alt_start != -1: + self.camera_token = mx.random.normal((1, 2, embed_dim)) * 0.02 + self.pos_embed = mx.zeros((1, num_patches + self.num_tokens, embed_dim)) + self.register_tokens = ( + mx.zeros((1, num_register_tokens, embed_dim)) + if num_register_tokens > 0 else None + ) + + # RoPE + self.rope = None + self.position_getter = None + if rope_start != -1: + if rope_freq > 0: + self.rope = RotaryPositionEmbedding2D(frequency=rope_freq) + self.position_getter = PositionGetter() + + # Build blocks + ffn_cls = "swiglu" if ffn_layer in ("swiglufused", "swiglu") else "mlp" + self.blocks = ModuleList([ + Block( + dim=embed_dim, + num_heads=num_heads, + mlp_ratio=mlp_ratio, + qkv_bias=qkv_bias, + proj_bias=proj_bias, + ffn_bias=ffn_bias, + init_values=init_values, + ffn_layer=ffn_cls, + qk_norm=(i >= qknorm_start) if qknorm_start != -1 else False, + rope=self.rope if (i >= rope_start and rope_start != -1) else None, + ) + for i in range(depth) + ]) + self.norm = nn.LayerNorm(embed_dim, eps=1e-6) + + def interpolate_pos_encoding(self, x: mx.array, w: int, h: int) -> mx.array: + """Interpolate position embeddings for arbitrary resolution.""" + npatch = x.shape[1] - 1 + N = self.pos_embed.shape[1] - 1 + if npatch == N and w == h: + return self.pos_embed + + pos_embed = self.pos_embed.astype(mx.float32) + class_pos_embed = pos_embed[:, 0:1] + patch_pos_embed = pos_embed[:, 1:] + dim = x.shape[-1] + w0 = w // self.patch_size + h0 = h // self.patch_size + M = int(math.sqrt(N)) + + # Reshape to spatial, interpolate, flatten + # (1, M, M, dim) -> (1, h0, w0, dim) + spatial = mx.reshape(patch_pos_embed, (1, M, M, dim)) + + if self.interpolate_offset: + sx = float(w0 + self.interpolate_offset) / M + sy = float(h0 + self.interpolate_offset) / M + # Compute target size + th = int(round(M * sy)) + tw = int(round(M * sx)) + else: + th, tw = h0, w0 + + from mlx_depth_anything_3.layers import bilinear_interpolate + spatial = bilinear_interpolate(spatial, th, tw, align_corners=False) + # Crop to exact size if needed + spatial = spatial[:, :h0, :w0, :] + patch_pos_embed = mx.reshape(spatial, (1, -1, dim)) + return mx.concatenate([class_pos_embed, patch_pos_embed], axis=1).astype(x.dtype) + + def prepare_tokens(self, x: mx.array) -> mx.array: + """Prepare tokens from multi-view images. + + Args: + x: (B, S, H, W, C) in NHWC. + + Returns: + (B, S, N+1, D) token tensor with cls_token prepended. + """ + B, S, H, W, C = x.shape + # Flatten batch and views + x_flat = mx.reshape(x, (B * S, H, W, C)) + x_flat = self.patch_embed(x_flat) # (B*S, num_patches, D) + + # Prepend cls token + cls = mx.broadcast_to(self.cls_token, (B * S, 1, self.embed_dim)) + x_flat = mx.concatenate([cls, x_flat], axis=1) + + # Add position embeddings + x_flat = x_flat + self.interpolate_pos_encoding(x_flat, W, H) + + # Add register tokens if present + if self.register_tokens is not None: + reg = mx.broadcast_to(self.register_tokens, (B * S, self.num_register_tokens, self.embed_dim)) + x_flat = mx.concatenate([x_flat[:, :1], reg, x_flat[:, 1:]], axis=1) + + # Reshape to (B, S, N, D) + N = x_flat.shape[1] + D = x_flat.shape[2] + x_out = mx.reshape(x_flat, (B, S, N, D)) + return x_out + + def _prepare_rope(self, B: int, S: int, H: int, W: int) -> tuple: + """Prepare RoPE position tensors.""" + if self.rope is None or self.position_getter is None: + return None, None + ph = H // self.patch_size + pw = W // self.patch_size + pos = self.position_getter(B * S, ph, pw) # (B*S, ph*pw, 2) + pos = mx.reshape(pos, (B, S, ph * pw, 2)) + pos_nodiff = mx.zeros_like(pos) + if self.patch_start_idx > 0: + pos = pos + 1 + special = mx.zeros((B, S, self.patch_start_idx, 2)) + pos = mx.concatenate([special, pos], axis=2) + pos_nodiff = pos_nodiff + 1 + pos_nodiff = mx.concatenate([mx.zeros((B, S, self.patch_start_idx, 2)), pos_nodiff], axis=2) + return pos, pos_nodiff + + def process_attention(self, x: mx.array, block: Block, attn_type: str, + pos: mx.array | None = None, + attn_mask: mx.array | None = None) -> mx.array: + """Run one block with local or global attention. + + Args: + x: (B, S, N, C) token tensor. + block: transformer block. + attn_type: "local" or "global". + pos: optional RoPE positions. + """ + b, s, n, c = x.shape + if attn_type == "local": + x_flat = mx.reshape(x, (b * s, n, c)) + if pos is not None: + pos_flat = mx.reshape(pos, (b * s, pos.shape[2], pos.shape[3])) + else: + pos_flat = None + x_flat = block(x_flat, pos=pos_flat, attn_mask=attn_mask) + return mx.reshape(x_flat, (b, s, n, c)) + else: # global + x_flat = mx.reshape(x, (b, s * n, c)) + if pos is not None: + pos_flat = mx.reshape(pos, (b, s * pos.shape[2], pos.shape[3])) + else: + pos_flat = None + x_flat = block(x_flat, pos=pos_flat, attn_mask=attn_mask) + return mx.reshape(x_flat, (b, s, n, c)) + + def get_intermediate_layers( + self, + x: mx.array, + out_layers: list[int], + cam_token: mx.array | None = None, + export_feat_layers: list[int] | None = None, + ref_view_strategy: str = "saddle_balanced", + ) -> tuple: + """Extract intermediate features from the backbone. + + Args: + x: (B, S, H, W, C) multi-view images in NHWC. + out_layers: list of layer indices to extract. + cam_token: (B, S, D) optional camera conditioning tokens. + export_feat_layers: additional layer indices for auxiliary features. + ref_view_strategy: reference view selection strategy. + + Returns: + (features, aux_features) tuple. + """ + if export_feat_layers is None: + export_feat_layers = [] + B, S, H, W, C = x.shape + + x_tok = self.prepare_tokens(x) # (B, S, N, D) + pos, pos_nodiff = self._prepare_rope(B, S, H, W) + + output = [] + aux_output = [] + blocks_to_take = set(out_layers) + local_x = x_tok + b_idx = None + + for i, blk in enumerate(self.blocks): + # RoPE positions + if i < self.rope_start or self.rope is None: + g_pos, l_pos = None, None + else: + g_pos = pos_nodiff + l_pos = pos + + # Reference view selection (before alt_start). + # Skipped when external camera tokens are provided, because the + # camera encoder already supplies per-view conditioning so + # feature-based reordering is unnecessary. + if (self.alt_start != -1 and i == self.alt_start - 1 + and x_tok.shape[1] >= THRESH_FOR_REF_SELECTION + and cam_token is None): + b_idx = select_reference_view(x_tok, strategy=ref_view_strategy) + x_tok = reorder_by_reference(x_tok, b_idx) + local_x = reorder_by_reference(local_x, b_idx) + + # Camera token injection + if self.alt_start != -1 and i == self.alt_start: + if cam_token is not None: + ct = cam_token + else: + ref_tok = mx.broadcast_to(self.camera_token[:, :1], (B, 1, self.embed_dim)) + src_tok = mx.broadcast_to(self.camera_token[:, 1:], (B, S - 1, self.embed_dim)) + ct = mx.concatenate([ref_tok, src_tok], axis=1) + # Replace cls token position with camera token + x_tok_list = [] + for bidx in range(B): + view_tokens = [] + for sidx in range(S): + tokens = x_tok[bidx, sidx] + tokens_new = mx.concatenate( + [mx.expand_dims(ct[bidx, sidx], 0), tokens[1:]], axis=0 + ) + view_tokens.append(tokens_new) + x_tok_list.append(mx.stack(view_tokens)) + x_tok = mx.stack(x_tok_list) + + # Local or global attention + if self.alt_start != -1 and i >= self.alt_start and i % 2 == 1: + x_tok = self.process_attention(x_tok, blk, "global", pos=g_pos) + else: + x_tok = self.process_attention(x_tok, blk, "local", pos=l_pos) + local_x = x_tok + + # Collect output features + if i in blocks_to_take: + if self.cat_token: + out_x = mx.concatenate([local_x, x_tok], axis=-1) + else: + out_x = x_tok + # Restore original order if reordering was applied + if (x_tok.shape[1] >= THRESH_FOR_REF_SELECTION + and self.alt_start != -1 and b_idx is not None): + out_x = restore_original_order(out_x, b_idx) + output.append((out_x[:, :, 0], out_x)) + + if i in export_feat_layers: + aux_output.append(x_tok) + + # Apply final norm + camera_tokens = [out[0] for out in output] + if output[0][1].shape[-1] == self.embed_dim: + outputs = [self.norm(out[1]) for out in output] + elif output[0][1].shape[-1] == self.embed_dim * 2: + outputs = [ + mx.concatenate( + [out[1][..., :self.embed_dim], self.norm(out[1][..., self.embed_dim:])], + axis=-1, + ) + for out in output + ] + else: + raise ValueError(f"Invalid output shape: {output[0][1].shape}") + + aux_output = [self.norm(out) for out in aux_output] + + # Drop cls + register tokens + skip = 1 + self.num_register_tokens + outputs = [out[..., skip:, :] for out in outputs] + aux_output = [out[..., skip:, :] for out in aux_output] + + return list(zip(outputs, camera_tokens)), aux_output + + +class DinoV2(nn.Module): + """DinoV2 wrapper matching the PyTorch API.""" + + def __init__( + self, + name: str, + out_layers: list[int], + alt_start: int = -1, + qknorm_start: int = -1, + rope_start: int = -1, + cat_token: bool = True, + ): + super().__init__() + self.name = name + self.out_layers = out_layers + encoder_map = { + "vits": dict(embed_dim=384, depth=12, num_heads=6), + "vitb": dict(embed_dim=768, depth=12, num_heads=12), + "vitl": dict(embed_dim=1024, depth=24, num_heads=16), + "vitg": dict(embed_dim=1536, depth=40, num_heads=24), + } + cfg = encoder_map[name] + ffn = "swiglu" if name == "vitg" else "mlp" + self.pretrained = DinoVisionTransformer( + img_size=518, + patch_size=14, + ffn_layer=ffn, + alt_start=alt_start, + qknorm_start=qknorm_start, + rope_start=rope_start, + cat_token=cat_token, + **cfg, + ) + + def __call__(self, x: mx.array, **kwargs) -> tuple: + return self.pretrained.get_intermediate_layers(x, self.out_layers, **kwargs) diff --git a/mlx_depth_anything_3/camera.py b/mlx_depth_anything_3/camera.py new file mode 100644 index 00000000..ac539d6b --- /dev/null +++ b/mlx_depth_anything_3/camera.py @@ -0,0 +1,108 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +"""Camera encoder and decoder modules for MLX.""" + +from __future__ import annotations + +import mlx.core as mx +import mlx.nn as nn + +from mlx_depth_anything_3.layers import CamBlock, Mlp, ModuleList +from mlx_depth_anything_3.transform import ( + affine_inverse, + extri_intri_to_pose_encoding, + pose_encoding_to_extri_intri, +) + + +class CameraEnc(nn.Module): + """Camera encoder: convert extrinsics+intrinsics to conditioning tokens. + + Uses an MLP to project the 9-D pose encoding to token space, then + refines via transformer blocks. + """ + + def __init__( + self, + dim_out: int = 1024, + dim_in: int = 9, + trunk_depth: int = 4, + num_heads: int = 16, + mlp_ratio: int = 4, + init_values: float = 0.01, + **kwargs, + ): + super().__init__() + self.pose_branch = Mlp( + in_features=dim_in, + hidden_features=dim_out // 2, + out_features=dim_out, + bias=True, + ) + self.token_norm = nn.LayerNorm(dim_out, eps=1e-6) + self.trunk = ModuleList([ + CamBlock(dim=dim_out, num_heads=num_heads, mlp_ratio=mlp_ratio, + init_values=init_values) + for _ in range(trunk_depth) + ]) + self.trunk_norm = nn.LayerNorm(dim_out, eps=1e-6) + + def __call__( + self, + ext: mx.array, + ixt: mx.array, + image_size: tuple[int, int], + ) -> mx.array: + """Encode camera parameters to conditioning tokens. + + Args: + ext: (B, S, 4, 4) world-to-camera extrinsics. + ixt: (B, S, 3, 3) intrinsics. + image_size: (H, W). + + Returns: + (B, S, dim_out) camera tokens. + """ + c2ws = affine_inverse(ext) + pose_encoding = extri_intri_to_pose_encoding(c2ws, ixt, image_size) + tokens = self.pose_branch(pose_encoding) + tokens = self.token_norm(tokens) + for blk in self.trunk: + tokens = blk(tokens) + tokens = self.trunk_norm(tokens) + return tokens + + +class CameraDec(nn.Module): + """Camera decoder: predict pose encoding from features. + + Not used for depth-only inference, but included for weight loading. + """ + + def __init__(self, dim_in: int = 1536, **kwargs): + super().__init__() + self.backbone_fc1 = nn.Linear(dim_in, dim_in) + self.backbone_fc2 = nn.Linear(dim_in, dim_in) + self.fc_t = nn.Linear(dim_in, 3) + self.fc_qvec = nn.Linear(dim_in, 4) + self.fc_fov_linear = nn.Linear(dim_in, 2) + + def __call__(self, feat: mx.array) -> mx.array: + """Predict 9-D pose encoding from features. + + Args: + feat: (B, S, D) features. + + Returns: + (B, S, 9) pose encoding [T(3), quat(4), fov(2)]. + """ + B, S = feat.shape[:2] + x = mx.reshape(feat, (B * S, -1)) + x = nn.relu(self.backbone_fc1(x)) + x = nn.relu(self.backbone_fc2(x)) + x_f = x.astype(mx.float32) + t = mx.reshape(self.fc_t(x_f), (B, S, 3)) + qvec = mx.reshape(self.fc_qvec(x_f), (B, S, 4)) + fov = nn.relu(mx.reshape(self.fc_fov_linear(x_f), (B, S, 2))) + return mx.concatenate([t, qvec, fov], axis=-1) diff --git a/mlx_depth_anything_3/dpt.py b/mlx_depth_anything_3/dpt.py new file mode 100644 index 00000000..6d61808c --- /dev/null +++ b/mlx_depth_anything_3/dpt.py @@ -0,0 +1,346 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +"""DualDPT depth prediction head for MLX. + +Only the main (depth) head is implemented; the auxiliary ray head +is omitted since we only export predicted depth. +""" + +from __future__ import annotations + +import mlx.core as mx +import mlx.nn as nn + +from mlx_depth_anything_3.layers import ( + ModuleList, + bilinear_interpolate, + create_uv_grid, + position_grid_to_embed, +) + + +# --------------------------------------------------------------------------- +# ResidualConvUnit +# --------------------------------------------------------------------------- +class ResidualConvUnit(nn.Module): + """Lightweight residual conv block used within fusion.""" + + def __init__(self, features: int): + super().__init__() + # MLX Conv2d: (N, H, W, C) in/out + self.conv1 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True) + self.conv2 = nn.Conv2d(features, features, kernel_size=3, stride=1, padding=1, bias=True) + + def __call__(self, x: mx.array) -> mx.array: + out = nn.relu(x) + out = self.conv1(out) + out = nn.relu(out) + out = self.conv2(out) + return out + x + + +# --------------------------------------------------------------------------- +# FeatureFusionBlock +# --------------------------------------------------------------------------- +class FeatureFusionBlock(nn.Module): + """Top-down fusion block: optional residual merge + upsample + 1x1 conv.""" + + def __init__(self, features: int, has_residual: bool = True): + super().__init__() + self.has_residual = has_residual + self.resConfUnit1 = ResidualConvUnit(features) if has_residual else None + self.resConfUnit2 = ResidualConvUnit(features) + self.out_conv = nn.Conv2d(features, features, kernel_size=1, stride=1, padding=0, bias=True) + + def __call__(self, *xs, size=None) -> mx.array: + """Forward. + + Args: + xs[0]: top-path input (N, H, W, C). + xs[1]: optional lateral input. + size: (target_h, target_w) for upsampling. + """ + y = xs[0] + if self.has_residual and len(xs) > 1 and self.resConfUnit1 is not None: + y = y + self.resConfUnit1(xs[1]) + y = self.resConfUnit2(y) + + # Upsample + if size is not None: + target_h, target_w = size + else: + target_h = y.shape[1] * 2 + target_w = y.shape[2] * 2 + y = bilinear_interpolate(y, target_h, target_w, align_corners=True) + y = self.out_conv(y) + return y + + +# --------------------------------------------------------------------------- +# Permute helper (for Sequential with LayerNorm on channel dim) +# --------------------------------------------------------------------------- +class PermuteToChannelsLast(nn.Module): + """(N, C, H, W) -> no-op since MLX already uses NHWC.""" + def __call__(self, x: mx.array) -> mx.array: + return x + + +# --------------------------------------------------------------------------- +# DualDPT – depth-only variant +# --------------------------------------------------------------------------- +class DualDPT(nn.Module): + """Dual-head DPT for depth prediction. + + Only the main depth head is computed (aux ray head is skipped). + + All convolutions work in NHWC (channels-last) layout. + """ + + def __init__( + self, + dim_in: int, + patch_size: int = 14, + output_dim: int = 2, + activation: str = "exp", + conf_activation: str = "expp1", + features: int = 256, + out_channels: list[int] | tuple[int, ...] = (256, 512, 1024, 1024), + pos_embed: bool = True, + down_ratio: int = 1, + **kwargs, + ): + super().__init__() + self.patch_size = patch_size + self.activation = activation + self.conf_activation = conf_activation + self.pos_embed = pos_embed + self.down_ratio = down_ratio + self.output_dim = output_dim + + # Token pre-norm + per-stage 1x1 projection + self.norm = nn.LayerNorm(dim_in, eps=1e-6) + self.projects = ModuleList([ + nn.Conv2d(dim_in, oc, kernel_size=1, stride=1, padding=0) + for oc in out_channels + ]) + + # Spatial resize layers (NHWC) + # Stage 0: 4x upsample via ConvTranspose + # Stage 1: 2x upsample via ConvTranspose + # Stage 2: identity + # Stage 3: 2x downsample via Conv stride=2 + self.resize_0 = nn.ConvTranspose2d(out_channels[0], out_channels[0], + kernel_size=4, stride=4, padding=0) + self.resize_1 = nn.ConvTranspose2d(out_channels[1], out_channels[1], + kernel_size=2, stride=2, padding=0) + # resize_2 = identity + self.resize_3 = nn.Conv2d(out_channels[3], out_channels[3], + kernel_size=3, stride=2, padding=1) + + # Scratch: layer adapters (3x3 conv, no bias) + self.layer1_rn = nn.Conv2d(out_channels[0], features, kernel_size=3, + stride=1, padding=1, bias=False) + self.layer2_rn = nn.Conv2d(out_channels[1], features, kernel_size=3, + stride=1, padding=1, bias=False) + self.layer3_rn = nn.Conv2d(out_channels[2], features, kernel_size=3, + stride=1, padding=1, bias=False) + self.layer4_rn = nn.Conv2d(out_channels[3], features, kernel_size=3, + stride=1, padding=1, bias=False) + + # Main fusion chain + self.refinenet4 = FeatureFusionBlock(features, has_residual=False) + self.refinenet3 = FeatureFusionBlock(features, has_residual=True) + self.refinenet2 = FeatureFusionBlock(features, has_residual=True) + self.refinenet1 = FeatureFusionBlock(features, has_residual=True) + + # Main output head: conv1 (3x3) -> conv2 (sequence: 3x3 + relu + 1x1) + head_features_1 = features + head_features_2 = 32 + self.output_conv1 = nn.Conv2d(head_features_1, head_features_1 // 2, + kernel_size=3, stride=1, padding=1) + # output_conv2 sequence + self.output_conv2_a = nn.Conv2d(head_features_1 // 2, head_features_2, + kernel_size=3, stride=1, padding=1) + self.output_conv2_b = nn.Conv2d(head_features_2, output_dim, + kernel_size=1, stride=1, padding=0) + + # Auxiliary fusion chain (loaded for weight compatibility only; not used in forward) + self.refinenet4_aux = FeatureFusionBlock(features, has_residual=False) + self.refinenet3_aux = FeatureFusionBlock(features, has_residual=True) + self.refinenet2_aux = FeatureFusionBlock(features, has_residual=True) + self.refinenet1_aux = FeatureFusionBlock(features, has_residual=True) + + # Aux output layers (loaded for weight compatibility only; not used in forward) + self.output_conv1_aux = ModuleList([ + self._make_aux_out1_module(head_features_1) + for _ in range(4) + ]) + + # Aux output conv2 per level (with LN) + self.output_conv2_aux = ModuleList([ + self._make_aux_out2_module(head_features_1 // 2, head_features_2) + for _ in range(4) + ]) + + def _make_aux_out1_module(self, in_ch: int): + """5-conv auxiliary pre-head stack.""" + class AuxOut1(nn.Module): + def __init__(self, in_ch): + super().__init__() + self.layers = [ + nn.Conv2d(in_ch, in_ch // 2, kernel_size=3, stride=1, padding=1), + nn.Conv2d(in_ch // 2, in_ch, kernel_size=3, stride=1, padding=1), + nn.Conv2d(in_ch, in_ch // 2, kernel_size=3, stride=1, padding=1), + nn.Conv2d(in_ch // 2, in_ch, kernel_size=3, stride=1, padding=1), + nn.Conv2d(in_ch, in_ch // 2, kernel_size=3, stride=1, padding=1), + ] + def __call__(self, x): + for layer in self.layers: + x = layer(x) + return x + return AuxOut1(in_ch) + + def _make_aux_out2_module(self, in_ch: int, mid_ch: int): + """Aux final projection: conv3x3 -> LN -> ReLU -> conv1x1.""" + class AuxOut2(nn.Module): + def __init__(self, in_ch, mid_ch): + super().__init__() + self.conv1 = nn.Conv2d(in_ch, mid_ch, kernel_size=3, stride=1, padding=1) + self.ln = nn.LayerNorm(mid_ch) + self.conv2 = nn.Conv2d(mid_ch, 7, kernel_size=1, stride=1, padding=0) + def __call__(self, x): + return self.conv2(nn.relu(self.ln(self.conv1(x)))) + return AuxOut2(in_ch, mid_ch) + + def __call__( + self, + feats: list, + H: int, + W: int, + patch_start_idx: int = 0, + ) -> dict[str, mx.array]: + """Forward pass. + + Args: + feats: list of 4 (features, camera_tokens) tuples. + features: (B, S, N, C). + H, W: original image dimensions. + patch_start_idx: index where patch tokens start in the sequence. + + Returns: + {"depth": (B, S, H, W), "depth_conf": (B, S, H, W)} + """ + B, S, N, C = feats[0][0].shape + feat_list = [feat[0] for feat in feats] + # Flatten B*S + feat_list = [mx.reshape(f, (B * S, N, C)) for f in feat_list] + + out_dict = self._forward_impl(feat_list, H, W, patch_start_idx) + # Reshape back to (B, S, ...) + out_dict = {k: mx.reshape(v, (B, S, *v.shape[1:])) for k, v in out_dict.items()} + return out_dict + + def _forward_impl( + self, + feats: list[mx.array], + H: int, + W: int, + patch_start_idx: int, + ) -> dict[str, mx.array]: + BS = feats[0].shape[0] + C = feats[0].shape[2] + ph = H // self.patch_size + pw = W // self.patch_size + + resized_feats = [] + for stage_idx in range(4): + x = feats[stage_idx][:, patch_start_idx:] # (BS, N_patch, C) + x = self.norm(x) + # Reshape to spatial NHWC: (BS, ph, pw, C) + x = mx.reshape(x, (BS, ph, pw, C)) + # 1x1 projection + x = self.projects[stage_idx](x) + if self.pos_embed: + x = self._add_pos_embed(x, W, H) + # Resize + if stage_idx == 0: + x = self.resize_0(x) + elif stage_idx == 1: + x = self.resize_1(x) + elif stage_idx == 2: + pass # identity + else: + x = self.resize_3(x) + resized_feats.append(x) + + # Fuse pyramid (main path only) + fused_main = self._fuse_main(resized_feats) + + # Upsample to target resolution + h_out = int(ph * self.patch_size / self.down_ratio) + w_out = int(pw * self.patch_size / self.down_ratio) + + fused_main = bilinear_interpolate(fused_main, h_out, w_out, align_corners=True) + if self.pos_embed: + fused_main = self._add_pos_embed(fused_main, W, H) + + # Main head: conv2 sequence + main_logits = self.output_conv2_a(fused_main) # (BS, H, W, 32) + main_logits = nn.relu(main_logits) + main_logits = self.output_conv2_b(main_logits) # (BS, H, W, output_dim) + + # Activation + depth = self._apply_activation(main_logits[..., :-1], self.activation) + conf = self._apply_activation(main_logits[..., -1], self.conf_activation) + # depth: (BS, H, W) or (BS, H, W, 1) -> squeeze + depth = mx.squeeze(depth, axis=-1) if depth.ndim == 4 else depth + + return {"depth": depth, "depth_conf": conf} + + def _fuse_main(self, feats: list[mx.array]) -> mx.array: + """Main fusion pyramid.""" + l1, l2, l3, l4 = feats + + l1_rn = self.layer1_rn(l1) + l2_rn = self.layer2_rn(l2) + l3_rn = self.layer3_rn(l3) + l4_rn = self.layer4_rn(l4) + + # 4 -> 3 -> 2 -> 1 + out = self.refinenet4(l4_rn, size=(l3_rn.shape[1], l3_rn.shape[2])) + out = self.refinenet3(out, l3_rn, size=(l2_rn.shape[1], l2_rn.shape[2])) + out = self.refinenet2(out, l2_rn, size=(l1_rn.shape[1], l1_rn.shape[2])) + out = self.refinenet1(out, l1_rn) + + out = self.output_conv1(out) + return out + + def _add_pos_embed(self, x: mx.array, W: int, H: int, + ratio: float = 0.1) -> mx.array: + """Add UV positional embedding to feature map (NHWC).""" + ph, pw = x.shape[1], x.shape[2] + C = x.shape[3] + pe = create_uv_grid(pw, ph, aspect_ratio=W / H, dtype=x.dtype) + pe = position_grid_to_embed(pe, C) * ratio # (ph, pw, C) + pe = mx.expand_dims(pe, 0) # (1, ph, pw, C) + return x + pe + + @staticmethod + def _apply_activation(x: mx.array, activation: str) -> mx.array: + act = activation.lower() + if act == "exp": + return mx.exp(x) + if act == "expp1": + return mx.exp(x) + 1.0 + if act == "expm1": + return mx.exp(x) - 1.0 + if act == "relu": + return nn.relu(x) + if act == "sigmoid": + return mx.sigmoid(x) + if act == "softplus": + return mx.log(1.0 + mx.exp(x)) + if act == "tanh": + return mx.tanh(x) + return x # linear diff --git a/mlx_depth_anything_3/inference.py b/mlx_depth_anything_3/inference.py new file mode 100644 index 00000000..7d5a37f3 --- /dev/null +++ b/mlx_depth_anything_3/inference.py @@ -0,0 +1,168 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +"""High-level inference API for Depth Anything 3 on MLX.""" + +from __future__ import annotations + +from pathlib import Path + +import mlx.core as mx +import mlx.nn as nn +import numpy as np + +from mlx_depth_anything_3.model import DepthAnything3Net + + +# ImageNet normalization constants +IMAGENET_MEAN = np.array([0.485, 0.456, 0.406], dtype=np.float32) +IMAGENET_STD = np.array([0.229, 0.224, 0.225], dtype=np.float32) + + +class DepthAnything3: + """Depth Anything 3 inference wrapper for MLX. + + Supports: + - Multiple input images (variable number of views). + - Optional extrinsics/intrinsics conditioning. + - Dynamic spatial resolution (any multiple of 14). + - Returns predicted depth for each input image. + + Example:: + + da3 = DepthAnything3("da3-small") + da3.load_weights("weights.safetensors") + images = [np.array(img1), np.array(img2)] # uint8 or float32, HWC + result = da3.predict(images) + depth_maps = result["depth"] # (N, H, W) numpy array + """ + + def __init__(self, config_name: str = "da3-small"): + self.config_name = config_name + self.model = DepthAnything3Net(config_name) + + def load_weights(self, path: str | Path) -> None: + """Load model weights from a safetensors or npz file.""" + path = str(path) + if path.endswith(".safetensors"): + weights = mx.load(path) + elif path.endswith(".npz"): + weights = dict(mx.load(path)) + else: + raise ValueError(f"Unsupported weight format: {path}") + + # The weight keys use dot-separated paths matching the model structure + self.model.load_weights(list(weights.items())) + mx.eval(self.model.parameters()) + + def predict( + self, + images: list[np.ndarray], + extrinsics: np.ndarray | None = None, + intrinsics: np.ndarray | None = None, + process_res: int = 504, + ref_view_strategy: str = "saddle_balanced", + ) -> dict[str, np.ndarray]: + """Run depth prediction on a list of images. + + Args: + images: list of N images as numpy arrays (H, W, 3) uint8 or float32. + extrinsics: optional (N, 4, 4) world-to-camera matrices. + intrinsics: optional (N, 3, 3) camera intrinsic matrices. + process_res: processing resolution (must be multiple of 14). + ref_view_strategy: reference view selection strategy. + + Returns: + {"depth": (N, H, W) numpy float32 array} + """ + # Preprocess images + processed = self._preprocess_images(images, process_res) + x = mx.array(processed) # (1, N, H, W, 3) + + # Prepare camera matrices + ext_mx = None + int_mx = None + if extrinsics is not None: + ext_np = extrinsics.astype(np.float32) + if ext_np.ndim == 3: + ext_np = ext_np[np.newaxis] # (1, N, 4, 4) + ext_mx = mx.array(ext_np) + if intrinsics is not None: + int_np = intrinsics.astype(np.float32) + if int_np.ndim == 3: + int_np = int_np[np.newaxis] # (1, N, 3, 3) + int_mx = mx.array(int_np) + + # Forward pass + output = self.model(x, extrinsics=ext_mx, intrinsics=int_mx, + ref_view_strategy=ref_view_strategy) + mx.eval(output["depth"]) + + # Convert to numpy + depth = np.array(output["depth"][0]) # (N, H, W) + return {"depth": depth} + + def _preprocess_images( + self, + images: list[np.ndarray], + process_res: int, + ) -> np.ndarray: + """Preprocess a list of images to model input format. + + Args: + images: list of (H, W, 3) images. + process_res: target resolution (must be multiple of 14). + + Returns: + (1, N, H, W, 3) float32 array, ImageNet-normalised, NHWC. + """ + assert process_res % 14 == 0, f"process_res must be multiple of 14, got {process_res}" + + processed = [] + for img in images: + if img.dtype == np.uint8: + img = img.astype(np.float32) / 255.0 + # Resize to process_res x process_res + img = self._resize_image(img, process_res, process_res) + # ImageNet normalise + img = (img - IMAGENET_MEAN) / IMAGENET_STD + processed.append(img) + + stacked = np.stack(processed, axis=0) # (N, H, W, 3) + return stacked[np.newaxis] # (1, N, H, W, 3) + + @staticmethod + def _resize_image(img: np.ndarray, target_h: int, target_w: int) -> np.ndarray: + """Resize image using bilinear interpolation (numpy only, no OpenCV dependency).""" + h, w = img.shape[:2] + if h == target_h and w == target_w: + return img + + # Simple bilinear resize + y_ratio = h / target_h + x_ratio = w / target_w + + y_coords = np.arange(target_h).astype(np.float32) * y_ratio + x_coords = np.arange(target_w).astype(np.float32) * x_ratio + + y0 = np.clip(np.floor(y_coords).astype(int), 0, h - 1) + y1 = np.clip(y0 + 1, 0, h - 1) + x0 = np.clip(np.floor(x_coords).astype(int), 0, w - 1) + x1 = np.clip(x0 + 1, 0, w - 1) + + wy = y_coords - y0.astype(np.float32) + wx = x_coords - x0.astype(np.float32) + + top_left = img[y0][:, x0] + top_right = img[y0][:, x1] + bot_left = img[y1][:, x0] + bot_right = img[y1][:, x1] + + wy = wy[:, np.newaxis, np.newaxis] + wx = wx[np.newaxis, :, np.newaxis] + + result = (top_left * (1 - wy) * (1 - wx) + + top_right * (1 - wy) * wx + + bot_left * wy * (1 - wx) + + bot_right * wy * wx) + return result.astype(np.float32) diff --git a/mlx_depth_anything_3/layers.py b/mlx_depth_anything_3/layers.py new file mode 100644 index 00000000..af31f0ad --- /dev/null +++ b/mlx_depth_anything_3/layers.py @@ -0,0 +1,492 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Copyright (c) Meta Platforms, Inc. and affiliates. +# Licensed under the Apache License, Version 2.0 + +"""Core neural network layers for DA3 in MLX.""" + +from __future__ import annotations + +import math +from typing import Optional + +import mlx.core as mx +import mlx.nn as nn + + +# --------------------------------------------------------------------------- +# ModuleList helper +# --------------------------------------------------------------------------- +class ModuleList(nn.Module): + """Simple wrapper to register a list of modules in MLX.""" + + def __init__(self, modules: list[nn.Module]): + super().__init__() + for i, m in enumerate(modules): + setattr(self, str(i), m) + self._len = len(modules) + + def __getitem__(self, i: int) -> nn.Module: + if i < 0: + i += self._len + if i < 0 or i >= self._len: + raise IndexError("ModuleList index out of range") + return getattr(self, str(i)) + + def __len__(self) -> int: + return self._len + + def __iter__(self): + for i in range(self._len): + yield self[i] + + +# --------------------------------------------------------------------------- +# LayerScale +# --------------------------------------------------------------------------- +class LayerScale(nn.Module): + """Learnable per-channel scaling.""" + + def __init__(self, dim: int, init_values: float = 1e-5): + super().__init__() + self.gamma = mx.full((dim,), init_values) + + def __call__(self, x: mx.array) -> mx.array: + return x * self.gamma + + +# --------------------------------------------------------------------------- +# Mlp (GELU-based 2-layer MLP) +# --------------------------------------------------------------------------- +class Mlp(nn.Module): + def __init__( + self, + in_features: int, + hidden_features: int | None = None, + out_features: int | None = None, + act_layer: str = "gelu", + bias: bool = True, + **kwargs, + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + self.fc1 = nn.Linear(in_features, hidden_features, bias=bias) + self.fc2 = nn.Linear(hidden_features, out_features, bias=bias) + self._act = act_layer + + def __call__(self, x: mx.array) -> mx.array: + x = self.fc1(x) + x = _apply_act(x, self._act) + x = self.fc2(x) + return x + + +# --------------------------------------------------------------------------- +# SwiGLU FFN +# --------------------------------------------------------------------------- +class SwiGLUFFN(nn.Module): + """SwiGLU feed-forward network.""" + + def __init__( + self, + in_features: int, + hidden_features: int | None = None, + out_features: int | None = None, + bias: bool = True, + **kwargs, + ): + super().__init__() + out_features = out_features or in_features + hidden_features = hidden_features or in_features + hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 + self.w12 = nn.Linear(in_features, 2 * hidden_features, bias=bias) + self.w3 = nn.Linear(hidden_features, out_features, bias=bias) + + def __call__(self, x: mx.array) -> mx.array: + x12 = self.w12(x) + half = x12.shape[-1] // 2 + x1 = x12[..., :half] + x2 = x12[..., half:] + hidden = nn.silu(x1) * x2 + return self.w3(hidden) + + +# --------------------------------------------------------------------------- +# 2-D Rotary Position Embedding +# --------------------------------------------------------------------------- +class PositionGetter: + """Generates and caches 2-D spatial positions for patches.""" + + def __init__(self): + self._cache: dict[tuple[int, int], mx.array] = {} + + def __call__(self, batch_size: int, height: int, width: int) -> mx.array: + key = (height, width) + if key not in self._cache: + y = mx.arange(height) + x = mx.arange(width) + # cartesian product: (H*W, 2) with (y, x) + yy = mx.repeat(y, width) + xx = mx.tile(x, (height,)) + self._cache[key] = mx.stack([yy, xx], axis=-1) + cached = self._cache[key] + return mx.broadcast_to(mx.expand_dims(cached, 0), (batch_size, height * width, 2)) + + +class RotaryPositionEmbedding2D(nn.Module): + """2-D RoPE that applies separate 1-D rotary embeddings for vertical/horizontal.""" + + def __init__(self, frequency: float = 100.0): + super().__init__() + self.base_frequency = frequency + self._freq_cache: dict = {} + + def _get_freq(self, dim: int, seq_len: int, dtype): + key = (dim, seq_len, dtype) + if key not in self._freq_cache: + exponents = mx.arange(0, dim, 2).astype(mx.float32) / dim + inv_freq = 1.0 / (self.base_frequency ** exponents) + positions = mx.arange(seq_len).astype(mx.float32) + angles = mx.expand_dims(positions, 1) * mx.expand_dims(inv_freq, 0) + angles = mx.concatenate([angles, angles], axis=-1).astype(dtype) + self._freq_cache[key] = (mx.cos(angles), mx.sin(angles)) + return self._freq_cache[key] + + @staticmethod + def _rotate(x: mx.array) -> mx.array: + half = x.shape[-1] // 2 + x1 = x[..., :half] + x2 = x[..., half:] + return mx.concatenate([-x2, x1], axis=-1) + + def _apply_1d(self, tokens, positions, cos_comp, sin_comp): + # positions: (B, N) int indices + cos = cos_comp[positions][:, None, :, :] # (B, 1, N, D) + sin = sin_comp[positions][:, None, :, :] + return tokens * cos + self._rotate(tokens) * sin + + def __call__(self, tokens: mx.array, positions: mx.array) -> mx.array: + """Apply 2D RoPE. + + Args: + tokens: (B, heads, N, dim) + positions: (B, N, 2) with [y, x] coords + """ + feat_dim = tokens.shape[-1] // 2 + max_pos = int(mx.max(positions).item()) + 1 + cos_comp, sin_comp = self._get_freq(feat_dim, max_pos, tokens.dtype) + + vert = tokens[..., :feat_dim] + horiz = tokens[..., feat_dim:] + + pos_y = positions[..., 0].astype(mx.int32) + pos_x = positions[..., 1].astype(mx.int32) + + vert = self._apply_1d(vert, pos_y, cos_comp, sin_comp) + horiz = self._apply_1d(horiz, pos_x, cos_comp, sin_comp) + return mx.concatenate([vert, horiz], axis=-1) + + +# --------------------------------------------------------------------------- +# Multi-head Attention +# --------------------------------------------------------------------------- +class Attention(nn.Module): + def __init__( + self, + dim: int, + num_heads: int = 8, + qkv_bias: bool = False, + proj_bias: bool = True, + qk_norm: bool = False, + rope: RotaryPositionEmbedding2D | None = None, + **kwargs, + ): + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim ** -0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = nn.LayerNorm(self.head_dim) if qk_norm else None + self.k_norm = nn.LayerNorm(self.head_dim) if qk_norm else None + self.proj = nn.Linear(dim, dim, bias=proj_bias) + self.rope = rope + + def __call__(self, x: mx.array, pos: mx.array | None = None, + attn_mask: mx.array | None = None) -> mx.array: + B, N, C = x.shape + qkv = self.qkv(x) # (B, N, 3*C) + qkv = mx.reshape(qkv, (B, N, 3, self.num_heads, self.head_dim)) + qkv = mx.transpose(qkv, axes=(2, 0, 3, 1, 4)) # (3, B, heads, N, head_dim) + q = qkv[0] + k = qkv[1] + v = qkv[2] + + if self.q_norm is not None: + q = self.q_norm(q) + if self.k_norm is not None: + k = self.k_norm(k) + if self.rope is not None and pos is not None: + q = self.rope(q, pos) + k = self.rope(k, pos) + + # Scaled dot-product attention + q = q * self.scale + attn = q @ mx.transpose(k, axes=(0, 1, 3, 2)) # (B, heads, N, N) + if attn_mask is not None: + # attn_mask: (B, N, N) boolean. True = attend, False = block. + mask = mx.expand_dims(attn_mask, 1) + mask = mx.broadcast_to(mask, attn.shape) + attn = attn + mx.where(mask, mx.zeros_like(attn), mx.full(attn.shape, -1e9)) + attn = mx.softmax(attn, axis=-1) + out = attn @ v # (B, heads, N, head_dim) + out = mx.transpose(out, axes=(0, 2, 1, 3)) # (B, N, heads, head_dim) + out = mx.reshape(out, (B, N, C)) + out = self.proj(out) + return out + + +# --------------------------------------------------------------------------- +# Transformer Block +# --------------------------------------------------------------------------- +class Block(nn.Module): + def __init__( + self, + dim: int, + num_heads: int, + mlp_ratio: float = 4.0, + qkv_bias: bool = False, + proj_bias: bool = True, + ffn_bias: bool = True, + init_values: float | None = None, + ffn_layer: str = "mlp", + qk_norm: bool = False, + rope: RotaryPositionEmbedding2D | None = None, + **kwargs, + ): + super().__init__() + self.norm1 = nn.LayerNorm(dim, eps=1e-6) + self.attn = Attention( + dim, num_heads=num_heads, qkv_bias=qkv_bias, + proj_bias=proj_bias, qk_norm=qk_norm, rope=rope, + ) + self.ls1 = LayerScale(dim, init_values) if init_values else None + self.norm2 = nn.LayerNorm(dim, eps=1e-6) + mlp_hidden = int(dim * mlp_ratio) + if ffn_layer == "swiglu": + self.mlp = SwiGLUFFN(in_features=dim, hidden_features=mlp_hidden, + out_features=dim, bias=ffn_bias) + else: + self.mlp = Mlp(in_features=dim, hidden_features=mlp_hidden, + out_features=dim, bias=ffn_bias) + self.ls2 = LayerScale(dim, init_values) if init_values else None + + def __call__(self, x: mx.array, pos: mx.array | None = None, + attn_mask: mx.array | None = None) -> mx.array: + attn_out = self.attn(self.norm1(x), pos=pos, attn_mask=attn_mask) + if self.ls1 is not None: + attn_out = self.ls1(attn_out) + x = x + attn_out + ffn_out = self.mlp(self.norm2(x)) + if self.ls2 is not None: + ffn_out = self.ls2(ffn_out) + x = x + ffn_out + return x + + +# --------------------------------------------------------------------------- +# Patch Embedding +# --------------------------------------------------------------------------- +class PatchEmbed(nn.Module): + """2-D image to patch embedding: (B, H, W, C) -> (B, N, D). + + Note: MLX Conv2d uses channels-last (NHWC) format. + """ + + def __init__(self, img_size: int = 224, patch_size: int = 16, + in_chans: int = 3, embed_dim: int = 768): + super().__init__() + self.patch_size = (patch_size, patch_size) if isinstance(patch_size, int) else patch_size + grid_h = img_size // self.patch_size[0] + grid_w = img_size // self.patch_size[1] + self.num_patches = grid_h * grid_w + self.proj = nn.Conv2d(in_chans, embed_dim, + kernel_size=self.patch_size[0], stride=self.patch_size[0]) + + def __call__(self, x: mx.array) -> mx.array: + # x: (B, H, W, C) in NHWC + x = self.proj(x) # (B, H/p, W/p, D) + B = x.shape[0] + x = mx.reshape(x, (B, -1, x.shape[-1])) # (B, N, D) + return x + + +# --------------------------------------------------------------------------- +# Camera Encoder Attention + Block (simpler variant from model/utils/) +# --------------------------------------------------------------------------- +class CamAttention(nn.Module): + """Attention for the camera encoder (simpler: always uses scaled_dot_product).""" + + def __init__(self, dim: int, num_heads: int = 8, qkv_bias: bool = True, + proj_bias: bool = True, qk_norm: bool = False): + super().__init__() + self.num_heads = num_heads + self.head_dim = dim // num_heads + self.scale = self.head_dim ** -0.5 + self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias) + self.q_norm = nn.LayerNorm(self.head_dim, eps=1e-6) if qk_norm else None + self.k_norm = nn.LayerNorm(self.head_dim, eps=1e-6) if qk_norm else None + self.proj = nn.Linear(dim, dim, bias=proj_bias) + + def __call__(self, x: mx.array) -> mx.array: + B, N, C = x.shape + qkv = self.qkv(x) + qkv = mx.reshape(qkv, (B, N, 3, self.num_heads, self.head_dim)) + qkv = mx.transpose(qkv, axes=(2, 0, 3, 1, 4)) + q, k, v = qkv[0], qkv[1], qkv[2] + if self.q_norm is not None: + q = self.q_norm(q) + if self.k_norm is not None: + k = self.k_norm(k) + q = q * self.scale + attn = q @ mx.transpose(k, axes=(0, 1, 3, 2)) + attn = mx.softmax(attn, axis=-1) + out = attn @ v + out = mx.transpose(out, axes=(0, 2, 1, 3)) + out = mx.reshape(out, (B, N, C)) + return self.proj(out) + + +class CamBlock(nn.Module): + """Transformer block for the camera encoder (no drop-path).""" + + def __init__(self, dim: int, num_heads: int, mlp_ratio: float = 4.0, + init_values: float | None = None): + super().__init__() + self.norm1 = nn.LayerNorm(dim, eps=1e-6) + self.attn = CamAttention(dim, num_heads=num_heads, qkv_bias=True) + self.ls1 = LayerScale(dim, init_values) if init_values else None + self.norm2 = nn.LayerNorm(dim, eps=1e-6) + hidden = int(dim * mlp_ratio) + self.mlp = Mlp(in_features=dim, hidden_features=hidden, out_features=dim, bias=True) + self.ls2 = LayerScale(dim, init_values) if init_values else None + + def __call__(self, x: mx.array) -> mx.array: + attn_out = self.attn(self.norm1(x)) + if self.ls1 is not None: + attn_out = self.ls1(attn_out) + x = x + attn_out + ffn_out = self.mlp(self.norm2(x)) + if self.ls2 is not None: + ffn_out = self.ls2(ffn_out) + x = x + ffn_out + return x + + +# --------------------------------------------------------------------------- +# Positional embedding helpers (used in DPT head) +# --------------------------------------------------------------------------- +def create_uv_grid(width: int, height: int, aspect_ratio: float | None = None, + dtype=mx.float32) -> mx.array: + """Create normalized UV grid (H, W, 2).""" + if aspect_ratio is None: + aspect_ratio = float(width) / float(height) + diag = (aspect_ratio ** 2 + 1.0) ** 0.5 + sx = aspect_ratio / diag + sy = 1.0 / diag + lx = -sx * (width - 1) / width + rx = sx * (width - 1) / width + ty = -sy * (height - 1) / height + by = sy * (height - 1) / height + xs = mx.linspace(lx, rx, width).astype(dtype) + ys = mx.linspace(ty, by, height).astype(dtype) + # meshgrid: (W,) x (H,) -> (H, W) + xx = mx.broadcast_to(mx.expand_dims(xs, 0), (height, width)) + yy = mx.broadcast_to(mx.expand_dims(ys, 1), (height, width)) + return mx.stack([xx, yy], axis=-1) + + +def make_sincos_pos_embed(embed_dim: int, pos: mx.array, + omega_0: float = 100.0) -> mx.array: + """1-D sinusoidal positional embedding.""" + omega = mx.arange(embed_dim // 2).astype(mx.float32) / (embed_dim / 2.0) + omega = 1.0 / (omega_0 ** omega) + pos_flat = mx.reshape(pos, (-1,)).astype(mx.float32) + out = mx.expand_dims(pos_flat, 1) * mx.expand_dims(omega, 0) + return mx.concatenate([mx.sin(out), mx.cos(out)], axis=1).astype(mx.float32) + + +def position_grid_to_embed(pos_grid: mx.array, embed_dim: int, + omega_0: float = 100.0) -> mx.array: + """Convert 2D position grid (H, W, 2) to sinusoidal embeddings (H, W, C).""" + H, W, _ = pos_grid.shape + pos_flat = mx.reshape(pos_grid, (-1, 2)) + emb_x = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 0], omega_0) + emb_y = make_sincos_pos_embed(embed_dim // 2, pos_flat[:, 1], omega_0) + emb = mx.concatenate([emb_x, emb_y], axis=-1) + return mx.reshape(emb, (H, W, embed_dim)) + + +# --------------------------------------------------------------------------- +# Bilinear interpolation helper +# --------------------------------------------------------------------------- +def bilinear_interpolate(x: mx.array, target_h: int, target_w: int, + align_corners: bool = True) -> mx.array: + """Bilinear interpolation for NHWC tensors. + + Args: + x: (N, H, W, C) input. + target_h, target_w: target spatial dimensions. + + Returns: + (N, target_h, target_w, C) interpolated tensor. + """ + N, H, W, C = x.shape + if H == target_h and W == target_w: + return x + + if align_corners and target_h > 1 and target_w > 1: + y_coords = mx.linspace(0, H - 1, target_h).astype(mx.float32) + x_coords = mx.linspace(0, W - 1, target_w).astype(mx.float32) + else: + y_coords = (mx.arange(target_h).astype(mx.float32) + 0.5) * H / target_h - 0.5 + x_coords = (mx.arange(target_w).astype(mx.float32) + 0.5) * W / target_w - 0.5 + + y_coords = mx.clip(y_coords, 0, H - 1) + x_coords = mx.clip(x_coords, 0, W - 1) + + y0 = mx.floor(y_coords).astype(mx.int32) + y1 = mx.minimum(y0 + 1, H - 1) + x0 = mx.floor(x_coords).astype(mx.int32) + x1 = mx.minimum(x0 + 1, W - 1) + + wy = y_coords - y0.astype(mx.float32) + wx = x_coords - x0.astype(mx.float32) + + # (target_h, target_w) grid indices + # Gather corners: x[:, y0[i], x0[j], :] + top_left = x[:, y0][:, :, x0] # (N, target_h, target_w, C) + top_right = x[:, y0][:, :, x1] + bot_left = x[:, y1][:, :, x0] + bot_right = x[:, y1][:, :, x1] + + wy = mx.reshape(wy, (1, target_h, 1, 1)) + wx = mx.reshape(wx, (1, 1, target_w, 1)) + + out = (top_left * (1 - wy) * (1 - wx) + + top_right * (1 - wy) * wx + + bot_left * wy * (1 - wx) + + bot_right * wy * wx) + return out + + +# --------------------------------------------------------------------------- +# Activation helper +# --------------------------------------------------------------------------- +def _apply_act(x: mx.array, act: str = "gelu") -> mx.array: + if act == "gelu": + return nn.gelu(x) + elif act == "relu": + return nn.relu(x) + elif act == "silu": + return nn.silu(x) + return x diff --git a/mlx_depth_anything_3/model.py b/mlx_depth_anything_3/model.py new file mode 100644 index 00000000..e014b510 --- /dev/null +++ b/mlx_depth_anything_3/model.py @@ -0,0 +1,147 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +"""Depth Anything 3 main model for MLX – depth-only output.""" + +from __future__ import annotations + +import mlx.core as mx +import mlx.nn as nn + +from mlx_depth_anything_3.backbone import DinoV2 +from mlx_depth_anything_3.camera import CameraDec, CameraEnc +from mlx_depth_anything_3.dpt import DualDPT +from mlx_depth_anything_3.transform import affine_inverse, pose_encoding_to_extri_intri + + +# --------------------------------------------------------------------------- +# Model configurations (matching PyTorch YAML configs) +# --------------------------------------------------------------------------- +MODEL_CONFIGS = { + "da3-small": dict( + backbone=dict(name="vits", out_layers=[5, 7, 9, 11], alt_start=4, + qknorm_start=4, rope_start=4, cat_token=True), + head=dict(dim_in=768, output_dim=2, features=64, + out_channels=[48, 96, 192, 384]), + cam_enc=dict(dim_out=384), + cam_dec=dict(dim_in=768), + ), + "da3-base": dict( + backbone=dict(name="vitb", out_layers=[5, 7, 9, 11], alt_start=4, + qknorm_start=4, rope_start=4, cat_token=True), + head=dict(dim_in=1536, output_dim=2, features=128, + out_channels=[96, 192, 384, 768]), + cam_enc=dict(dim_out=768), + cam_dec=dict(dim_in=1536), + ), + "da3-large": dict( + backbone=dict(name="vitl", out_layers=[11, 15, 19, 23], alt_start=8, + qknorm_start=8, rope_start=8, cat_token=True), + head=dict(dim_in=2048, output_dim=2, features=256, + out_channels=[256, 512, 1024, 1024]), + cam_enc=dict(dim_out=1024), + cam_dec=dict(dim_in=2048), + ), + "da3-giant": dict( + backbone=dict(name="vitg", out_layers=[19, 27, 33, 39], alt_start=13, + qknorm_start=13, rope_start=13, cat_token=True), + head=dict(dim_in=3072, output_dim=2, features=256, + out_channels=[256, 512, 1024, 1024]), + cam_enc=dict(dim_out=1536), + cam_dec=dict(dim_in=3072), + ), +} + + +class DepthAnything3Net(nn.Module): + """Depth Anything 3 network for depth estimation. + + Supports: + - Multiple input images (B=1, S=1..N views). + - Optional extrinsics/intrinsics conditioning. + - Dynamic spatial resolution (must be multiple of 14). + - Outputs only predicted depth (and confidence) for each image. + """ + + PATCH_SIZE = 14 + + def __init__(self, config_name: str = "da3-small"): + super().__init__() + cfg = MODEL_CONFIGS[config_name] + + # Backbone + bcfg = cfg["backbone"] + self.backbone = DinoV2( + name=bcfg["name"], + out_layers=bcfg["out_layers"], + alt_start=bcfg["alt_start"], + qknorm_start=bcfg["qknorm_start"], + rope_start=bcfg["rope_start"], + cat_token=bcfg["cat_token"], + ) + + # Depth head + hcfg = cfg["head"] + self.head = DualDPT( + dim_in=hcfg["dim_in"], + output_dim=hcfg["output_dim"], + features=hcfg["features"], + out_channels=hcfg["out_channels"], + ) + + # Camera encoder (for extrinsics/intrinsics conditioning) + ecfg = cfg["cam_enc"] + self.cam_enc = CameraEnc(dim_out=ecfg["dim_out"]) + + # Camera decoder (for pose prediction; kept for weight compat) + dcfg = cfg["cam_dec"] + self.cam_dec = CameraDec(dim_in=dcfg["dim_in"]) + + def __call__( + self, + x: mx.array, + extrinsics: mx.array | None = None, + intrinsics: mx.array | None = None, + ref_view_strategy: str = "saddle_balanced", + ) -> dict[str, mx.array]: + """Forward pass – returns predicted depth for each view. + + Args: + x: (B, S, H, W, 3) input images in NHWC, float32, ImageNet-normalised. + extrinsics: (B, S, 4, 4) optional world-to-camera matrices. + intrinsics: (B, S, 3, 3) optional camera intrinsic matrices. + ref_view_strategy: reference view selection strategy. + + Returns: + {"depth": (B, S, H, W), "depth_conf": (B, S, H, W)} + """ + B, S, H, W, C = x.shape + + # Camera encoding + cam_token = None + if extrinsics is not None and intrinsics is not None: + cam_token = self.cam_enc( + extrinsics.astype(mx.float32), + intrinsics.astype(mx.float32), + (H, W), + ) + + # Backbone feature extraction + feats, _ = self.backbone( + x, cam_token=cam_token, + ref_view_strategy=ref_view_strategy, + ) + + # Depth head + output = self.head(feats, H, W, patch_start_idx=0) + + # Refined camera parameters + # Take the camera token from the last backbone layer + cam_token_refined = feats[-1][1] + pose_encoding = self.cam_dec(cam_token_refined) + refined_ext, refined_int = pose_encoding_to_extri_intri(pose_encoding, (H, W)) + + output["intrinsics"] = refined_int + output["extrinsics"] = refined_ext + + return output diff --git a/mlx_depth_anything_3/reference_view.py b/mlx_depth_anything_3/reference_view.py new file mode 100644 index 00000000..c1b84b06 --- /dev/null +++ b/mlx_depth_anything_3/reference_view.py @@ -0,0 +1,138 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +"""Reference view selection strategies for multi-view depth estimation (MLX).""" + +import mlx.core as mx + + +THRESH_FOR_REF_SELECTION = 3 + + +def select_reference_view( + x: mx.array, + strategy: str = "saddle_balanced", +) -> mx.array: + """Select a reference view from multiple views. + + Args: + x: (B, S, N, C) feature tensor. + strategy: "first", "middle", "saddle_balanced", or "saddle_sim_range". + + Returns: + (B,) tensor of selected view indices. + """ + B, S, N, C = x.shape + if S <= 1: + return mx.zeros((B,), dtype=mx.int32) + if strategy == "first": + return mx.zeros((B,), dtype=mx.int32) + if strategy == "middle": + return mx.full((B,), S // 2, dtype=mx.int32) + + # Feature-based strategies: normalized class tokens + cls_feat = x[:, :, 0] # (B, S, C) + cls_norm = mx.sqrt(mx.sum(cls_feat * cls_feat, axis=-1, keepdims=True) + 1e-12) + img_class_feat = cls_feat / cls_norm # (B, S, C) + + if strategy == "saddle_balanced": + sim = img_class_feat @ mx.transpose(img_class_feat, axes=(0, 2, 1)) # (B, S, S) + eye = mx.broadcast_to(mx.expand_dims(mx.eye(S), 0), sim.shape) + sim_no_diag = sim - eye + sim_score = mx.sum(sim_no_diag, axis=-1) / (S - 1) + + feat_norm = mx.sqrt(mx.sum(cls_feat * cls_feat, axis=-1) + 1e-12) + feat_var = mx.var(img_class_feat, axis=-1) + + def normalize_metric(m): + mn = mx.min(m, axis=1, keepdims=True) + mx_val = mx.max(m, axis=1, keepdims=True) + return (m - mn) / (mx_val - mn + 1e-8) + + sim_n = normalize_metric(sim_score) + norm_n = normalize_metric(feat_norm) + var_n = normalize_metric(feat_var) + + balance = mx.abs(sim_n - 0.5) + mx.abs(norm_n - 0.5) + mx.abs(var_n - 0.5) + return mx.argmin(balance, axis=1) + + elif strategy == "saddle_sim_range": + sim = img_class_feat @ mx.transpose(img_class_feat, axes=(0, 2, 1)) + eye = mx.broadcast_to(mx.expand_dims(mx.eye(S), 0), sim.shape) + sim_no_diag = sim - eye + sim_max = mx.max(sim_no_diag, axis=-1) + sim_min = mx.min(sim_no_diag, axis=-1) + sim_range = sim_max - sim_min + return mx.argmax(sim_range, axis=1) + + raise ValueError(f"Unknown strategy: {strategy}") + + +def reorder_by_reference(x: mx.array, b_idx: mx.array) -> mx.array: + """Reorder views so that the reference view is at index 0. + + Args: + x: (B, S, ...) tensor. + b_idx: (B,) reference view indices. + + Returns: + Reordered tensor. + """ + B, S = x.shape[0], x.shape[1] + if S <= 1: + return x + + positions = mx.broadcast_to(mx.arange(S).reshape(1, S), (B, S)) + b_idx_expanded = mx.expand_dims(b_idx, 1) + + reorder_indices = mx.where( + (positions > 0) & (positions <= b_idx_expanded), + positions - 1, + positions, + ) + # Set position 0 to ref_idx + col0 = mx.expand_dims(b_idx, 1) + reorder_indices = mx.concatenate([col0, reorder_indices[:, 1:]], axis=1) + # Fix: for positions > 0 and <= b_idx, we already set positions-1, which is correct + # For position 0, we set b_idx, which is correct + + # Gather using take_along_axis + # Expand indices to match x's extra dimensions + idx_shape = list(reorder_indices.shape) + [1] * (x.ndim - 2) + idx_expanded = mx.reshape(reorder_indices, idx_shape) + idx_expanded = mx.broadcast_to(idx_expanded, list(reorder_indices.shape) + list(x.shape[2:])) + return mx.take_along_axis(x, idx_expanded, axis=1) + + +def restore_original_order(x: mx.array, b_idx: mx.array) -> mx.array: + """Restore original view order after reference-first reordering. + + Args: + x: (B, S, ...) reordered tensor. + b_idx: (B,) original reference view indices. + + Returns: + Tensor with original order restored. + """ + B, S = x.shape[0], x.shape[1] + if S <= 1: + return x + + positions = mx.broadcast_to(mx.arange(S).reshape(1, S), (B, S)) + b_idx_expanded = mx.expand_dims(b_idx, 1) + + restore_indices = mx.where( + positions < b_idx_expanded, + positions + 1, + positions, + ) + # Position = ref_idx comes from position 0 + # Use a mask to set the right position + ref_mask = positions == b_idx_expanded + restore_indices = mx.where(ref_mask, mx.zeros_like(restore_indices), restore_indices) + + # Gather using take_along_axis + idx_shape = list(restore_indices.shape) + [1] * (x.ndim - 2) + idx_expanded = mx.reshape(restore_indices, idx_shape) + idx_expanded = mx.broadcast_to(idx_expanded, list(restore_indices.shape) + list(x.shape[2:])) + return mx.take_along_axis(x, idx_expanded, axis=1) diff --git a/mlx_depth_anything_3/transform.py b/mlx_depth_anything_3/transform.py new file mode 100644 index 00000000..4ca142ec --- /dev/null +++ b/mlx_depth_anything_3/transform.py @@ -0,0 +1,186 @@ +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +"""Quaternion math and camera pose encoding/decoding utilities for MLX.""" + +import mlx.core as mx + + +def extri_intri_to_pose_encoding( + extrinsics: mx.array, + intrinsics: mx.array, + image_size_hw: tuple[int, int], +) -> mx.array: + """Convert camera extrinsics (c2w) and intrinsics to a compact 9-D pose encoding. + + Args: + extrinsics: (B, S, 4, 4) camera-to-world matrices. + intrinsics: (B, S, 3, 3) camera intrinsic matrices. + image_size_hw: (H, W) of the images. + + Returns: + (B, S, 9) pose encoding: [T_x, T_y, T_z, q_x, q_y, q_z, q_w, fov_h, fov_w]. + """ + R = extrinsics[:, :, :3, :3] + T = extrinsics[:, :, :3, 3] + quat = mat_to_quat(R) + H, W = image_size_hw + fov_h = 2.0 * mx.arctan(mx.array(H / 2.0) / intrinsics[:, :, 1, 1]) + fov_w = 2.0 * mx.arctan(mx.array(W / 2.0) / intrinsics[:, :, 0, 0]) + pose_encoding = mx.concatenate( + [T, quat, mx.expand_dims(fov_h, -1), mx.expand_dims(fov_w, -1)], axis=-1 + ).astype(mx.float32) + return pose_encoding + + +def pose_encoding_to_extri_intri( + pose_encoding: mx.array, + image_size_hw: tuple[int, int], +) -> tuple[mx.array, mx.array]: + """Convert 9-D pose encoding back to extrinsics and intrinsics. + + Returns: + extrinsics: (B, S, 3, 4) + intrinsics: (B, S, 3, 3) + """ + T = pose_encoding[..., :3] + quat = pose_encoding[..., 3:7] + fov_h = pose_encoding[..., 7] + fov_w = pose_encoding[..., 8] + R = quat_to_mat(quat) + extrinsics = mx.concatenate([R, mx.expand_dims(T, -1)], axis=-1) + H, W = image_size_hw + fy = (H / 2.0) / mx.clip(mx.tan(fov_h / 2.0), a_min=1e-6, a_max=None) + fx = (W / 2.0) / mx.clip(mx.tan(fov_w / 2.0), a_min=1e-6, a_max=None) + B_dim, S_dim = pose_encoding.shape[:2] + intrinsics = mx.zeros((*pose_encoding.shape[:2], 3, 3)) + # Build intrinsics row by row + row0 = mx.stack([fx, mx.zeros_like(fx), mx.full_like(fx, W / 2.0)], axis=-1) + row1 = mx.stack([mx.zeros_like(fy), fy, mx.full_like(fy, H / 2.0)], axis=-1) + row2 = mx.stack( + [mx.zeros_like(fx), mx.zeros_like(fx), mx.ones_like(fx)], axis=-1 + ) + intrinsics = mx.stack([row0, row1, row2], axis=-2) + return extrinsics, intrinsics + + +def quat_to_mat(quaternions: mx.array) -> mx.array: + """Quaternion (XYZW, scalar-last) to rotation matrix. + + Args: + quaternions: (..., 4) in [i, j, k, r] order. + + Returns: + (..., 3, 3) rotation matrices. + """ + i = quaternions[..., 0] + j = quaternions[..., 1] + k = quaternions[..., 2] + r = quaternions[..., 3] + two_s = 2.0 / mx.sum(quaternions * quaternions, axis=-1) + o = mx.stack( + [ + 1 - two_s * (j * j + k * k), + two_s * (i * j - k * r), + two_s * (i * k + j * r), + two_s * (i * j + k * r), + 1 - two_s * (i * i + k * k), + two_s * (j * k - i * r), + two_s * (i * k - j * r), + two_s * (j * k + i * r), + 1 - two_s * (i * i + j * j), + ], + axis=-1, + ) + return mx.reshape(o, (*quaternions.shape[:-1], 3, 3)) + + +def mat_to_quat(matrix: mx.array) -> mx.array: + """Rotation matrix to quaternion (XYZW, scalar-last). + + Args: + matrix: (..., 3, 3) rotation matrices. + + Returns: + (..., 4) quaternions in [i, j, k, r] order. + """ + batch_dim = matrix.shape[:-2] + flat = mx.reshape(matrix, (*batch_dim, 9)) + m00 = flat[..., 0] + m01 = flat[..., 1] + m02 = flat[..., 2] + m10 = flat[..., 3] + m11 = flat[..., 4] + m12 = flat[..., 5] + m20 = flat[..., 6] + m21 = flat[..., 7] + m22 = flat[..., 8] + + q_abs = _sqrt_positive_part( + mx.stack( + [ + 1.0 + m00 + m11 + m22, + 1.0 + m00 - m11 - m22, + 1.0 - m00 + m11 - m22, + 1.0 - m00 - m11 + m22, + ], + axis=-1, + ) + ) + + quat_by_rijk = mx.stack( + [ + mx.stack([q_abs[..., 0] ** 2, m21 - m12, m02 - m20, m10 - m01], axis=-1), + mx.stack([m21 - m12, q_abs[..., 1] ** 2, m10 + m01, m02 + m20], axis=-1), + mx.stack([m02 - m20, m10 + m01, q_abs[..., 2] ** 2, m12 + m21], axis=-1), + mx.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3] ** 2], axis=-1), + ], + axis=-2, + ) + + flr = mx.array(0.1) + denom = 2.0 * mx.maximum(mx.expand_dims(q_abs, -1), flr) + quat_candidates = quat_by_rijk / denom + + # Select best-conditioned quaternion + best_idx = mx.argmax(q_abs, axis=-1) + # Manual one-hot: compare each position with best_idx + indices = mx.arange(4) + # Broadcast: best_idx (...,1) vs indices (4,) + one_hot = (mx.expand_dims(best_idx, -1) == indices).astype(q_abs.dtype) + mask = one_hot > 0.5 + # Gather the selected quaternion for each batch element + out = mx.sum(quat_candidates * mx.expand_dims(mask, -1), axis=-2) + + # Convert from [r, i, j, k] to [i, j, k, r] + out = mx.concatenate([out[..., 1:2], out[..., 2:3], out[..., 3:4], out[..., 0:1]], axis=-1) + out = standardize_quaternion(out) + return out + + +def _sqrt_positive_part(x: mx.array) -> mx.array: + """sqrt(max(0, x)) with zero subgradient at x=0.""" + return mx.sqrt(mx.maximum(x, 0.0)) + + +def standardize_quaternion(quaternions: mx.array) -> mx.array: + """Ensure the real (last) component of the quaternion is non-negative.""" + return mx.where(quaternions[..., 3:4] < 0, -quaternions, quaternions) + + +def affine_inverse(A: mx.array) -> mx.array: + """Efficient inverse of an affine (rigid-body) transformation matrix. + + Args: + A: (..., 4, 4) affine matrix. + + Returns: + (..., 4, 4) inverse matrix. + """ + R = A[..., :3, :3] + T = A[..., :3, 3:] + P = A[..., 3:, :] + Rt = mx.transpose(R, axes=list(range(R.ndim - 2)) + [-1, -2]) + return mx.concatenate( + [mx.concatenate([Rt, -(Rt @ T)], axis=-1), P], axis=-2 + ) diff --git a/run_pt.py b/run_pt.py new file mode 100644 index 00000000..b6292c82 --- /dev/null +++ b/run_pt.py @@ -0,0 +1,62 @@ +import sys +import os +import torch +import numpy as np +from PIL import Image +from omegaconf import OmegaConf + +sys.path.append("/Users/aweise/dev/sunscape2/da3conv/Depth-Anything-3-coreml/src") +from depth_anything_3.cfg import create_object + +def load_image(path, size=(518, 518)): + img = Image.open(path).convert('RGB').resize(size) + img_np = np.array(img).astype(np.float32) / 255.0 + # normalize using ImageNet mean/std + mean = np.array([0.485, 0.456, 0.406], dtype=np.float32) + std = np.array([0.229, 0.224, 0.225], dtype=np.float32) + img_np = (img_np - mean) / std + return img_np + +def run_pytorch(): + config_name = "da3-base" + weights_path = "/Users/aweise/.cache/huggingface/hub/models--depth-anything--DA3-BASE/snapshots/f4a6c9b3c95e41c82048423d3493a81ec3fa810e/model.safetensors" + + cfg_path = f"/Users/aweise/dev/sunscape2/da3conv/Depth-Anything-3-coreml/src/depth_anything_3/configs/{config_name}.yaml" + cfg = OmegaConf.load(cfg_path) + model = create_object(cfg) + model.eval() + + from safetensors.torch import load_file + state_dict = load_file(weights_path) + new_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("model."): + new_state_dict[k[6:]] = v + else: + new_state_dict[k] = v + model.load_state_dict(new_state_dict, strict=False) + + dir_path = "/Users/aweise/dev/sunscape2/sunscape2/DepthAnything3MLX/Tests/DepthAnything3MLXTests/" + img1 = load_image(os.path.join(dir_path, "sv1.png")) + img2 = load_image(os.path.join(dir_path, "sv2.png")) + img3 = load_image(os.path.join(dir_path, "sv3.png")) + + # images_np: (1, N, H, W, 3) + images_np = np.stack([img1, img2, img3])[None, ...] + images_np.astype(np.float32).tofile(os.path.join(dir_path, "input_images.bin")) + np.save(os.path.join(dir_path, "input_images.npy"), images_np.astype(np.float32)) + print("Saved input_images.bin and input_images.npy") + + images_pt = torch.from_numpy(images_np).permute(0, 1, 4, 2, 3).float() + print("Running pytorch inference with shape", images_pt.shape) + + with torch.no_grad(): + output = model(images_pt, extrinsics=None, intrinsics=None, export_feat_layers=[], infer_gs=False, use_ray_pose=False) + + depth = output["depth"].numpy()[0] # (N, H, W) + depth.astype(np.float32).tofile(os.path.join(dir_path, "expected_depth.bin")) + np.save(os.path.join(dir_path, "expected_depth.npy"), depth.astype(np.float32)) + print("Saved expected_depth.bin and expected_depth.npy") + +if __name__ == "__main__": + run_pytorch() diff --git a/verify_parity.py b/verify_parity.py new file mode 100644 index 00000000..811e3320 --- /dev/null +++ b/verify_parity.py @@ -0,0 +1,242 @@ +#!/usr/bin/env python3 +# Copyright (c) 2025 ByteDance Ltd. and/or its affiliates +# Licensed under the Apache License, Version 2.0 + +"""Verify parity between PyTorch and MLX Depth Anything 3 models. + +This script: +1. Creates both PyTorch and MLX models with matching configurations. +2. Generates identical random inputs (images, extrinsics, intrinsics). +3. Runs forward passes on both models. +4. Compares depth outputs and reports max/mean absolute differences. + +Usage: + python verify_parity.py --model da3-small [--weights pytorch_weights.safetensors] + +If --weights is provided, it loads the PyTorch weights, converts them to MLX, +and runs both models with real weights. Without --weights, it initialises +both models from scratch with identical random weights (useful for checking +the architecture mapping). +""" + +from __future__ import annotations + +import argparse +import sys + +import numpy as np + +# ---- helpers --------------------------------------------------------------- + +def make_random_input( + num_views: int = 2, + height: int = 504, + width: int = 504, + seed: int = 42, + with_cameras: bool = True, +): + """Generate deterministic random inputs for both frameworks.""" + rng = np.random.RandomState(seed) + images = rng.randn(1, num_views, height, width, 3).astype(np.float32) + + extrinsics = None + intrinsics = None + if with_cameras: + # Simple extrinsics: identity + small perturbation + extrinsics = np.tile(np.eye(4, dtype=np.float32), (1, num_views, 1, 1)) + for i in range(num_views): + extrinsics[0, i, :3, 3] = rng.randn(3) * 0.1 + # Simple intrinsics + fx = fy = 500.0 + cx, cy = width / 2, height / 2 + K = np.array([[fx, 0, cx], [0, fy, cy], [0, 0, 1]], dtype=np.float32) + intrinsics = np.tile(K, (1, num_views, 1, 1)) + return images, extrinsics, intrinsics + + +# ---- PyTorch forward ------------------------------------------------------- + +def run_pytorch( + images_np: np.ndarray, + extrinsics_np: np.ndarray | None, + intrinsics_np: np.ndarray | None, + config_name: str, + weights_path: str | None, +) -> np.ndarray: + """Run PyTorch model and return depth as numpy. + + images_np: (1, N, H, W, 3) NHWC float32. + Returns: (N, H, W) depth. + """ + import torch + from omegaconf import OmegaConf + + from depth_anything_3.cfg import create_object + from depth_anything_3.utils.geometry import affine_inverse + + cfg_path = f"src/depth_anything_3/configs/{config_name}.yaml" + cfg = OmegaConf.load(cfg_path) + model = create_object(cfg) + model.eval() + + if weights_path is not None: + from safetensors.torch import load_file + state_dict = load_file(weights_path) + + # Strip 'model.' prefix if present + new_state_dict = {} + for k, v in state_dict.items(): + if k.startswith("model."): + new_state_dict[k[6:]] = v + else: + new_state_dict[k] = v + + model.load_state_dict(new_state_dict, strict=False) + print(f"[PyTorch] Loaded weights from {weights_path}") + + # Convert NHWC -> NCHW for PyTorch + # images_np: (1, N, H, W, 3) -> torch (1, N, 3, H, W) + images_pt = torch.from_numpy(images_np).permute(0, 1, 4, 2, 3).float() + + ext_pt = None + int_pt = None + if extrinsics_np is not None: + ext_pt = torch.from_numpy(extrinsics_np).float() + if intrinsics_np is not None: + int_pt = torch.from_numpy(intrinsics_np).float() + + with torch.no_grad(): + output = model(images_pt, extrinsics=ext_pt, intrinsics=int_pt, + export_feat_layers=[], infer_gs=False, use_ray_pose=False) + + depth = output["depth"].numpy() # (1, N, H, W) + return depth[0] # (N, H, W) + + +# ---- MLX forward ----------------------------------------------------------- + +def run_mlx( + images_np: np.ndarray, + extrinsics_np: np.ndarray | None, + intrinsics_np: np.ndarray | None, + config_name: str, + weights_path: str | None, +) -> np.ndarray: + """Run MLX model and return depth as numpy. + + images_np: (1, N, H, W, 3) NHWC float32. + Returns: (N, H, W) depth. + """ + import mlx.core as mx + import mlx.nn as nn + + from mlx_depth_anything_3.model import DepthAnything3Net + + model = DepthAnything3Net(config_name) + + if weights_path is not None: + from convert_to_mlx import convert_weights, load_pytorch_weights + pt_weights = load_pytorch_weights(weights_path) + mlx_weights = convert_weights(pt_weights) + + def flatten_params(params, prefix=""): + flat = {} + for k, v in params.items(): + name = f"{prefix}.{k}" if prefix else k + if isinstance(v, dict): + flat.update(flatten_params(v, name)) + else: + flat[name] = v + return flat + + model_params = flatten_params(model.parameters()) + filtered_weights = {k: v for k, v in mlx_weights.items() if k in model_params} + + # Load the filtered weights + model.load_weights(list(filtered_weights.items()), strict=False) + mx.eval(model.parameters()) + print(f"[MLX] Loaded and converted {len(filtered_weights)} weights from {weights_path}") + + x = mx.array(images_np) + ext_mx = mx.array(extrinsics_np) if extrinsics_np is not None else None + int_mx = mx.array(intrinsics_np) if intrinsics_np is not None else None + + output = model(x, extrinsics=ext_mx, intrinsics=int_mx) + mx.eval(output["depth"]) + + depth = np.array(output["depth"][0]) # (N, H, W) + return depth + + +# ---- Comparison ------------------------------------------------------------ + +def compare(pt_depth: np.ndarray, mlx_depth: np.ndarray, tol: float = 1e-3): + """Compare PyTorch and MLX depth outputs.""" + assert pt_depth.shape == mlx_depth.shape, ( + f"Shape mismatch: PyTorch {pt_depth.shape} vs MLX {mlx_depth.shape}" + ) + diff = np.abs(pt_depth - mlx_depth) + max_diff = diff.max() + mean_diff = diff.mean() + rel_diff = diff / (np.abs(pt_depth) + 1e-8) + max_rel = rel_diff.max() + mean_rel = rel_diff.mean() + + print(f"\n{'='*60}") + print(f"Parity Verification Results") + print(f"{'='*60}") + print(f" Output shape: {pt_depth.shape}") + print(f" Absolute error:") + print(f" Max: {max_diff:.6e}") + print(f" Mean: {mean_diff:.6e}") + print(f" Relative error:") + print(f" Max: {max_rel:.6e}") + print(f" Mean: {mean_rel:.6e}") + print(f" PyTorch range: [{pt_depth.min():.4f}, {pt_depth.max():.4f}]") + print(f" MLX range: [{mlx_depth.min():.4f}, {mlx_depth.max():.4f}]") + + passed = max_diff < tol + print(f"\n Status: {'PASS ✓' if passed else 'FAIL ✗'} (tol={tol})") + print(f"{'='*60}\n") + return passed + + +def main(): + parser = argparse.ArgumentParser(description="Verify DA3 PyTorch/MLX parity") + parser.add_argument("--model", type=str, default="da3-small", + help="Model config name") + parser.add_argument("--weights", type=str, default=None, + help="Path to PyTorch weights file") + parser.add_argument("--num-views", type=int, default=2, + help="Number of input views") + parser.add_argument("--height", type=int, default=504, + help="Image height") + parser.add_argument("--width", type=int, default=504, + help="Image width") + parser.add_argument("--no-cameras", action="store_true", + help="Disable camera conditioning") + parser.add_argument("--tol", type=float, default=1e-3, + help="Tolerance for parity check") + args = parser.parse_args() + + print(f"Generating random inputs: {args.num_views} views, " + f"{args.height}x{args.width}, cameras={'off' if args.no_cameras else 'on'}") + images, extrinsics, intrinsics = make_random_input( + num_views=args.num_views, + height=args.height, + width=args.width, + with_cameras=not args.no_cameras, + ) + + print(f"\nRunning PyTorch model ({args.model})...") + pt_depth = run_pytorch(images, extrinsics, intrinsics, args.model, args.weights) + + print(f"Running MLX model ({args.model})...") + mlx_depth = run_mlx(images, extrinsics, intrinsics, args.model, args.weights) + + passed = compare(pt_depth, mlx_depth, tol=args.tol) + sys.exit(0 if passed else 1) + + +if __name__ == "__main__": + main()