Initial commit

.gitattributes

@@ -0,0 +1,2 @@
+# Auto detect text files and perform LF normalization
+* text=auto

.gitignore

@@ -0,0 +1,160 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/

__init__.py

@@ -0,0 +1,3 @@
+from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS
+
+__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"]

meshflow/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+

meshflow/models/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+

meshflow/models/condition_encoder.py

@@ -0,0 +1,137 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+import random
+from typing import Union
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from PIL import Image
+from torchvision import transforms
+
+ImageType = Union[Image.Image, torch.Tensor]
+
+_HUB_EMBED_DIMS = {
+    "dinov3_vits16": 384,
+    "dinov3_vits16plus": 384,
+    "dinov3_vitb16": 768,
+    "dinov3_vitl16": 1024,
+    "dinov3_vitl16plus": 1024,
+    "dinov3_vith16plus": 1280,
+    "dinov3_vit7b16": 4096,
+}
+
+
+def empty_visual_embeds_shape(image_size: int, embed_dim: int) -> tuple[int, int]:
+    seq_len = 1 + 4 + (image_size // 16) ** 2
+    return seq_len, embed_dim
+
+
+def make_empty_visual_embeds(image_size: int, embed_dim: int) -> torch.Tensor:
+    seq_len, dim = empty_visual_embeds_shape(image_size, embed_dim)
+    return torch.zeros(1, seq_len, dim)
+
+
+class DINOv3Encoder(nn.Module):
+    """ViT-based visual encoder."""
+
+    def __init__(
+        self,
+        image_size: int,
+        hub_model: str,
+        hub_dir: str | None = None,
+        hub_weights: str | None = None,
+        pretrained: bool = True,
+        freeze: bool = True,
+        empty_embeds_ratio: float = 0.1,
+    ):
+        super().__init__()
+        self.empty_embeds_ratio = empty_embeds_ratio
+
+        if hub_dir is None:
+            hub_dir = os.path.expanduser("~/.cache/torch/hub/facebookresearch_dinov3_main")
+        elif "/root/.cache/" in hub_dir:
+            hub_dir = hub_dir.replace("/root/.cache/", os.path.expanduser("~/.cache/"))
+
+        hub_kwargs: dict = {"pretrained": pretrained}
+        if hub_weights is not None:
+            hub_kwargs["weights"] = hub_weights
+        self.dino_model = torch.hub.load(
+            hub_dir, hub_model, source="local", **hub_kwargs
+        )
+        embed_dim = self.dino_model.embed_dim
+
+        self.transform = transforms.Compose([
+            transforms.Resize(image_size, transforms.InterpolationMode.BICUBIC, antialias=True),
+            transforms.CenterCrop(image_size),
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+        seq_len, _ = empty_visual_embeds_shape(image_size, embed_dim)
+        self.register_buffer(
+            "empty_image_embeds",
+            torch.zeros(1, seq_len, embed_dim),
+            persistent=False,
+        )
+
+        self.dino_model.eval()
+        if freeze:
+            self.dino_model.requires_grad_(False)
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def encode_image(self, image: ImageType) -> torch.Tensor:
+        if isinstance(image, torch.Tensor):
+            x = self.transform(image.permute(0, 3, 1, 2))
+        else:
+            x = self.transform(transforms.ToTensor()(image).unsqueeze(0))
+
+        ref = next(self.dino_model.parameters())
+        x = x.to(device=ref.device, dtype=ref.dtype)
+        with torch.autocast(device_type=ref.device.type, enabled=False):
+            out = self.dino_model.forward_features(x)["x_prenorm"]
+        return F.layer_norm(out, out.shape[-1:])
+
+    def forward(self, batch):
+        bs = batch["image"].shape[0]
+        if self.training and random.random() < self.empty_embeds_ratio:
+            return self.empty_image_embeds.repeat(bs, 1, 1).to(batch["image"].device)
+        return self.encode_image(batch["image"])
+
+
+def build_dinov3_encoder(
+    encoder_type: str,
+    cfg: dict,
+) -> DINOv3Encoder:
+    if encoder_type != "dinov3-encoder":
+        raise ValueError(f"Unknown visual_condition_type: {encoder_type!r}")
+
+    return DINOv3Encoder(
+        image_size=cfg["image_size"],
+        hub_model=cfg["hub_model"],
+        hub_dir=cfg.get("hub_dir", os.path.expanduser("~/.cache/torch/hub/facebookresearch_dinov3_main")),
+        hub_weights=cfg.get("hub_weights"),
+        pretrained=cfg.get("pretrained", True),
+        freeze=cfg.get("freeze", True),
+        empty_embeds_ratio=cfg.get("empty_embeds_ratio", 0.1),
+    )
+
+
+def resolve_visual_embed_dim(cfg: dict, fallback: int | None = None) -> int:
+    if "embed_dim" in cfg:
+        return int(cfg["embed_dim"])
+    hub_model = cfg.get("hub_model")
+    if hub_model in _HUB_EMBED_DIMS:
+        return _HUB_EMBED_DIMS[hub_model]
+    if fallback is not None:
+        return int(fallback)
+    raise ValueError(
+        f"Cannot infer DINOv3 embed_dim for hub_model={hub_model!r}; "
+        "set visual_condition.embed_dim in config.yaml"
+    )

meshflow/models/meshflow_dit.py

@@ -0,0 +1,887 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from dataclasses import dataclass
+from typing import Any, Dict, Optional, Tuple
+
+import numpy as np
+import torch
+from torch import nn
+import torch.nn.functional as F
+from diffusers.models.modeling_utils import ModelMixin
+from einops import rearrange
+
+
+def voxelize_pc(pc: torch.Tensor, resolution: int):
+    """
+    Voxelize the point cloud in [-1, 1]^3 with fixed resolution.
+    Each point is replaced by the center of the voxel that contains it.
+
+    Args:
+        pc: (N, 3) or (bs, N, 3) point cloud in range [-1, 1].
+        resolution: number of voxels per axis (e.g. 32 -> 32^3 grid).
+
+    Returns:
+        centers: same shape as pc, voxel center coordinates in [-1, 1], float32.
+        vox_idx: same shape as pc, integer voxel indices in [0, resolution - 1].
+    """
+    pc = pc.to(dtype=torch.float32)
+    t = (pc + 1.0) * 0.5 * resolution
+    vox_idx = t.long().clamp(0, resolution - 1)
+    centers = -1.0 + (vox_idx + 0.5) * 2.0 / resolution
+    return centers.to(dtype=torch.float32), vox_idx
+
+
+@dataclass
+class Transformer1DModelOutput:
+    """Output of the 1D transformer; single field `sample`."""
+
+    sample: torch.FloatTensor
+
+
+def modulate(x, shift, scale):
+    """Apply affine modulation: x * (1 + scale) + shift (broadcast over sequence)."""
+    return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+def get_1d_sincos_pos_embed_from_grid(embed_dim: int, pos: np.ndarray) -> np.ndarray:
+    """
+    Build 1D sin/cos position embeddings.
+
+    Args:
+        embed_dim: Output dimension per position; must be even.
+        pos: Position indices, shape (M,) or flattened to (M,).
+
+    Returns:
+        Position embeddings of shape (M, embed_dim).
+    """
+    assert embed_dim % 2 == 0
+    omega = np.arange(embed_dim // 2, dtype=np.float64)
+    omega /= embed_dim / 2.0
+    omega = 1.0 / 10000**omega
+
+    pos = pos.reshape(-1)
+    out = np.einsum("m,d->md", pos, omega)
+    emb_sin = np.sin(out)
+    emb_cos = np.cos(out)
+    return np.concatenate([emb_sin, emb_cos], axis=1)
+
+
+class DropPath(nn.Module):
+    """Stochastic depth per sample (drop entire residual branch)."""
+
+    def __init__(self, drop_prob: float = 0.0) -> None:
+        super().__init__()
+        self.drop_prob = float(drop_prob)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.drop_prob == 0.0 or not self.training:
+            return x
+        keep_prob = 1.0 - self.drop_prob
+        shape = (x.shape[0],) + (1,) * (x.ndim - 1)
+        mask = x.new_empty(shape).bernoulli_(keep_prob)
+        if keep_prob > 0.0:
+            mask.div_(keep_prob)
+        return x * mask
+
+
+class RotaryPositionalEmbeddings(nn.Module):
+    """
+    Rotary Positional Embeddings (RoPE) using a "Divide and Conquer" strategy.
+    
+    This implementation splits the head dimension equally among the input dimensions (ndim).
+    Each axis (1D, 2D, or 3D) is assigned an independent chunk of the head_dim, 
+    ensuring that rotational signals are applied consistently across all spatial axes.
+    """
+    
+    def __init__(self, dim: int, base: float = 10_000.0, ndim: int = 1) -> None:
+        """
+        Args:
+            dim: Total dimension of the attention head (head_dim).
+            base: Base for the geometric progression of rotation frequencies.
+            ndim: Number of spatial dimensions (e.g., 1 for sequence, 2 for 2D grid).
+        """
+        super().__init__()
+        assert dim % ndim == 0, f"head_dim ({dim}) must be divisible by ndim ({ndim})"
+        
+        self.dim = dim
+        self.base = base
+        self.ndim = ndim
+        self.sub_dim = dim // ndim  # Dimension assigned to each axis
+        assert (
+            self.sub_dim % 2 == 0
+        ), f"sub_dim ({self.sub_dim}) must be even to form rotation pairs"
+        
+        # Precompute the inverse frequency for each axis chunk.
+        # Each sub_dim chunk will be treated as sub_dim // 2 rotation pairs.
+        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.sub_dim, 2).float() / self.sub_dim))
+        self.register_buffer("inv_freq", inv_freq, persistent=False)
+
+    def _compute_cos_sin(self, pos_axis: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        """
+        Compute cos and sin for a specific axis based on its position coordinates.
+        
+        Args:
+            pos_axis: Position indices for one axis, shape [..., 1].
+        Returns:
+            cos, sin: Both shape [..., sub_dim // 2].
+        """
+        # Outer product: [..., 1] * [sub_dim // 2] -> [..., sub_dim // 2]
+        inv_freq = self.inv_freq.to(device=pos_axis.device)
+        args = pos_axis.float() * inv_freq.view(1, -1)
+        return torch.cos(args), torch.sin(args)
+
+    @staticmethod
+    def _apply_rotary(x: torch.Tensor, cos: torch.Tensor, sin: torch.Tensor) -> torch.Tensor:
+        """
+        Apply the rotation matrix to a chunk of the features.
+        
+        Args:
+            x: Input feature chunk, shape [..., sub_dim].
+            cos: Cosine values, shape [..., sub_dim // 2].
+            sin: Sine values, shape [..., sub_dim // 2].
+        Returns:
+            Rotated features, shape [..., sub_dim].
+        """
+        
+        # Reshape to [..., sub_dim // 2, 2] to apply 2D rotation for each pair
+        x_float = x.float().reshape(*x.shape[:-1], -1, 2)
+        x1, x2 = x_float[..., 0], x_float[..., 1]
+        
+        # Apply the rotation:
+        # [x1_new] = [cos, -sin] * [x1]
+        # [x2_new]   [sin,  cos]   [x2]
+        x_out = torch.stack([
+            x1 * cos - x2 * sin,
+            x1 * sin + x2 * cos
+        ], dim=-1).flatten(-2)
+        
+        return x_out.type_as(x)
+
+    def forward(self, x: torch.Tensor, pos: torch.Tensor) -> torch.Tensor:
+        """
+        Apply RoPE to the input tensor.
+        
+        Args:
+            x: Input tensor [batch, seq_len, num_heads, head_dim].
+            pos: Spatial coordinates [batch, seq_len, ndim].
+            
+        Returns:
+            Rotated tensor of the same shape as x.
+        """
+        if pos.shape[-1] != self.ndim:
+            raise ValueError(f"pos last dim must be {self.ndim}, got {pos.shape[-1]}")
+        # 1. Split x into ndim chunks along the head_dim
+        # Each chunk has shape [..., sub_dim]
+        B, N, H, D = x.shape
+        x_split = x.chunk(self.ndim, dim=-1)
+
+        out_list = []
+        for i in range(self.ndim):
+            # 2. Extract coordinates for the current axis: [..., 1]
+            pos_axis = pos[..., i:i+1]
+            
+            # 3. Compute rotation frequencies for this specific axis
+            cos, sin = self._compute_cos_sin(pos_axis)
+            cos = cos.unsqueeze(2).expand(-1, -1, H, -1)
+            sin = sin.unsqueeze(2).expand(-1, -1, H, -1)
+
+            # 4. Apply rotation to the corresponding chunk
+            out_list.append(self._apply_rotary(x_split[i], cos, sin))
+            
+        # 5. Concatenate all rotated chunks back to original head_dim
+        return torch.cat(out_list, dim=-1)
+
+
+class Timesteps(nn.Module):
+    """Sinusoidal timestep embedding (sin/cos over frequency bands)."""
+
+    def __init__(
+        self,
+        num_channels: int,
+        downscale_freq_shift: float = 0.0,
+        scale: int = 1000,
+        max_period: int = 10000,
+    ):
+        super().__init__()
+        self.num_channels = num_channels
+        self.downscale_freq_shift = downscale_freq_shift
+        self.scale = scale
+        self.max_period = max_period
+
+        half_dim = num_channels // 2
+        exponent = -math.log(max_period) * torch.arange(
+            start=0, end=half_dim, dtype=torch.float32
+        )
+        exponent = exponent / (half_dim - downscale_freq_shift)
+        self.register_buffer("freq_embedding", torch.exp(exponent), persistent=False)
+
+    def forward(self, timesteps: torch.Tensor) -> torch.Tensor:
+        """Embed 1D timestep tensor (shape [B]) to [B, num_channels]."""
+        assert len(timesteps.shape) == 1, "Timesteps should be a 1d-array"
+        if self.scale == 1000:
+            assert timesteps.max() <= 1.0, "Timesteps should be in the range [0, 1]"
+        emb = timesteps[:, None].float() * self.freq_embedding[None, :].to(
+            timesteps.device
+        )
+        emb = self.scale * emb
+        emb = torch.cat([torch.sin(emb), torch.cos(emb)], dim=-1)
+        if self.num_channels % 2 == 1:
+            emb = torch.nn.functional.pad(emb, (0, 1, 0, 0))
+        return emb
+
+
+class TimestepEmbedder(nn.Module):
+    """Embed scalar timesteps to a vector [B, out_size] for AdaLN (no extra sequence token)."""
+
+    def __init__(
+        self,
+        hidden_size,
+        frequency_embedding_size=256,
+        cond_proj_dim=None,
+        out_size=None,
+    ):
+        super().__init__()
+        if out_size is None:
+            out_size = hidden_size
+        self.mlp = nn.Sequential(
+            nn.Linear(hidden_size, frequency_embedding_size, bias=True),
+            nn.GELU(),
+            nn.Linear(frequency_embedding_size, out_size, bias=True),
+        )
+        self.frequency_embedding_size = frequency_embedding_size
+
+        if cond_proj_dim is not None:
+            self.cond_proj = nn.Linear(cond_proj_dim, hidden_size, bias=False)
+
+        self.time_embed = Timesteps(hidden_size)
+
+    def forward(self, t: torch.Tensor, condition: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """Return timestep embedding [B, out_size]."""
+        t_freq = self.time_embed(t).type(self.mlp[0].weight.dtype)
+
+        if condition is not None:
+            t_freq = t_freq + self.cond_proj(condition)
+
+        return self.mlp(t_freq)
+
+
+class MLP(nn.Module):
+    """Two-layer FFN: Linear -> GELU -> Linear with 4x expansion."""
+
+    def __init__(self, *, width: int):
+        super().__init__()
+        self.width = width
+        self.fc1 = nn.Linear(width, width * 4)
+        self.fc2 = nn.Linear(width * 4, width)
+        self.gelu = nn.GELU()
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.fc2(self.gelu(self.fc1(x)))
+
+
+class CrossAttention(nn.Module):
+    """Multi-head cross-attention (Q from x, K/V from y) with optional QK norm and RoPE on Q."""
+
+    def __init__(
+        self,
+        qdim,
+        kdim,
+        num_heads,
+        qkv_bias=True,
+        qk_norm=False,
+        norm_layer=nn.LayerNorm,
+        use_rope=False,
+        rope_input_ndim: Optional[int] = None,
+        **kwargs,
+    ):
+        super().__init__()
+        self.qdim = qdim
+        self.kdim = kdim
+        self.num_heads = num_heads
+        assert self.qdim % num_heads == 0, "self.qdim must be divisible by num_heads"
+        self.head_dim = self.qdim // num_heads
+        assert (
+            self.head_dim % 8 == 0 and self.head_dim <= 128
+        ), "Only support head_dim <= 128 and divisible by 8"
+        self.scale = self.head_dim**-0.5
+
+        self.to_q = nn.Linear(qdim, qdim, bias=qkv_bias)
+        self.to_k = nn.Linear(kdim, qdim, bias=qkv_bias)
+        self.to_v = nn.Linear(kdim, qdim, bias=qkv_bias)
+        self.q_norm = (
+            norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6)
+            if qk_norm
+            else nn.Identity()
+        )
+        self.k_norm = (
+            norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6)
+            if qk_norm
+            else nn.Identity()
+        )
+        self.use_rope = use_rope
+        if self.use_rope:
+            self.rotary_emb_q = RotaryPositionalEmbeddings(self.head_dim, ndim=rope_input_ndim)
+        self.out_proj = nn.Linear(qdim, qdim, bias=True)
+
+    def forward(
+        self, 
+        x: torch.Tensor, 
+        y: torch.Tensor, 
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """
+        Cross-attention: queries from x, keys/values from y.
+
+        Args:
+            x: Query sequence [batch, seq_len_q, qdim].
+            y: Key/value sequence [batch, seq_len_k, kdim].
+
+        Returns:
+            Output [batch, seq_len_q, qdim].
+        """
+        b, s1, c = x.shape
+        _, s2, c = y.shape
+
+        q = rearrange(self.to_q(x), "b n (h d) -> b n h d", h=self.num_heads)
+        k = rearrange(self.to_k(y), "b n (h d) -> b n h d", h=self.num_heads)
+        v = rearrange(self.to_v(y), "b n (h d) -> b n h d", h=self.num_heads)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if self.use_rope:
+            if joint_attention_kwargs is not None and joint_attention_kwargs.get("rope_input") is not None:
+                rope_in = joint_attention_kwargs["rope_input"]
+                q = self.rotary_emb_q(q, rope_in)
+            else:
+                pos_1d = torch.arange(s1, dtype=torch.long, device=q.device).view(1, -1, 1).expand(b, -1, -1)
+                q = self.rotary_emb_q(q, pos_1d)
+
+        with torch.nn.attention.sdpa_kernel(
+            backends=[
+                torch.nn.attention.SDPBackend.FLASH_ATTENTION,
+                torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION,
+                torch.nn.attention.SDPBackend.MATH,
+            ]
+        ):
+            q, k, v = map(
+                lambda t: rearrange(t, "b n h d -> b h n d", h=self.num_heads),
+                (q, k, v),
+            )
+            context = (
+                F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
+                .transpose(1, 2)
+                .reshape(b, s1, -1)
+            )
+
+        return self.out_proj(context)
+
+
+class Attention(nn.Module):
+    """Multi-head self-attention with optional QK norm and RoPE."""
+
+    def __init__(
+        self,
+        dim,
+        num_heads,
+        qkv_bias=True,
+        qk_norm=False,
+        norm_layer=nn.LayerNorm,
+        use_rope=False,
+        rope_input_ndim: Optional[int] = None,
+    ):
+        super().__init__()
+        self.dim = dim
+        self.num_heads = num_heads
+        assert self.dim % num_heads == 0, "dim should be divisible by num_heads"
+        self.head_dim = self.dim // num_heads
+        assert (
+            self.head_dim % 8 == 0 and self.head_dim <= 128
+        ), "Only support head_dim <= 128 and divisible by 8"
+        self.scale = self.head_dim**-0.5
+
+        self.to_q = nn.Linear(dim, dim, bias=qkv_bias)
+        self.to_k = nn.Linear(dim, dim, bias=qkv_bias)
+        self.to_v = nn.Linear(dim, dim, bias=qkv_bias)
+        self.q_norm = (
+            norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6)
+            if qk_norm
+            else nn.Identity()
+        )
+        self.k_norm = (
+            norm_layer(self.head_dim, elementwise_affine=True, eps=1e-6)
+            if qk_norm
+            else nn.Identity()
+        )
+        self.use_rope = use_rope
+        if self.use_rope:
+            self.rotary_emb_q = RotaryPositionalEmbeddings(self.head_dim, ndim=rope_input_ndim)
+            self.rotary_emb_k = RotaryPositionalEmbeddings(self.head_dim, ndim=rope_input_ndim)
+
+        self.out_proj = nn.Linear(dim, dim)
+
+    def forward(
+        self, 
+        x: torch.Tensor, 
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Self-attention over sequence; input/output [batch, seq_len, dim]."""
+        b, s, _ = x.shape
+
+        q = rearrange(self.to_q(x), "b n (h d) -> b n h d", h=self.num_heads)
+        k = rearrange(self.to_k(x), "b n (h d) -> b n h d", h=self.num_heads)
+        v = rearrange(self.to_v(x), "b n (h d) -> b n h d", h=self.num_heads)
+
+        q = self.q_norm(q)
+        k = self.k_norm(k)
+
+        if self.use_rope:
+            if joint_attention_kwargs is not None and joint_attention_kwargs.get("rope_input") is not None:
+                rope_in = joint_attention_kwargs["rope_input"]
+                q = self.rotary_emb_q(q, rope_in)
+                k = self.rotary_emb_k(k, rope_in)
+            else:
+                pos_1d = torch.arange(s, dtype=torch.long, device=q.device).view(1, -1, 1).expand(b, -1, -1)
+                q = self.rotary_emb_q(q, pos_1d)
+                k = self.rotary_emb_k(k, pos_1d)
+
+        with torch.nn.attention.sdpa_kernel(
+            backends=[
+                torch.nn.attention.SDPBackend.FLASH_ATTENTION,
+                torch.nn.attention.SDPBackend.EFFICIENT_ATTENTION,
+                torch.nn.attention.SDPBackend.MATH,
+            ]
+        ):
+            q, k, v = map(
+                lambda t: rearrange(t, "b n h d -> b h n d", h=self.num_heads),
+                (q, k, v),
+            )
+            context = (
+                F.scaled_dot_product_attention(q, k, v, attn_mask=attention_mask)
+                .transpose(1, 2)
+                .reshape(b, s, -1)
+            )
+
+        return self.out_proj(context)
+
+
+def _zero_init_adaln_modulation(linear: nn.Linear) -> None:
+    nn.init.constant_(linear.weight, 0)
+    nn.init.constant_(linear.bias, 0)
+
+
+class MeshFlowDiTBlock(nn.Module):
+    """
+    Transformer block with AdaLN-Zero (7-way: self-attn + MLP + gate_cross for cross-attn),
+    cross-attention, optional U-Net skip merge (learnable gate on residual path, default init 0.5),
+    and DropPath on gated residuals.
+    """
+
+    def __init__(
+        self,
+        hidden_size,
+        c_emb_size,
+        num_heads,
+        context_dim=1024,
+        qk_norm=False,
+        norm_layer=nn.LayerNorm,
+        qk_norm_layer=nn.RMSNorm,
+        qkv_bias=True,
+        skip_connection=True,
+        use_ele_affine: bool = True,
+        use_rope: bool = False,
+        rope_input_ndim: Optional[int] = None,
+        use_rope_in_cross_attention: bool = False,
+        skip_zero_init_gate: bool = True,
+        drop_path_rate: float = 0.0,
+        **kwargs,
+    ):
+        super().__init__()
+        self.use_ele_affine = use_ele_affine
+        self.use_rope = use_rope
+        self.use_rope_in_cross_attention = use_rope_in_cross_attention
+
+        self.norm1 = norm_layer(
+            hidden_size, elementwise_affine=use_ele_affine, eps=1e-6
+        )
+        self.attn1 = Attention(
+            hidden_size,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            norm_layer=qk_norm_layer,
+            use_rope=use_rope,
+            rope_input_ndim=rope_input_ndim,
+        )
+        self.norm2 = norm_layer(
+            hidden_size, elementwise_affine=use_ele_affine, eps=1e-6
+        )
+        self.attn2 = CrossAttention(
+            hidden_size,
+            context_dim,
+            num_heads=num_heads,
+            qkv_bias=qkv_bias,
+            qk_norm=qk_norm,
+            norm_layer=qk_norm_layer,
+            use_rope=use_rope and use_rope_in_cross_attention,
+            rope_input_ndim=rope_input_ndim,
+        )
+        self.norm3 = norm_layer(hidden_size, elementwise_affine=use_ele_affine, eps=1e-6)
+
+        self.adaLN_modulation = nn.Sequential(nn.SiLU(), nn.Linear(c_emb_size, 7 * hidden_size, bias=True))
+        _zero_init_adaln_modulation(self.adaLN_modulation[-1])
+
+        if skip_connection:
+            self.skip_norm = norm_layer(hidden_size, elementwise_affine=True, eps=1e-6)
+            self.skip_linear = nn.Linear(2 * hidden_size, hidden_size)
+            self.skip_zero_init_gate = skip_zero_init_gate
+            if skip_zero_init_gate:
+                # Skip residual carries shallow 1D structure; 0.5 avoids fully suppressing it at init.
+                self.skip_gate = nn.Parameter(torch.full((1,), 0.5))
+            else:
+                self.skip_gate = None
+        else:
+            self.skip_linear = None
+            self.skip_norm = None
+            self.skip_gate = None
+
+        self.mlp = MLP(width=hidden_size)
+        self.drop_path_attn = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+        self.drop_path_mlp = DropPath(drop_path_rate) if drop_path_rate > 0 else nn.Identity()
+
+    def _forward(
+        self,
+        hidden_states: torch.Tensor,
+        conditions: Optional[torch.Tensor] = None,
+        temb: Optional[torch.Tensor] = None,
+        skip_value: Optional[torch.Tensor] = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.skip_linear is not None:
+            assert self.skip_norm is not None
+            cat = torch.cat([skip_value, hidden_states], dim=-1)
+            skip_out = self.skip_linear(cat)
+            if self.skip_gate is not None:
+                hidden_states = self.skip_norm(skip_out + self.skip_gate * hidden_states)
+            else:
+                hidden_states = self.skip_norm(skip_out)
+
+        assert temb is not None
+        (
+            shift_msa,
+            scale_msa,
+            gate_msa,
+            shift_mlp,
+            scale_mlp,
+            gate_mlp,
+            gate_cross,
+        ) = self.adaLN_modulation(temb).chunk(7, dim=1)
+
+        attn_in = modulate(self.norm1(hidden_states), shift_msa, scale_msa)
+        hidden_states = hidden_states + self.drop_path_attn(
+            gate_msa.unsqueeze(1)
+            * self.attn1(
+                attn_in,
+                joint_attention_kwargs=joint_attention_kwargs,
+                attention_mask=attention_mask,
+            )
+        )
+        if conditions is not None:
+            hidden_states = hidden_states + gate_cross.unsqueeze(1) * self.attn2(
+                self.norm2(hidden_states),
+                conditions,
+                joint_attention_kwargs=joint_attention_kwargs,
+                attention_mask=cross_attention_mask,
+            )
+        mlp_in = modulate(self.norm3(hidden_states), shift_mlp, scale_mlp)
+        hidden_states = hidden_states + self.drop_path_mlp(
+            gate_mlp.unsqueeze(1) * self.mlp(mlp_in)
+        )
+
+        return hidden_states
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        conditions: Optional[torch.Tensor] = None,
+        temb: Optional[torch.Tensor] = None,
+        skip_value: Optional[torch.Tensor] = None,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        cross_attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if skip_value is None and attention_mask is None and cross_attention_mask is None:
+            return self._forward(
+                hidden_states,
+                conditions,
+                temb,
+                skip_value=None,
+                joint_attention_kwargs=joint_attention_kwargs,
+                attention_mask=attention_mask,
+                cross_attention_mask=cross_attention_mask,
+            )
+
+        def _run(
+            _hidden_states: torch.Tensor,
+            _conditions: Optional[torch.Tensor],
+            _temb: Optional[torch.Tensor],
+            _skip_value: Optional[torch.Tensor],
+            _attention_mask: Optional[torch.Tensor],
+            _cross_attention_mask: Optional[torch.Tensor],
+        ) -> torch.Tensor:
+            return self._forward(
+                _hidden_states,
+                _conditions,
+                _temb,
+                skip_value=_skip_value,
+                joint_attention_kwargs=joint_attention_kwargs,
+                attention_mask=_attention_mask,
+                cross_attention_mask=_cross_attention_mask,
+            )
+
+        return torch.utils.checkpoint.checkpoint(
+            _run,
+            hidden_states,
+            conditions,
+            temb,
+            skip_value,
+            attention_mask,
+            cross_attention_mask,
+            use_reentrant=False,
+        )
+
+
+class FinalLayer(nn.Module):
+    """Linear head; apply ``MeshFlowDiT.final_norm`` before this module."""
+
+    def __init__(self, final_hidden_size: int, out_channels: int):
+        super().__init__()
+        self.linear = nn.Linear(final_hidden_size, out_channels, bias=True)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return self.linear(x)
+
+
+class MeshFlowDiT(ModelMixin):
+    """
+    MeshFlow DiT: AdaLN-Zero, optional sin-cos PE, RoPE, visual cross-attention,
+    and a final norm + linear head.
+    """
+
+    def __init__(
+        self,
+        input_size: int = 1024,
+        in_channels: int = 4,
+        hidden_size: int = 1024,
+        context_dim: int = 1024,
+        depth: int = 24,
+        num_heads: int = 16,
+        mlp_ratio: float = 4.0,
+        norm_type: str = "layer",
+        qk_norm_type: str = "rms",
+        qk_norm: bool = False,
+        qkv_bias: bool = True,
+        use_pos_emb: bool = False,
+        use_rope: bool = False,
+        rope_input_ndim: Optional[int] = None,
+        use_rope_in_cross_attention: bool = False,
+        use_proj_cond_on_temb: bool = False,
+        proj_cond_dim: int = 1,
+        visual_condition_dim: Optional[int] = None,
+        skip_zero_init_gate: bool = True,
+        drop_path_rate: float = 0.0,
+        use_skip_connection: bool = True,
+        zero_init: bool = False,
+        pretrained_model_name_or_path: Optional[str] = None,
+        strict_load: bool = True,
+    ):
+        super().__init__()
+        self.input_size = input_size
+        self.depth = depth
+        self.in_channels = in_channels
+        self.out_channels = in_channels
+        self.num_heads = num_heads
+
+        self.hidden_size = hidden_size
+        self.norm = nn.LayerNorm if norm_type == "layer" else nn.RMSNorm
+        self.qk_norm = nn.RMSNorm if qk_norm_type == "rms" else nn.LayerNorm
+        self.context_dim = context_dim
+        self.use_pos_emb = use_pos_emb
+        self.use_rope = use_rope
+        self.use_rope_in_cross_attention = use_rope_in_cross_attention
+        self.use_proj_cond_on_temb = use_proj_cond_on_temb
+        self.proj_cond_dim = proj_cond_dim
+        self.drop_path_rate = drop_path_rate
+
+        self.x_embedder = nn.Linear(in_channels, hidden_size, bias=True)
+        self.t_embedder = TimestepEmbedder(
+            hidden_size,
+            hidden_size * 4,
+            cond_proj_dim=proj_cond_dim if use_proj_cond_on_temb else None,
+        )
+        if self.use_pos_emb:
+            self.register_buffer("pos_embed", torch.zeros(1, input_size, hidden_size))
+            pos = np.arange(self.input_size, dtype=np.float32)
+            pos_embed = get_1d_sincos_pos_embed_from_grid(self.pos_embed.shape[-1], pos)
+            self.pos_embed.data.copy_(torch.from_numpy(pos_embed).float().unsqueeze(0))
+
+        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, depth)] if depth > 0 else []
+        self.blocks = nn.ModuleList(
+            [
+                MeshFlowDiTBlock(
+                    hidden_size=hidden_size,
+                    c_emb_size=hidden_size,
+                    num_heads=num_heads,
+                    mlp_ratio=mlp_ratio,
+                    context_dim=context_dim,
+                    qk_norm=qk_norm,
+                    norm_layer=self.norm,
+                    qk_norm_layer=self.qk_norm,
+                    skip_connection=use_skip_connection and layer > depth // 2,
+                    qkv_bias=qkv_bias,
+                    use_rope=use_rope,
+                    rope_input_ndim=rope_input_ndim,
+                    use_rope_in_cross_attention=use_rope_in_cross_attention,
+                    skip_zero_init_gate=skip_zero_init_gate,
+                    drop_path_rate=dpr[layer] if dpr else 0.0,
+                )
+                for layer in range(depth)
+            ]
+        )
+
+        self.final_norm = self.norm(hidden_size, elementwise_affine=True, eps=1e-6)
+        self.final_layer = FinalLayer(hidden_size, self.out_channels)
+
+        visual_dim = visual_condition_dim if visual_condition_dim is not None else context_dim
+        if visual_dim != context_dim:
+            self.proj_visual = nn.Sequential(nn.RMSNorm(visual_dim), nn.Linear(visual_dim, context_dim))
+        else:
+            self.proj_visual = None
+
+        if zero_init:
+            for block in self.blocks:
+                for name in ("attn1", "attn2"):
+                    attn = getattr(block, name)
+                    nn.init.constant_(attn.out_proj.weight, 0)
+                    nn.init.constant_(attn.out_proj.bias, 0)
+                nn.init.constant_(block.mlp.fc2.weight, 0)
+                nn.init.constant_(block.mlp.fc2.bias, 0)
+
+        if pretrained_model_name_or_path:
+            ckpt = torch.load(pretrained_model_name_or_path, map_location="cpu")
+            state = ckpt["state_dict"] if isinstance(ckpt, dict) and "state_dict" in ckpt else ckpt
+            filtered = {}
+            for k, v in state.items():
+                if k.startswith("denoiser_model.dit_model."):
+                    filtered[k[len("denoiser_model.dit_model.") :]] = v
+                elif k.startswith("denoiser_model."):
+                    rel = k[len("denoiser_model.") :]
+                    if not rel.startswith("dit_model."):
+                        filtered[rel] = v
+            if filtered:
+                self.load_state_dict(filtered, strict=strict_load)
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def forward(
+        self,
+        model_input,
+        timestep,
+        visual_condition: Optional[torch.Tensor] = None,
+        attention_mask=None,
+        cross_attention_mask=None,
+        joint_attention_kwargs=None,
+        return_dict=True,
+        **kwargs,
+    ):
+        assert timestep.max() <= 1.0
+        conditions = None
+        if visual_condition is not None:
+            conditions = self.proj_visual(visual_condition) if self.proj_visual is not None else visual_condition
+
+        x = model_input
+        t = timestep
+        if self.use_proj_cond_on_temb:
+            temb = self.t_embedder(t, condition=kwargs["proj_cond_on_temb"])
+        else:
+            temb = self.t_embedder(t, condition=None)
+
+        hidden_states = self.x_embedder(x)
+        if self.use_pos_emb:
+            pos_embed = self.pos_embed.to(hidden_states.dtype)
+            hidden_states = hidden_states + pos_embed
+
+        if attention_mask is not None:
+            if attention_mask.dim() == 2:
+                attention_mask = attention_mask.bool().unsqueeze(1).unsqueeze(1)
+            else:
+                attention_mask = attention_mask.bool()
+
+        if cross_attention_mask is not None:
+            if cross_attention_mask.dim() == 2:
+                cross_attention_mask = cross_attention_mask.bool().unsqueeze(1).unsqueeze(1)
+            else:
+                cross_attention_mask = cross_attention_mask.bool()
+
+        skip_value_list = []
+        for layer, block in enumerate(self.blocks):
+            skip_value = None if layer <= self.depth // 2 else skip_value_list.pop()
+            hidden_states = block(
+                hidden_states,
+                conditions,
+                temb,
+                skip_value=skip_value,
+                joint_attention_kwargs=joint_attention_kwargs,
+                attention_mask=attention_mask,
+                cross_attention_mask=cross_attention_mask,
+            )
+            if layer < self.depth // 2:
+                skip_value_list.append(hidden_states)
+
+        hidden_states = self.final_norm(hidden_states)
+        hidden_states = self.final_layer(hidden_states)
+        if not return_dict:
+            return hidden_states
+        return Transformer1DModelOutput(sample=hidden_states)
+
+
+def build_meshflow_dit(d: dict) -> MeshFlowDiT:
+    """Map denoiser_model YAML dict to MeshFlowDiT constructor kwargs."""
+    context_dim = d.get("condition_dim", 1024)
+    return MeshFlowDiT(
+        input_size=d["input_size"],
+        in_channels=d["input_channels"],
+        hidden_size=d["width"],
+        context_dim=context_dim,
+        depth=d["layers"],
+        num_heads=d["num_heads"],
+        mlp_ratio=d.get("mlp_ratio", 4.0),
+        use_pos_emb=d.get("use_pos_emb", False),
+        use_rope=d.get("use_rope", True),
+        rope_input_ndim=d.get("rope_input_ndim", 3),
+        use_rope_in_cross_attention=d.get("use_rope_in_cross_attention", False),
+        norm_type=d.get("norm_type", "rms"),
+        qk_norm_type=d.get("qk_norm_type", "rms"),
+        qk_norm=d.get("qk_norm", True),
+        qkv_bias=d.get("qkv_bias", False),
+        visual_condition_dim=d.get("visual_condition_dim", context_dim),
+        use_proj_cond_on_temb=d.get("use_proj_cond_on_temb", False),
+        proj_cond_dim=d.get("proj_cond_dim", 1),
+        skip_zero_init_gate=d.get("skip_zero_init_gate", True),
+        use_skip_connection=d.get("use_skip_connection", True),
+        drop_path_rate=d.get("drop_path_rate", 0.0),
+        zero_init=d.get("zero_init", False),
+        pretrained_model_name_or_path=d.get("pretrained_model_name_or_path"),
+        strict_load=d.get("strict_load", True),
+    )

meshflow/models/meshflow_vae.py

@@ -0,0 +1,786 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Any, Dict, List, Optional
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.utils.checkpoint import checkpoint
+
+from meshflow.utils.mesh import extract_mesh_from_verts_normals_edges
+
+
+class FP32LayerNorm(nn.LayerNorm):
+    """LayerNorm computed in float32 for numerical stability."""
+
+    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
+        origin_dtype = inputs.dtype
+        return F.layer_norm(
+            inputs.float(),
+            self.normalized_shape,
+            self.weight.float() if self.weight is not None else None,
+            self.bias.float() if self.bias is not None else None,
+            self.eps,
+        ).to(origin_dtype)
+
+
+class RMSNorm(nn.Module):
+    """RMSNorm on the last dimension."""
+
+    def __init__(self, dim: int, eps: float = 1e-5, elementwise_affine: bool = True):
+        super().__init__()
+        self.eps = eps
+        self.weight = nn.Parameter(torch.ones(dim)) if elementwise_affine else None
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        input_dtype = hidden_states.dtype
+        variance = hidden_states.to(torch.float32).pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        if self.weight is not None:
+            if self.weight.dtype in (torch.float16, torch.bfloat16):
+                hidden_states = hidden_states.to(self.weight.dtype)
+            hidden_states = hidden_states * self.weight
+        else:
+            hidden_states = hidden_states.to(input_dtype)
+        return hidden_states
+
+
+class GELUProj(nn.Module):
+    """Linear projection followed by GELU."""
+
+    def __init__(self, dim_in: int, dim_out: int, bias: bool = True):
+        super().__init__()
+        self.proj = nn.Linear(dim_in, dim_out, bias=bias)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        return F.gelu(self.proj(hidden_states))
+
+
+class FeedForward(nn.Module):
+    """Two-layer feed-forward network with GELU activation."""
+
+    def __init__(
+        self,
+        dim: int,
+        dropout: float = 0.0,
+        activation_fn: str = "gelu",
+        final_dropout: bool = False,
+        inner_dim: Optional[int] = None,
+        bias: bool = True,
+    ):
+        super().__init__()
+        if activation_fn != "gelu":
+            raise ValueError(f"Only activation_fn='gelu' is supported, got {activation_fn!r}")
+        inner_dim = int(dim * 4) if inner_dim is None else inner_dim
+        layers: List[nn.Module] = [
+            GELUProj(dim, inner_dim, bias=bias),
+            nn.Dropout(dropout),
+            nn.Linear(inner_dim, dim, bias=bias),
+        ]
+        if final_dropout:
+            layers.append(nn.Dropout(dropout))
+        self.net = nn.ModuleList(layers)
+
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        for module in self.net:
+            hidden_states = module(hidden_states)
+        return hidden_states
+
+
+class Attention(nn.Module):
+    """
+    Multi-head attention with optional QK normalization.
+    Uses parameter names to_q, to_k, to_v, norm_q, norm_k, and to_out.
+    """
+
+    def __init__(
+        self,
+        query_dim: int,
+        cross_attention_dim: Optional[int] = None,
+        heads: int = 8,
+        dim_head: int = 64,
+        dropout: float = 0.0,
+        bias: bool = False,
+        out_bias: bool = True,
+        qk_norm: Optional[str] = None,
+        eps: float = 1e-5,
+    ):
+        super().__init__()
+        self.heads = heads
+        self.dim_head = dim_head
+        inner_dim = dim_head * heads
+        kv_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
+
+        self.to_q = nn.Linear(query_dim, inner_dim, bias=bias)
+        self.to_k = nn.Linear(kv_dim, inner_dim, bias=bias)
+        self.to_v = nn.Linear(kv_dim, inner_dim, bias=bias)
+
+        if qk_norm is None:
+            self.norm_q = None
+            self.norm_k = None
+        elif qk_norm == "rms_norm":
+            self.norm_q = RMSNorm(dim_head, eps=eps)
+            self.norm_k = RMSNorm(dim_head, eps=eps)
+        elif qk_norm in ("layer_norm", "fp32_layer_norm"):
+            norm_cls = FP32LayerNorm if qk_norm == "fp32_layer_norm" else nn.LayerNorm
+            self.norm_q = norm_cls(dim_head, eps=eps)
+            self.norm_k = norm_cls(dim_head, eps=eps)
+        else:
+            raise ValueError(f"Unsupported qk_norm: {qk_norm!r}")
+
+        self.to_out = nn.Sequential(
+            nn.Linear(inner_dim, query_dim, bias=out_bias),
+            nn.Dropout(dropout),
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        batch_size, seq_len_q, _ = hidden_states.shape
+
+        if encoder_hidden_states is None:
+            encoder_hidden_states = hidden_states
+
+        query = self.to_q(hidden_states)
+        key = self.to_k(encoder_hidden_states)
+        value = self.to_v(encoder_hidden_states)
+
+        query = query.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
+        key = key.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
+        value = value.view(batch_size, -1, self.heads, self.dim_head).transpose(1, 2)
+
+        if self.norm_q is not None:
+            query = self.norm_q(query)
+        if self.norm_k is not None:
+            key = self.norm_k(key)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, seq_len_q, -1)
+        hidden_states = self.to_out[0](hidden_states)
+        hidden_states = self.to_out[1](hidden_states)
+        return hidden_states
+
+
+class BasicTransformerBlock(nn.Module):
+    def __init__(
+        self,
+        dim: int,
+        num_attention_heads: int,
+        use_self_attention: bool = True,
+        use_cross_attention: bool = False,
+        self_attention_norm_type: Optional[str] = "layer_norm",
+        cross_attention_dim: Optional[int] = None,
+        cross_attention_norm_type: Optional[str] = None,
+        dropout: float = 0.0,
+        activation_fn: str = "gelu",
+        attention_bias: bool = False,
+        norm_elementwise_affine: bool = True,
+        norm_eps: float = 1e-5,
+        final_dropout: bool = False,
+        ff_inner_dim: Optional[int] = None,
+        ff_bias: bool = True,
+        attention_out_bias: bool = True,
+        qk_norm: Optional[str] = None,
+    ):
+        super().__init__()
+        head_dim = dim // num_attention_heads
+        attn_kwargs = dict(
+            query_dim=dim,
+            heads=num_attention_heads,
+            dim_head=head_dim,
+            dropout=dropout,
+            bias=attention_bias,
+            out_bias=attention_out_bias,
+            qk_norm=qk_norm,
+            eps=norm_eps,
+        )
+
+        def _make_norm(norm_type: Optional[str]) -> nn.Module:
+            if norm_type == "fp32_layer_norm":
+                return FP32LayerNorm(dim, norm_eps, norm_elementwise_affine)
+            return nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+
+        if use_self_attention:
+            self.norm1 = _make_norm(self_attention_norm_type)
+            self.attn1 = Attention(**attn_kwargs)
+        else:
+            self.norm1 = None
+            self.attn1 = None
+
+        if use_cross_attention:
+            cross_norm = cross_attention_norm_type or self_attention_norm_type
+            self.norm2 = _make_norm(cross_norm)
+            self.attn2 = Attention(
+                cross_attention_dim=cross_attention_dim, **attn_kwargs
+            )
+        else:
+            self.norm2 = None
+            self.attn2 = None
+
+        self.norm3 = nn.LayerNorm(dim, norm_eps, norm_elementwise_affine)
+        self.ff = FeedForward(
+            dim,
+            dropout=dropout,
+            activation_fn=activation_fn,
+            final_dropout=final_dropout,
+            inner_dim=ff_inner_dim,
+            bias=ff_bias,
+        )
+
+    def forward(
+        self,
+        hidden_states: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        if self.attn1 is not None:
+            norm_hidden_states = self.norm1(hidden_states)
+            hidden_states = (
+                self.attn1(norm_hidden_states, attention_mask=attention_mask)
+                + hidden_states
+            )
+            if hidden_states.ndim == 4:
+                hidden_states = hidden_states.squeeze(1)
+
+        if self.attn2 is not None:
+            norm_hidden_states = self.norm2(hidden_states)
+            hidden_states = (
+                self.attn2(
+                    norm_hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    attention_mask=encoder_attention_mask,
+                )
+                + hidden_states
+            )
+
+        norm_hidden_states = self.norm3(hidden_states)
+        return self.ff(norm_hidden_states) + hidden_states
+
+
+class DiagonalGaussianDistribution:
+    def __init__(self, parameters, deterministic=False, feat_dim=1):
+        self.feat_dim = feat_dim
+        if isinstance(parameters, list):
+            self.mean, self.logvar = parameters[0], parameters[1]
+        else:
+            self.mean, self.logvar = torch.chunk(parameters, 2, dim=feat_dim)
+        self.logvar = torch.clamp(self.logvar, -30.0, 20.0)
+        self.deterministic = deterministic
+        self.std = torch.exp(0.5 * self.logvar)
+        self.var = torch.exp(self.logvar)
+        if deterministic:
+            self.var = self.std = torch.zeros_like(self.mean)
+
+    def sample(self):
+        return self.mean + self.std * torch.randn_like(self.mean)
+
+    def kl(self, other=None, dims=(1, 2)):
+        if self.deterministic:
+            return torch.tensor(0.0, device=self.mean.device)
+        if other is None:
+            return 0.5 * torch.mean(
+                torch.pow(self.mean, 2) + self.var - 1.0 - self.logvar, dim=dims
+            )
+        return 0.5 * torch.mean(
+            torch.pow(self.mean - other.mean, 2) / other.var
+            + self.var / other.var
+            - 1.0
+            - self.logvar
+            + other.logvar,
+            dim=dims,
+        )
+
+    def mode(self):
+        return self.mean
+
+
+class FourierEmbedder(nn.Module):
+    def __init__(
+        self,
+        num_freqs: int = 12,
+        logspace: bool = True,
+        input_dim: int = 3,
+        include_input: bool = True,
+        include_pi: bool = False,
+    ):
+        super().__init__()
+        if logspace:
+            frequencies = 2.0 ** torch.arange(num_freqs, dtype=torch.float32)
+        else:
+            frequencies = torch.linspace(
+                1.0, 2.0 ** (num_freqs - 1), num_freqs, dtype=torch.float32
+            )
+        if include_pi:
+            frequencies *= torch.pi
+
+        self.register_buffer("frequencies", frequencies, persistent=False)
+        self.include_input = include_input
+        self.num_freqs = num_freqs
+        temp = 1 if include_input or num_freqs == 0 else 0
+        self.out_dim = input_dim * (num_freqs * 2 + temp)
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        if self.num_freqs == 0:
+            return x
+        embed = (x[..., None].contiguous() * self.frequencies).view(*x.shape[:-1], -1)
+        if self.include_input:
+            return torch.cat((x, embed.sin(), embed.cos()), dim=-1)
+        return torch.cat((embed.sin(), embed.cos()), dim=-1)
+
+
+class Encoder(nn.Module):
+    def __init__(
+        self,
+        num_enc_latents: int,
+        in_channels: int,
+        num_latents: int = 1024,
+        latents_dim: int = 64,
+        dim: int = 512,
+        num_attention_heads: int = 8,
+        num_layers: int = 8,
+        qk_norm: Optional[str] = None,
+    ):
+        super().__init__()
+        self.num_enc_latents = num_enc_latents
+        self.in_channels = in_channels
+        self.num_latents = num_latents
+        self.latents_dim = latents_dim
+        self.dim = dim
+        self.num_attention_heads = num_attention_heads
+        self.qk_norm = qk_norm
+
+        assert num_enc_latents % num_latents == 0
+        self.downsample_ratio = num_enc_latents // num_latents
+
+        self.proj_in = nn.Linear(self.in_channels, dim, bias=True)
+        self.downsample_mlp = nn.Linear(
+            self.dim * self.downsample_ratio,
+            self.dim,
+            bias=True,
+        )
+
+        self.cross_attn_block = BasicTransformerBlock(
+            dim=dim,
+            num_attention_heads=num_attention_heads,
+            use_self_attention=False,
+            use_cross_attention=True,
+            cross_attention_dim=dim,
+            cross_attention_norm_type="fp32_layer_norm",
+            qk_norm=qk_norm,
+        )
+
+        self.blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    dim=dim,
+                    num_attention_heads=num_attention_heads,
+                    use_self_attention=True,
+                    self_attention_norm_type="fp32_layer_norm",
+                    use_cross_attention=False,
+                    qk_norm=qk_norm,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm_out = FP32LayerNorm(dim)
+
+    def _get_init_query(self, hidden_states, pc):
+        bs = hidden_states.shape[0]
+        pc = pc.reshape(bs, -1, 3 * self.downsample_ratio)
+        q = self.downsample_mlp(
+            hidden_states.reshape(bs, -1, self.dim * self.downsample_ratio)
+        )
+        return q, pc
+
+    def _forward(
+        self,
+        x: torch.FloatTensor,
+        pc: torch.FloatTensor,
+    ):
+        _, num_tokens, _ = x.shape
+        assert self.num_enc_latents == num_tokens
+
+        hidden_states = self.proj_in(x)
+        query, pc = self._get_init_query(hidden_states, pc)
+
+        hidden_states = self.cross_attn_block(
+            query,
+            encoder_hidden_states=hidden_states,
+        )
+
+        for block in self.blocks:
+            hidden_states = block(hidden_states, attention_mask=None)
+        hidden_states = self.norm_out(hidden_states)
+
+        return hidden_states, pc
+
+    def forward(
+        self,
+        x: torch.FloatTensor,
+        pc: torch.FloatTensor,
+    ):
+        if self.training:
+            return checkpoint(self._forward, x, pc, use_reentrant=False)
+        return self._forward(x, pc)
+
+
+class Decoder(nn.Module):
+    def __init__(
+        self,
+        num_dec_latents: int,
+        num_latents: int = 1024,
+        adjacency_emb_dim: int = 32,
+        dim: int = 512,
+        num_attention_heads: int = 8,
+        num_layers: int = 16,
+        qk_norm: Optional[str] = "rms_norm",
+        use_verts_head: bool = True,
+        use_verts_normal_head: bool = True,
+        use_mask_head: bool = True,
+        use_adjacency_head: bool = True,
+    ):
+        super().__init__()
+        self.num_dec_latents = num_dec_latents
+        self.num_latents = num_latents
+        self.adjacency_emb_dim = adjacency_emb_dim
+        self.dim = dim
+        self.num_attention_heads = num_attention_heads
+        self.num_layers = num_layers
+        self.qk_norm = qk_norm
+        self.use_verts_head = use_verts_head
+        self.use_verts_normal_head = use_verts_normal_head
+        self.use_mask_head = use_mask_head
+        self.use_adjacency_head = use_adjacency_head
+
+        assert num_dec_latents % num_latents == 0
+        self.upsample_ratio = num_dec_latents // num_latents
+        self.upsample_mlp = nn.Linear(
+            self.dim // self.upsample_ratio,
+            self.dim,
+            bias=True,
+        )
+
+        self.cross_attn_block = BasicTransformerBlock(
+            dim=dim,
+            num_attention_heads=num_attention_heads,
+            use_self_attention=False,
+            use_cross_attention=True,
+            cross_attention_norm_type="fp32_layer_norm",
+            qk_norm=qk_norm,
+        )
+
+        self.blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    dim=dim,
+                    num_attention_heads=num_attention_heads,
+                    use_self_attention=True,
+                    self_attention_norm_type="fp32_layer_norm",
+                    use_cross_attention=False,
+                    qk_norm=qk_norm,
+                )
+                for _ in range(num_layers)
+            ]
+        )
+        self.norm_out = FP32LayerNorm(dim)
+
+        if use_mask_head:
+            self.mask_head = nn.Sequential(
+                nn.Linear(dim, dim // 2),
+                nn.LeakyReLU(),
+                nn.Linear(dim // 2, 1),
+            )
+        if use_verts_head:
+            self.verts_head = nn.Sequential(
+                nn.Linear(dim, dim // 2), nn.LeakyReLU(), nn.Linear(dim // 2, 3)
+            )
+        if use_verts_normal_head:
+            self.verts_normal_head = nn.Sequential(
+                nn.Linear(dim, dim // 2),
+                nn.LeakyReLU(),
+                nn.Linear(dim // 2, 3),
+                nn.Tanh(),
+            )
+        if use_adjacency_head:
+            self.adjacency_head = nn.Sequential(
+                nn.Linear(dim, dim // 2),
+                nn.LeakyReLU(),
+                nn.Linear(dim // 2, adjacency_emb_dim),
+            )
+
+    def _get_init_query(self, hidden_states):
+        bs = hidden_states.shape[0]
+        return self.upsample_mlp(
+            hidden_states.reshape(bs, -1, self.dim // self.upsample_ratio)
+        )
+
+    def _forward_features(self, latents: torch.Tensor) -> Dict[str, torch.Tensor]:
+        for block in self.blocks:
+            latents = block(latents, attention_mask=None)
+
+        query = self._get_init_query(latents)
+        latents = self.cross_attn_block(query, encoder_hidden_states=latents)
+        out = self.norm_out(latents)
+
+        ret: Dict[str, torch.Tensor] = {}
+        if self.use_mask_head:
+            ret["mask"] = self.mask_head(out)
+        if self.use_verts_head:
+            ret["verts"] = self.verts_head(out)
+        if self.use_verts_normal_head:
+            ret["verts_normal"] = self.verts_normal_head(out)
+        if self.use_adjacency_head:
+            ret["adjacency_emb"] = self.adjacency_head(out)
+        return ret
+
+    def forward(self, latents: torch.Tensor) -> Dict[str, torch.Tensor]:
+        if self.training:
+            return checkpoint(self._forward_features, latents, use_reentrant=False)
+        return self._forward_features(latents)
+
+
+class MeshFlowVAE(nn.Module):
+    def __init__(
+        self,
+        *,
+        n_samples: int,
+        num_latents: int = 4096,
+        latents_dim: int = 64,
+        width_encoder: int = 1024,
+        width_decoder: int = 1024,
+        num_attention_heads: int = 16,
+        num_layers_encoder: int = 8,
+        num_layers_decoder: int = 8,
+        qk_norm: Optional[str] = "rms_norm",
+        adjacency_emb_dim: int = 32,
+        max_degree: int = 50,
+        point_feats_type: List[str] = ["verts_normal"],
+        embedder_kwargs: Optional[Dict[str, Any]] = None,
+        use_verts_head: bool = True,
+        use_verts_normal_head: bool = True,
+        use_mask_head: bool = True,
+        use_adjacency_head: bool = True,
+        use_learnable_tau: bool = False,
+        init_tau: float = 0.6,
+        adjacency_emb_rescale: float = 1.0,
+        mean: Optional[float] = None,
+        std: Optional[float] = None,
+        pretrained_model_name_or_path: Optional[str] = None,
+        strict_load: bool = True,
+    ):
+        super().__init__()
+        self.latents_dim = latents_dim
+        self.width_encoder = width_encoder
+        self.point_feats_type = point_feats_type
+        self.adjacency_emb_rescale = float(adjacency_emb_rescale)
+        self.register_buffer(
+            "latent_mean",
+            torch.tensor(0.0 if mean is None else float(mean), dtype=torch.float32),
+            persistent=False,
+        )
+        self.register_buffer(
+            "latent_std",
+            torch.tensor(1.0 if std is None else float(std), dtype=torch.float32),
+            persistent=False,
+        )
+
+        if latents_dim > 0:
+            self.post_kl = nn.Linear(latents_dim, width_decoder)
+            self.latent_shape = (num_latents, latents_dim)
+        else:
+            self.latent_shape = (num_latents, width_encoder)
+
+        self.embedder = FourierEmbedder(**(embedder_kwargs or {}))
+        if latents_dim > 0:
+            self.pre_kl = nn.Linear(width_encoder, latents_dim * 2)
+
+        self.feature_dim = {
+            "verts": 3,
+            "verts_normal": 3,
+            "neighbor_points": max_degree * 3,
+            "adjacency_matrix": n_samples,
+            "degree": 1,
+            "verts_mask": 1,
+            "adjacency_emb": adjacency_emb_dim,
+        }
+        in_channels = self.embedder.out_dim
+        for k in self.point_feats_type:
+            key = k.split("@")[-1]
+            if key not in self.feature_dim:
+                raise KeyError(
+                    f"Unknown point feature '{key}' in point_feats_type; "
+                    f"known keys: {sorted(self.feature_dim)}"
+                )
+            d = self.feature_dim[key]
+            if k.startswith("embed@"):
+                temp = 1 if self.embedder.include_input or self.embedder.num_freqs == 0 else 0
+                in_channels += d * (self.embedder.num_freqs * 2 + temp)
+            else:
+                in_channels += d
+        self.in_channels = in_channels
+
+        self.encoder = Encoder(
+            num_enc_latents=n_samples,
+            in_channels=self.in_channels,
+            num_latents=num_latents,
+            latents_dim=latents_dim,
+            dim=width_encoder,
+            num_attention_heads=num_attention_heads,
+            num_layers=num_layers_encoder,
+            qk_norm=qk_norm,
+        )
+        self.decoder = Decoder(
+            num_dec_latents=n_samples,
+            num_latents=num_latents,
+            adjacency_emb_dim=adjacency_emb_dim,
+            dim=width_decoder,
+            num_attention_heads=num_attention_heads,
+            num_layers=num_layers_decoder,
+            qk_norm=qk_norm,
+            use_verts_head=use_verts_head,
+            use_verts_normal_head=use_verts_normal_head,
+            use_mask_head=use_mask_head,
+            use_adjacency_head=use_adjacency_head,
+        )
+        self.tau = nn.Parameter(torch.ones(1) * init_tau) if use_learnable_tau else init_tau
+
+        if pretrained_model_name_or_path:
+            self.load_pretrained_model(pretrained_model_name_or_path, strict=strict_load)
+
+    def load_pretrained_model(self, path: str, strict: bool = True) -> None:
+        ckpt = torch.load(path, map_location="cpu")
+        state = ckpt["state_dict"] if isinstance(ckpt, dict) and "state_dict" in ckpt else ckpt
+        filtered = {}
+        for k, v in state.items():
+            for prefix in ("mesh_model.vae.", "mesh_model."):
+                if k.startswith(prefix):
+                    filtered[k[len(prefix) :]] = v
+                    break
+        if filtered:
+            self.load_state_dict(filtered, strict=strict)
+
+    @property
+    def dtype(self):
+        return next(self.parameters()).dtype
+
+    def normalize(self, z: torch.Tensor) -> torch.Tensor:
+        return (z - self.latent_mean.to(z)) / self.latent_std.to(z)
+
+    def denormalize(self, z: torch.Tensor) -> torch.Tensor:
+        return z * self.latent_std.to(z) + self.latent_mean.to(z)
+
+    def encode(self, x, pc, sample_posterior=True):
+        shape_latents, encode_pc = self.encoder(x, pc)
+        assert shape_latents.shape[-1] == self.width_encoder
+        if self.latents_dim > 0:
+            posterior = DiagonalGaussianDistribution(self.pre_kl(shape_latents), feat_dim=-1)
+            z = posterior.sample() if sample_posterior else posterior.mode()
+        else:
+            z, posterior = shape_latents, None
+        return {
+            "shape_latents": shape_latents,
+            "encode_pc": encode_pc,
+            "z": z,
+            "posterior": posterior,
+        }
+
+    def decode(self, latents):
+        if self.latents_dim > 0:
+            latents = self.post_kl(latents)
+        return self.decoder(latents)
+
+    def collect_input(self, batch: dict) -> torch.Tensor:
+        parts = [self.embedder(batch["padded_verts"])]
+        for k in self.point_feats_type:
+            raw_key = k.split("@")[-1]
+            v = batch[raw_key]
+            if k.startswith("embed@"):
+                v = self.embedder(v)
+            if raw_key == "adjacency_emb" and self.adjacency_emb_rescale != 1.0:
+                v = v * self.adjacency_emb_rescale
+            parts.append(v)
+        return torch.cat(parts, dim=-1)
+
+    def forward(self, batch, sample_posterior=False, **kwargs):
+        pc = batch["padded_verts"]
+        enc = self.encode(self.collect_input(batch), pc, sample_posterior=sample_posterior)
+        return {**enc, "pc": pc, "dec": self.decode(enc["z"])}
+
+    @staticmethod
+    def compute_spacetime_distances(adjacency_embeddings: torch.Tensor, k_s: int) -> torch.Tensor:
+        x_s = adjacency_embeddings[..., :k_s]
+        x_t = adjacency_embeddings[..., k_s:]
+        return torch.cdist(x_s, x_s, p=2) ** 2 - torch.cdist(x_t, x_t, p=2) ** 2
+
+    def extract_mesh(
+        self,
+        pred: Dict[str, torch.Tensor],
+        fill_holes: bool = False,
+        max_cycle_size: int = 20,
+        return_full: bool = False,
+    ) -> List:
+        """Extract one mesh per batch item. pred tensors are [B, N, ...] or [N, ...]."""
+        if pred["mask"].dim() == 2:
+            pred = {k: v.unsqueeze(0) for k, v in pred.items()}
+
+        out: List = []
+        for i in range(pred["mask"].shape[0]):
+            pred_i = {k: v[i] for k, v in pred.items()}
+            mask = pred_i["mask"].squeeze(-1) > 0.0
+            verts = pred_i["verts"][mask].float()
+            vnorm = pred_i["verts_normal"][mask].float().clamp(-1, 1)
+            adj = pred_i["adjacency_emb"][mask].float()
+            if self.adjacency_emb_rescale != 1.0:
+                adj = adj / self.adjacency_emb_rescale
+
+            k_s = adj.shape[-1] // 2
+            dist = self.compute_spacetime_distances(adj, k_s)
+            triu = torch.triu(torch.ones_like(dist, dtype=torch.bool), diagonal=1)
+            tau = self.tau if isinstance(self.tau, (int, float)) else float(self.tau.item())
+            edges = torch.nonzero((torch.sigmoid(dist - tau) > 0.5) & triu, as_tuple=False)
+
+            mesh = extract_mesh_from_verts_normals_edges(
+                verts, vnorm, edges, fill_holes=fill_holes, max_cycle_size=max_cycle_size
+            )
+            out.append((mesh, verts, vnorm, edges) if return_full else mesh)
+        return out
+
+
+def build_mesh_model(cfg: dict) -> MeshFlowVAE:
+    """Build MeshFlowVAE from a mesh_model config dict."""
+    return MeshFlowVAE(
+        n_samples=cfg["n_samples"],
+        num_latents=cfg.get("num_latents", 4096),
+        latents_dim=cfg.get("latents_dim", 64),
+        width_encoder=cfg.get("width_encoder", 1024),
+        width_decoder=cfg.get("width_decoder", 1024),
+        num_attention_heads=cfg.get("num_attention_heads", 16),
+        num_layers_encoder=cfg.get("num_layers_encoder", 8),
+        num_layers_decoder=cfg.get("num_layers_decoder", 8),
+        qk_norm=cfg.get("qk_norm", "rms_norm"),
+        adjacency_emb_dim=cfg.get("adjacency_emb_dim", 32),
+        max_degree=cfg.get("max_degree", 50),
+        point_feats_type=cfg.get("point_feats_type", ["verts_normal"]),
+        embedder_kwargs=cfg.get("embedder_kwargs"),
+        use_verts_head=cfg.get("use_verts_head", True),
+        use_verts_normal_head=cfg.get("use_verts_normal_head", True),
+        use_mask_head=cfg.get("use_mask_head", True),
+        use_adjacency_head=cfg.get("use_adjacency_head", True),
+        use_learnable_tau=cfg.get("use_learnable_tau", False),
+        init_tau=cfg.get("init_tau", 0.6),
+        adjacency_emb_rescale=cfg.get("adjacency_emb_rescale", 1.0),
+        mean=cfg.get("mean"),
+        std=cfg.get("std"),
+        pretrained_model_name_or_path=cfg.get("pretrained_model_name_or_path"),
+        strict_load=cfg.get("strict_load", True),
+    )

meshflow/pipelines/__init__.py

@@ -0,0 +1,10 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .meshflow_pipeline import MeshFlowPipeline
+from .utils import flow_sample
+
+__all__ = ["MeshFlowPipeline", "flow_sample"]

meshflow/pipelines/meshflow_pipeline.py

@@ -0,0 +1,463 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Callable, Optional, Union
+
+import torch
+import torch.nn as nn
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from omegaconf import OmegaConf
+from PIL import Image
+
+from ..models.condition_encoder import (
+    build_dinov3_encoder,
+    make_empty_visual_embeds,
+    resolve_visual_embed_dim,
+)
+from ..models.meshflow_dit import build_meshflow_dit, voxelize_pc
+from ..models.meshflow_vae import build_mesh_model
+from ..utils.dtype import resolve_torch_dtype, torch_dtype_to_name
+from ..utils.mesh import Mesh
+from .utils import flow_sample, preprocess_image as rmbg_preprocess_image
+
+
+class MeshFlowPipeline:
+    """
+    MeshFlow DiT image-to-mesh inference pipeline.
+
+    Conditions:
+      - RoPE: surface points sampled from an input mesh, voxelized to 3D indices
+      - Visual: optional reference image encoded by the visual encoder (or empty tokens)
+    """
+
+    model_names_to_load = ["mesh_model", "visual_encoder", "denoiser"]
+
+    def __init__(
+        self,
+        models: dict[str, nn.Module] = None,
+        scheduler: FlowMatchEulerDiscreteScheduler = None,
+        num_verts: int = 4096,
+        voxelize_resolution: int = 32,
+        num_inference_steps: int = 28,
+        guidance_scale: float = 2.5,
+        dtype: torch.dtype = torch.float16,
+        use_rmbg: bool = False,
+        rmbg_preprocess: Optional[Callable[[Image.Image], Image.Image]] = None,
+    ):
+        if models is None:
+            return
+
+        self.models = models
+        self.scheduler = scheduler
+        self.num_verts = int(num_verts)
+        self.voxelize_resolution = int(voxelize_resolution)
+        self.num_inference_steps = int(num_inference_steps)
+        self.guidance_scale = float(guidance_scale)
+        self.dtype = dtype
+        self.use_rmbg = use_rmbg
+        self.rmbg_preprocess = rmbg_preprocess or rmbg_preprocess_image
+        self._device = torch.device("cpu")
+        self._visual_encoder_type: str | None = None
+        self._visual_encoder_cfg: dict | None = None
+        self._empty_visual_embeds: torch.Tensor | None = None
+        self._compile_models = False
+        self.num_latents: int | None = None
+
+    @property
+    def device(self) -> torch.device:
+        return self._device
+
+    @classmethod
+    def from_pretrained(
+        cls,
+        model_path: str,
+        device: Union[str, torch.device] = "cuda",
+        dtype: Union[str, torch.dtype] = "fp16",
+        compile_models: bool = False,
+        num_verts: Optional[int] = None,
+    ) -> "MeshFlowPipeline":
+        """
+        Load MeshFlow DiT + VAE + visual encoder from a model bundle directory.
+
+        Args:
+            model_path: Directory that must contain `config.yaml` and `model.pth`.
+            device: Target torch device.
+            dtype: `bf16`, `fp16`, or `fp32`.
+            compile_models: Whether to `torch.compile` hot paths (CUDA only).
+            num_verts: Target value for ``proj_cond_on_temb`` (``num_verts / num_latents``).
+                Also resizes the latent sequence when ``use_proj_cond_on_temb`` is enabled.
+                ``num_latents`` for normalization is always read from
+                ``mesh_model.num_latents`` in config.
+        """
+        root = Path(model_path)
+        if not root.is_dir():
+            raise NotADirectoryError(f"model_path must be a directory: {root}")
+
+        config_path = root / "config.yaml"
+        ckpt_path = root / "model.pth"
+        if not config_path.is_file():
+            raise FileNotFoundError(f"Config not found: {config_path}")
+        if not ckpt_path.is_file():
+            raise FileNotFoundError(f"Checkpoint not found: {ckpt_path}")
+
+        print(f"[MeshFlowPipeline] Loading from {root.resolve()}")
+        print(f"  config: {config_path.name}")
+        print(f"  weights: {ckpt_path.name} ({ckpt_path.stat().st_size / (1024 ** 3):.2f} GB)")
+
+        cfg = OmegaConf.load(config_path)
+        data_cfg, scfg = cfg.data, cfg.system
+
+        config_num_verts = int(data_cfg.n_verts)
+        n_samples = int(scfg.mesh_model.n_samples)
+        input_size = int(scfg.denoiser_model.input_size)
+        use_proj_cond_on_temb = bool(scfg.denoiser_model.get("use_proj_cond_on_temb", False))
+
+        num_latents = int(scfg.mesh_model.num_latents)
+        if use_proj_cond_on_temb:
+            if num_verts is not None:
+                num_verts = int(num_verts)
+                config_num_verts = n_samples = input_size = num_verts
+        else:
+            if num_verts is not None:
+                print(
+                    "[MeshFlowPipeline] Ignoring num_verts: "
+                    "denoiser_model.use_proj_cond_on_temb is disabled in config"
+                )
+            num_verts = None
+            num_latents = None
+
+        dtype = resolve_torch_dtype(dtype)
+        device = torch.device(device if torch.cuda.is_available() or device == "cpu" else "cpu")
+        if num_verts is not None:
+            print(
+                f"  num_verts control: num_verts={config_num_verts} "
+                f"n_samples={n_samples} input_size={input_size}"
+            )
+
+        print(f"[MeshFlowPipeline] Reading checkpoint on CPU...")
+        ckpt = torch.load(ckpt_path, map_location="cpu", weights_only=False)
+        state = ckpt["state_dict"] if isinstance(ckpt, dict) and "state_dict" in ckpt else ckpt
+        print(f"  state_dict: {len(state)} tensors")
+
+        mesh_cfg = OmegaConf.to_container(scfg.mesh_model, resolve=True)
+        mesh_cfg["pretrained_model_name_or_path"] = str(ckpt_path)
+        mesh_cfg["n_samples"] = n_samples
+        if num_verts is not None:
+            mesh_cfg["num_latents"] = num_verts
+        print(f"[MeshFlowPipeline] Building mesh_model (VAE)...")
+        mesh_model = build_mesh_model(mesh_cfg).to(device).eval()
+        print(f"  mesh_model: {type(mesh_model).__name__} on {device}")
+
+        visual_cfg = OmegaConf.to_container(scfg.visual_condition, resolve=True)
+        embed_dim = resolve_visual_embed_dim(
+            visual_cfg,
+            fallback=int(scfg.denoiser_model.visual_condition_dim),
+        )
+        empty_visual_embeds = make_empty_visual_embeds(int(visual_cfg["image_size"]), embed_dim)
+        print(
+            "[MeshFlowPipeline] visual_encoder deferred "
+            f"(loads on first reference image; zero cond shape={tuple(empty_visual_embeds.shape)})"
+        )
+
+        print(f"[MeshFlowPipeline] Building denoiser (DiT)...")
+        denoiser_cfg = OmegaConf.to_container(scfg.denoiser_model, resolve=True)
+        denoiser_cfg["pretrained_model_name_or_path"] = str(ckpt_path)
+        denoiser_cfg["input_size"] = input_size
+        denoiser = build_meshflow_dit(denoiser_cfg).to(device).eval()
+        if hasattr(denoiser, "set_rope_input_resolution"):
+            denoiser.set_rope_input_resolution(int(scfg.voxelize_resolution))
+        print(
+            f"  denoiser: {type(denoiser).__name__} "
+            f"(layers={denoiser.depth}, input_size={denoiser.input_size}) on {device}"
+        )
+
+        if compile_models:
+            if device.type != "cuda":
+                print("[MeshFlowPipeline] torch.compile skipped (CUDA required)")
+            else:
+                print("[MeshFlowPipeline] torch.compile enabled:")
+                print("  compiling mesh_model.decoder ...")
+                mesh_model.decoder = torch.compile(mesh_model.decoder)
+                print("  compiling denoiser ...")
+                denoiser = torch.compile(denoiser)
+                print("[MeshFlowPipeline] torch.compile done (first run will warm up)")
+
+        scheduler = FlowMatchEulerDiscreteScheduler(
+            **OmegaConf.to_container(scfg.denoise_scheduler, resolve=True)
+        )
+
+        pipeline = cls(
+            models={
+                "mesh_model": mesh_model,
+                "visual_encoder": None,
+                "denoiser": denoiser,
+            },
+            scheduler=scheduler,
+            num_verts=config_num_verts,
+            voxelize_resolution=int(scfg.voxelize_resolution),
+            num_inference_steps=int(scfg.num_inference_steps),
+            guidance_scale=float(scfg.guidance_scale),
+            dtype=dtype,
+        )
+        pipeline._device = device
+        pipeline._visual_encoder_type = str(scfg.visual_condition_type)
+        pipeline._visual_encoder_cfg = visual_cfg
+        pipeline._empty_visual_embeds = empty_visual_embeds
+        pipeline._compile_models = compile_models
+        pipeline.num_latents = num_latents
+        num_latents_repr = num_latents if num_latents is not None else "n/a"
+        print(
+            f"[MeshFlowPipeline] Ready  device={device} dtype={torch_dtype_to_name(dtype)} "
+            f"num_verts={config_num_verts} num_latents={num_latents_repr} "
+            f"use_proj_cond_on_temb={use_proj_cond_on_temb} "
+            f"steps={pipeline.num_inference_steps} "
+            f"guidance_scale={pipeline.guidance_scale} compile={compile_models}"
+        )
+        return pipeline
+
+    def to(self, device: Union[str, torch.device]) -> "MeshFlowPipeline":
+        device = torch.device(device)
+        self._device = device
+        for model in self.models.values():
+            if model is not None:
+                model.to(device)
+        return self
+
+    def _ensure_visual_encoder(self) -> nn.Module:
+        visual_encoder = self.models.get("visual_encoder")
+        if visual_encoder is not None:
+            return visual_encoder
+        if self._visual_encoder_type is None or self._visual_encoder_cfg is None:
+            raise RuntimeError("visual encoder config is not initialized")
+
+        print(f"[MeshFlowPipeline] Loading visual_encoder ({self._visual_encoder_type})...")
+        visual_encoder = build_dinov3_encoder(
+            self._visual_encoder_type,
+            self._visual_encoder_cfg,
+        ).to(self._device).eval()
+        if self._compile_models and self._device.type == "cuda":
+            print("  compiling visual_encoder.dino_model ...")
+            visual_encoder.dino_model = torch.compile(visual_encoder.dino_model)
+        self.models["visual_encoder"] = visual_encoder
+        print(f"  visual_encoder: {type(visual_encoder).__name__} on {self._device}")
+        return visual_encoder
+
+    def preprocess_image(self, image: Image.Image) -> Image.Image:
+        """Optional RMBG + foreground crop for reference images."""
+        if not self.use_rmbg:
+            return image.convert("RGB")
+        output = self.rmbg_preprocess(image)
+        if isinstance(output, list):
+            output = output[0]
+        return output.convert("RGB")
+
+    def sample_surface_points(
+        self,
+        mesh: Union[str, Mesh],
+        num_verts: Optional[int] = None,
+    ) -> torch.Tensor:
+        """
+        Build points for RoPE conditioning.
+
+        Meshes are surface-sampled; point clouds are normalized and randomly resampled directly,
+        then voxelized in ``get_rope_cond``.
+
+        Returns:
+            Tensor of shape (1, N, 3) on pipeline device.
+        """
+        num_verts = num_verts or self.num_verts
+        if isinstance(mesh, str):
+            surface_pts = Mesh.load_rope_points(
+                mesh,
+                n_verts=num_verts,
+                normalize_by="bsphere",
+                obj_scale=2.0,
+                preprocess=True,
+                device=self.device,
+            )
+        elif int(mesh.faces.shape[0]) == 0:
+            from ..utils.mesh import _normalize_points, _resample_point_cloud
+
+            points = _normalize_points(mesh.verts, normalize_by="bsphere", size=2.0)
+            surface_pts = _resample_point_cloud(points, num_verts)
+        else:
+            surface_pts, _ = mesh.sample_surface(num_verts)
+        return surface_pts.unsqueeze(0).to(self.device)
+
+    def get_rope_cond(self, surface_pc: torch.Tensor) -> dict:
+        """Voxelize surface points for RoPE conditioning."""
+        _, vox_idx = voxelize_pc(surface_pc, self.voxelize_resolution)
+        return {"rope_input": vox_idx.to(self.device)}
+
+    def get_visual_cond(
+        self,
+        image: Optional[Image.Image] = None,
+        guidance_scale: Optional[float] = None,
+        preprocess_image: bool = True,
+    ) -> torch.Tensor:
+        """
+        Build visual tokens for DiT cross-attention.
+
+        When `guidance_scale != 1`, returns concatenated [uncond, cond] for CFG.
+        """
+        guidance_scale = self.guidance_scale if guidance_scale is None else float(guidance_scale)
+        use_cfg = guidance_scale != 1.0
+        empty_visual_embeds = self._empty_visual_embeds
+        if empty_visual_embeds is None:
+            raise RuntimeError("empty visual embeds are not initialized")
+
+        with torch.autocast(device_type=self.device.type, enabled=False):
+            if image is not None:
+                if preprocess_image:
+                    image = self.preprocess_image(image)
+                visual_encoder = self._ensure_visual_encoder()
+                visual_cond = visual_encoder.encode_image(image.convert("RGB"))
+            else:
+                visual_cond = empty_visual_embeds
+            visual_cond = visual_cond.to(device=self.device, dtype=self.dtype)
+
+        if use_cfg:
+            uncond = empty_visual_embeds.to(visual_cond).expand(
+                visual_cond.shape[0], -1, -1
+            )
+            visual_cond = torch.cat([uncond, visual_cond], dim=0)
+        return visual_cond
+
+    def get_proj_cond_on_temb(
+        self,
+        num_verts: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+    ) -> Optional[torch.Tensor]:
+        """Build normalized vertex-count DiT control for timestep embedding projection."""
+        denoiser = self.models["denoiser"]
+        if not denoiser.use_proj_cond_on_temb:
+            return None
+        if self.num_latents is None:
+            raise RuntimeError("num_latents is required when use_proj_cond_on_temb is enabled")
+
+        num_verts = int(num_verts or self.num_verts)
+        guidance_scale = self.guidance_scale if guidance_scale is None else float(guidance_scale)
+        proj_cond = torch.tensor(
+            [[num_verts / self.num_latents]],
+            dtype=torch.float32,
+            device=self.device,
+        )
+        if guidance_scale != 1.0:
+            proj_cond = proj_cond.repeat(2, 1)
+        return proj_cond
+
+    def sample_latent(
+        self,
+        visual_cond: torch.Tensor,
+        joint_attention_kwargs: dict,
+        steps: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+        seed: int = 42,
+        disable_prog: bool = False,
+        num_verts: Optional[int] = None,
+    ) -> torch.Tensor:
+        """Run flow-matching sampling and return final latents."""
+        denoiser = self.models["denoiser"]
+        steps = self.num_inference_steps if steps is None else int(steps)
+        guidance_scale = self.guidance_scale if guidance_scale is None else float(guidance_scale)
+        use_cfg = guidance_scale != 1.0
+        sample_shape = (denoiser.input_size, denoiser.in_channels)
+        generator = torch.Generator(device=self.device).manual_seed(int(seed))
+        denoiser_model_kwargs = {}
+        proj_cond_on_temb = self.get_proj_cond_on_temb(
+            num_verts=num_verts,
+            guidance_scale=guidance_scale,
+        )
+        if proj_cond_on_temb is not None:
+            denoiser_model_kwargs["proj_cond_on_temb"] = proj_cond_on_temb
+
+        latents = None
+        with torch.autocast(device_type=self.device.type, dtype=self.dtype):
+            for latents, _ in flow_sample(
+                self.scheduler,
+                denoiser,
+                sample_shape,
+                steps=steps,
+                visual_cond=visual_cond,
+                guidance_scale=guidance_scale,
+                do_classifier_free_guidance=use_cfg,
+                generator=generator,
+                device=self.device,
+                disable_prog=disable_prog,
+                joint_attention_kwargs=joint_attention_kwargs,
+                denoiser_model_kwargs=denoiser_model_kwargs or None,
+            ):
+                pass
+        return latents
+
+    def decode_latent(
+        self,
+        latents: torch.Tensor,
+        fill_holes: bool = False,
+    ) -> Mesh:
+        """Decode sampled latents with the MeshFlow VAE and extract a mesh."""
+        mesh_model = self.models["mesh_model"]
+        with torch.autocast(device_type=self.device.type, dtype=self.dtype):
+            decoded = mesh_model.decode(mesh_model.denormalize(latents))
+            out_mesh = mesh_model.extract_mesh(
+                decoded,
+                fill_holes=fill_holes,
+            )[0]
+        return out_mesh
+
+    @torch.no_grad()
+    def run(
+        self,
+        mesh: Union[str, Mesh],
+        image: Optional[Image.Image] = None,
+        steps: Optional[int] = None,
+        guidance_scale: Optional[float] = None,
+        seed: int = 42,
+        preprocess_image: bool = True,
+        return_latent: bool = False,
+        disable_prog: bool = False,
+        num_verts: Optional[int] = None,
+    ):
+        """
+        End-to-end MeshFlow inference.
+
+        Args:
+            mesh: Input mesh/point-cloud path or `Mesh` instance for RoPE surface sampling.
+            image: Optional reference image for visual conditioning.
+            steps: Flow sampling steps (default: from config).
+            guidance_scale: CFG scale (default: from config).
+            seed: Random seed.
+            preprocess_image: Whether to run image preprocessing before encoding.
+            return_latent: If True, return `(mesh, latents)` instead of mesh only.
+            disable_prog: Disable tqdm progress bar during sampling.
+            num_verts: Override vertex count for proj_cond_on_temb only (default:
+                pipeline ``num_verts``). RoPE surface sampling always uses
+                ``pipeline.num_verts`` to match denoiser ``input_size``.
+        """
+        proj_num_verts = num_verts or self.num_verts
+        surface_pc = self.sample_surface_points(mesh, num_verts=self.num_verts)
+        joint_kwargs = self.get_rope_cond(surface_pc)
+        visual_cond = self.get_visual_cond(
+            image,
+            guidance_scale=guidance_scale,
+            preprocess_image=preprocess_image,
+        )
+        latents = self.sample_latent(
+            visual_cond,
+            joint_kwargs,
+            steps=steps,
+            guidance_scale=guidance_scale,
+            seed=seed,
+            disable_prog=disable_prog,
+            num_verts=proj_num_verts,
+        )
+        out_mesh = self.decode_latent(latents)
+        if return_latent:
+            return out_mesh, latents
+        return out_mesh

meshflow/pipelines/utils.py

@@ -0,0 +1,226 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import inspect
+import os
+import threading
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import torch
+from diffusers.schedulers import FlowMatchEulerDiscreteScheduler
+from PIL import Image
+from tqdm import tqdm
+from torchvision import transforms
+
+DEFAULT_RMBG_PATH = "/mnt/jfs-ssd/pretrained_models/RMBG-2.0"
+
+_rmbg_model = None
+_rmbg_lock = threading.Lock()
+
+
+def preprocess_image(
+    images: Union[Image.Image, List[Image.Image]],
+    background_color: Optional[List[int]] = None,
+    foreground_ratio: float = 0.9,
+    rmbg_path: str = DEFAULT_RMBG_PATH,
+) -> Union[Image.Image, List[Image.Image]]:
+    """Remove background (if needed), crop foreground, center and square-pad."""
+    global _rmbg_model
+
+    bg = tuple(background_color or [255, 255, 255])
+    single = isinstance(images, Image.Image)
+    batch = [images] if single else list(images)
+    out: List[Image.Image] = []
+
+    for im in batch:
+        if im.mode not in ("RGB", "RGBA"):
+            im = im.convert("RGBA" if "A" in im.mode else "RGB")
+
+        # Run matting when alpha channel is not fully opaque.
+        if not (im.mode == "RGBA" and im.getextrema()[3][0] >= 255):
+            with _rmbg_lock:
+                if _rmbg_model is None:
+                    if not os.path.isdir(rmbg_path):
+                        raise FileNotFoundError(f"RMBG model not found: {rmbg_path}")
+                    from transformers import AutoModelForImageSegmentation
+
+                    device = "cuda" if torch.cuda.is_available() else "cpu"
+                    _rmbg_model = AutoModelForImageSegmentation.from_pretrained(
+                        rmbg_path, trust_remote_code=True
+                    ).eval().to(device)
+
+            model = _rmbg_model
+            device = next(model.parameters()).device
+            rgb = im.convert("RGB")
+            norm = transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
+            inp = norm(transforms.ToTensor()(transforms.Resize((1024, 1024))(rgb))).unsqueeze(0).to(device)
+            with torch.autocast(device.type, enabled=device.type == "cuda"), torch.no_grad():
+                mask = model(inp)[-1].sigmoid()[0, 0].cpu()
+            im = rgb.copy()
+            im.putalpha(transforms.ToPILImage()(mask).resize(rgb.size))
+
+        alpha = im.split()[-1]
+        bbox = alpha.getbbox()
+        if bbox is None:
+            raise ValueError("No foreground found after background removal.")
+
+        w0, h0 = im.size
+        x1, y1, x2, y2 = bbox
+        dw, dh = x2 - x1, y2 - y1
+        scale = min(h0 * foreground_ratio / dh, w0 * foreground_ratio / dw)
+        tw, th = max(1, int(dw * scale)), max(1, int(dh * scale))
+
+        fg = Image.alpha_composite(Image.new("RGBA", im.size, (*bg, 255)), im).crop(bbox)
+        fg = fg.resize((tw, th))
+        alpha = alpha.crop(bbox).resize((tw, th))
+
+        canvas = Image.new("RGBA", (w0, h0), (*bg, 255))
+        canvas.paste(fg, ((w0 - tw) // 2, (h0 - th) // 2), alpha)
+        side = max(w0, h0)
+        if w0 != h0:
+            square = Image.new("RGBA", (side, side), (*bg, 255))
+            square.paste(canvas, ((side - w0) // 2, (side - h0) // 2))
+            canvas = square
+        out.append(canvas)
+
+    return out[0] if single else out
+
+
+def retrieve_timesteps(
+    scheduler,
+    num_inference_steps: Optional[int] = None,
+    device: Optional[Union[str, torch.device]] = None,
+    timesteps: Optional[List[int]] = None,
+    sigmas: Optional[List[float]] = None,
+    **kwargs,
+):
+    if timesteps is not None and sigmas is not None:
+        raise ValueError("Only one of `timesteps` or `sigmas` can be passed.")
+    if timesteps is not None:
+        accepts_timesteps = "timesteps" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accepts_timesteps:
+            raise ValueError(
+                f"{scheduler.__class__}'s `set_timesteps` does not support custom timesteps."
+            )
+        scheduler.set_timesteps(timesteps=timesteps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    elif sigmas is not None:
+        accept_sigmas = "sigmas" in set(
+            inspect.signature(scheduler.set_timesteps).parameters.keys()
+        )
+        if not accept_sigmas:
+            raise ValueError(
+                f"{scheduler.__class__}'s `set_timesteps` does not support custom sigmas."
+            )
+        scheduler.set_timesteps(sigmas=sigmas, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+        num_inference_steps = len(timesteps)
+    else:
+        scheduler.set_timesteps(num_inference_steps, device=device, **kwargs)
+        timesteps = scheduler.timesteps
+    return timesteps, num_inference_steps
+
+
+@torch.no_grad()
+def flow_sample(
+    scheduler: FlowMatchEulerDiscreteScheduler,
+    diffusion_model: torch.nn.Module,
+    shape: Union[List[int], Tuple[int]],
+    steps: int,
+    visual_cond: Optional[torch.Tensor] = None,
+    shape_cond: Optional[torch.Tensor] = None,
+    guidance_scale: float = 3.0,
+    do_classifier_free_guidance: bool = True,
+    generator: Optional[torch.Generator] = None,
+    device: torch.device = "cuda:0",
+    disable_prog: bool = True,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    denoiser_model_kwargs: Optional[Dict[str, Any]] = None,
+):
+    assert steps > 0, f"{steps} must > 0."
+
+    bsz_from_doubled_cond = False
+    if visual_cond is not None:
+        bsz = visual_cond.shape[0]
+        device = visual_cond.device
+        dtype = visual_cond.dtype
+        bsz_from_doubled_cond = do_classifier_free_guidance
+    elif shape_cond is not None:
+        bsz = shape_cond.shape[0]
+        device = shape_cond.device
+        dtype = shape_cond.dtype
+        bsz_from_doubled_cond = do_classifier_free_guidance
+    elif joint_attention_kwargs is not None:
+        bsz, device, dtype = None, None, None
+        if "rope_input" in joint_attention_kwargs:
+            v = joint_attention_kwargs["rope_input"]
+            if isinstance(v, torch.Tensor) and v.dim() >= 1:
+                bsz, device, dtype = v.shape[0], v.device, v.dtype
+        if bsz is None:
+            for v in joint_attention_kwargs.values():
+                if isinstance(v, torch.Tensor) and v.dim() >= 1:
+                    bsz, device, dtype = v.shape[0], v.device, v.dtype
+                    break
+        if bsz is None:
+            raise ValueError(
+                "flow_sample: joint_attention_kwargs must contain a batch-sized tensor."
+            )
+        bsz_from_doubled_cond = False
+    else:
+        raise ValueError("flow_sample: at least one condition tensor is required.")
+
+    if do_classifier_free_guidance and bsz_from_doubled_cond:
+        bsz = bsz // 2
+
+    latents = torch.randn((bsz, *shape), generator=generator, device=device, dtype=dtype)
+    try:
+        latents = latents * scheduler.init_noise_sigma
+    except AttributeError:
+        pass
+
+    joint_attention_kwargs_for_forward = joint_attention_kwargs
+    if joint_attention_kwargs_for_forward is not None and do_classifier_free_guidance:
+        joint_attention_kwargs_for_forward = {
+            k: (
+                torch.cat([v, v], dim=0)
+                if isinstance(v, torch.Tensor) and v.shape[0] == bsz
+                else v
+            )
+            for k, v in joint_attention_kwargs.items()
+        }
+
+    timesteps, _ = retrieve_timesteps(scheduler, steps + 1, device)
+    distance = (timesteps[:-1] - timesteps[1:]) / scheduler.config.num_train_timesteps
+
+    for i, t in enumerate(
+        tqdm(timesteps[:-1], disable=disable_prog, desc="Flow Sampling:", leave=False)
+    ):
+        latent_model_input = (
+            torch.cat([latents] * 2) if do_classifier_free_guidance else latents
+        )
+        timestep_tensor = torch.tensor([t], dtype=latents.dtype, device=device)
+        timestep_tensor = timestep_tensor.expand(latent_model_input.shape[0])
+        timestep_tensor = timestep_tensor / scheduler.config.num_train_timesteps
+
+        noise_pred = diffusion_model(
+            latent_model_input,
+            timestep_tensor,
+            visual_condition=visual_cond,
+            joint_attention_kwargs=joint_attention_kwargs_for_forward,
+            **(denoiser_model_kwargs or {}),
+        ).sample
+
+        if do_classifier_free_guidance:
+            noise_pred_uncond, noise_pred_cond = noise_pred.chunk(2)
+            noise_pred = noise_pred_uncond + guidance_scale * (
+                noise_pred_cond - noise_pred_uncond
+            )
+
+        latents = latents - distance[i] * noise_pred
+        yield latents, t

meshflow/utils/__init__.py

@@ -0,0 +1,6 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+

meshflow/utils/config.py

@@ -0,0 +1,75 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+import os
+from dataclasses import dataclass, field, fields
+from datetime import datetime
+from typing import Any, List, Optional
+
+from omegaconf import OmegaConf
+
+
+@dataclass
+class ExperimentConfig:
+    name: str = "default"
+    tag: str = ""
+    use_timestamp: bool = True
+    timestamp: Optional[str] = None
+    exp_root_dir: str = "outputs"
+    exp_dir: str = "outputs/default"
+    trial_name: str = "exp"
+    trial_dir: str = "outputs/default/exp"
+    n_gpus: int = 1
+    resume: Optional[str] = None
+    data: dict = field(default_factory=dict)
+    system_type: str = ""
+    system: dict = field(default_factory=dict)
+    trainer: dict = field(default_factory=dict)
+    checkpoint: dict = field(default_factory=dict)
+
+    def __post_init__(self):
+        if not self.tag and not self.use_timestamp:
+            raise ValueError("Either tag is specified or use_timestamp is True.")
+        self.trial_name = self.tag
+        if self.timestamp is None:
+            self.timestamp = ""
+            if self.use_timestamp:
+                if self.n_gpus > 1:
+                    print(
+                        "Timestamp is disabled when using multiple GPUs; set a unique tag."
+                    )
+                else:
+                    self.timestamp = datetime.now().strftime("@%Y%m%d-%H%M%S")
+        self.trial_name += self.timestamp
+        self.exp_dir = os.path.join(self.exp_root_dir, self.name)
+        self.trial_dir = os.path.join(self.exp_dir, self.trial_name)
+
+
+def _to_dict(cfg: Any) -> dict:
+    if isinstance(cfg, dict):
+        return cfg
+    return OmegaConf.to_container(cfg, resolve=True)
+
+
+def structured_config(cls, cfg: Any):
+    """Build a structured config object from a dataclass type and plain dict."""
+    return OmegaConf.structured(cls(**_to_dict(cfg)))
+
+
+def load_config(
+    path: str,
+    cli_args: Optional[List[str]] = None,
+    **overrides,
+) -> ExperimentConfig:
+    cfg = OmegaConf.load(path)
+    if cli_args:
+        cfg = OmegaConf.merge(cfg, OmegaConf.from_cli(cli_args))
+    if overrides:
+        cfg = OmegaConf.merge(cfg, overrides)
+    OmegaConf.resolve(cfg)
+    raw = _to_dict(cfg)
+    valid = {f.name for f in fields(ExperimentConfig)}
+    return ExperimentConfig(**{k: v for k, v in raw.items() if k in valid})

meshflow/utils/dtype.py

@@ -0,0 +1,44 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+from typing import Union
+
+import torch
+
+AUTOCAST_DTYPE_CHOICES = ("bf16", "fp16", "fp32")
+
+_DTYPE_ALIASES = {
+    "bf16": torch.bfloat16,
+    "fp16": torch.float16,
+    "fp32": torch.float32,
+}
+
+
+def resolve_torch_dtype(dtype: Union[str, torch.dtype]) -> torch.dtype:
+    """Map ``bf16`` / ``fp16`` / ``fp32`` (or matching ``torch.dtype``) to autocast dtype."""
+    if isinstance(dtype, str):
+        key = dtype.lower()
+        if key not in _DTYPE_ALIASES:
+            raise ValueError(
+                f"Unsupported dtype: {dtype!r}. Choose from: {', '.join(AUTOCAST_DTYPE_CHOICES)}"
+            )
+        return _DTYPE_ALIASES[key]
+    if dtype in _DTYPE_ALIASES.values():
+        return dtype
+    raise ValueError(
+        f"Unsupported dtype: {dtype!r}. Choose from: {', '.join(AUTOCAST_DTYPE_CHOICES)}"
+    )
+
+
+def torch_dtype_to_name(dtype: torch.dtype) -> str:
+    for name, torch_dtype in _DTYPE_ALIASES.items():
+        if dtype == torch_dtype:
+            return name
+    raise ValueError(
+        f"Unsupported dtype: {dtype!r}. Choose from: {', '.join(AUTOCAST_DTYPE_CHOICES)}"
+    )

meshflow/utils/mesh.py

@@ -0,0 +1,626 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the
+# LICENSE file in the root directory of this source tree.
+
+from __future__ import annotations
+
+import os
+from pathlib import Path
+from typing import Any, Dict, List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn.functional as F
+import trimesh
+
+MESH_EXTS = (".stl", ".obj", ".ply", ".glb", ".gltf")
+POINT_CLOUD_EXTS = (".pcd", ".xyz", ".pts", ".npy", ".npz")
+GEOMETRY_EXTS = MESH_EXTS + POINT_CLOUD_EXTS
+DEFAULT_NUM_VERTS = 2048
+
+
+def resolve_num_verts_for_mesh(
+    mesh_path: Union[str, Path],
+    num_verts: int,
+    num_latents: int | None,
+) -> int:
+    """GLB proj_cond value: file verts if < num_latents, else num_verts (not for RoPE)."""
+    path = Path(mesh_path)
+    if path.suffix.lower() != ".glb" or Mesh.is_point_cloud_file(str(path)):
+        return int(num_verts)
+    if num_latents is None:
+        return int(num_verts)
+    mesh = Mesh.load_mesh(str(path), normalize=False, preprocess=True)
+    file_num_verts = int(mesh.verts.shape[0])
+    if file_num_verts < int(num_latents):
+        return file_num_verts
+    return int(num_verts)
+
+
+def _normalize_points(
+    points: torch.Tensor,
+    normalize_by: str = "bsphere",
+    size: float = 2.0,
+) -> torch.Tensor:
+    if normalize_by == "bbox":
+        bbox_min, bbox_max = points.min(0).values, points.max(0).values
+        return (points - (bbox_min + bbox_max) / 2) * (size / (bbox_max - bbox_min).max())
+    if normalize_by == "bsphere":
+        center = (points.min(0).values + points.max(0).values) / 2
+        radius = torch.linalg.norm(points - center, dim=-1).max() * 2
+        return (points - center) * (size / radius)
+    raise NotImplementedError(f"Invalid normalize_by: {normalize_by}")
+
+
+def _resample_points(points: torch.Tensor, num_points: int) -> torch.Tensor:
+    if points.shape[0] == 0:
+        raise ValueError("Point cloud is empty")
+    if points.shape[0] >= num_points:
+        idx = torch.randperm(points.shape[0], device=points.device)[:num_points]
+        return points[idx]
+    idx = torch.randint(0, points.shape[0], (num_points,), device=points.device)
+    return points[idx]
+
+
+def _resample_point_cloud(points: torch.Tensor, num_points: int) -> torch.Tensor:
+    if points.shape[0] == 0:
+        raise ValueError("Point cloud is empty")
+    if points.shape[0] < num_points:
+        raise ValueError(
+            f"Point cloud must have at least {num_points} points, got {int(points.shape[0])}"
+        )
+    idx = torch.randperm(points.shape[0], device=points.device)[:num_points]
+    return points[idx]
+
+
+def _read_point_cloud_file(filename: str) -> torch.Tensor:
+    ext = os.path.splitext(filename)[1].lower()
+    if ext == ".npy":
+        arr = np.load(filename)
+    elif ext == ".npz":
+        data = np.load(filename)
+        for key in ("points", "pts", "xyz", "vertices", "data"):
+            if key in data:
+                arr = data[key]
+                break
+        else:
+            raise ValueError(f"No point array found in {filename}")
+    elif ext in (".xyz", ".pts"):
+        arr = np.loadtxt(filename, dtype=np.float32)
+    elif ext == ".pcd":
+        with open(filename, "r", encoding="utf-8", errors="ignore") as f:
+            lines = [line.strip() for line in f if line.strip() and not line.startswith("#")]
+        start = 0
+        for i, line in enumerate(lines):
+            header = line.upper()
+            if header.startswith("POINTS") or header.startswith("WIDTH"):
+                start = i + 1
+                break
+        arr = np.loadtxt(lines[start:], dtype=np.float32)
+    else:
+        raise NotImplementedError(f"Unsupported point cloud format: {ext}")
+
+    arr = np.asarray(arr, dtype=np.float32)
+    if arr.ndim == 1:
+        if arr.shape[0] % 3 != 0:
+            raise ValueError(f"Invalid point array shape in {filename}")
+        arr = arr.reshape(-1, 3)
+    if arr.shape[-1] < 3:
+        raise ValueError(f"Expected at least 3 coordinates per point in {filename}")
+    return torch.from_numpy(arr[..., :3].copy())
+
+
+def _try_read_point_cloud_ply(filename: str) -> Optional[torch.Tensor]:
+    loaded = trimesh.load(filename, process=False)
+    if isinstance(loaded, trimesh.PointCloud):
+        return torch.from_numpy(np.asarray(loaded.vertices, dtype=np.float32))
+    if isinstance(loaded, trimesh.Trimesh) and len(loaded.faces) == 0:
+        return torch.from_numpy(np.asarray(loaded.vertices, dtype=np.float32))
+    if isinstance(loaded, trimesh.Scene):
+        points = []
+        for geometry in loaded.geometry.values():
+            if isinstance(geometry, trimesh.PointCloud):
+                points.append(geometry.vertices)
+            elif isinstance(geometry, trimesh.Trimesh) and len(geometry.faces) == 0:
+                points.append(geometry.vertices)
+        if points:
+            return torch.from_numpy(np.concatenate(points, axis=0).astype(np.float32))
+    return None
+
+
+def detect_and_triangulate_holes(
+    triangles: List[Tuple[int, int, int]],
+    max_cycle_size: int = 20,
+) -> List[Tuple[int, int, int]]:
+    """Find boundary loops and fan-triangulate holes in a triangle soup."""
+    edge_to_triangles: Dict[Tuple[int, int], List[Tuple[int, int, int]]] = {}
+    for tri in triangles:
+        edges_in_tri = [
+            tuple(sorted([tri[0], tri[1]])),
+            tuple(sorted([tri[1], tri[2]])),
+            tuple(sorted([tri[0], tri[2]])),
+        ]
+        for edge in edges_in_tri:
+            edge_to_triangles.setdefault(edge, []).append(tri)
+
+    boundary_edges = {edge for edge, tris in edge_to_triangles.items() if len(tris) == 1}
+    if not boundary_edges:
+        return []
+
+    boundary_vertices: set[int] = set()
+    boundary_adjacency: Dict[int, List[int]] = {}
+    for v1, v2 in boundary_edges:
+        boundary_vertices.update((v1, v2))
+        boundary_adjacency.setdefault(v1, []).append(v2)
+        boundary_adjacency.setdefault(v2, []).append(v1)
+
+    def find_cycles_dfs(start_vertex: int) -> List[List[int]]:
+        cycles: List[List[int]] = []
+        visited_in_path: set[int] = set()
+        path: List[int] = []
+
+        def dfs(current: int) -> None:
+            visited_in_path.add(current)
+            path.append(current)
+            has_unvisited_neighbor = False
+            for neighbor in boundary_adjacency.get(current, []):
+                if neighbor == start_vertex and len(path) >= 3:
+                    if 4 <= len(path) <= max_cycle_size:
+                        cycles.append(path[:])
+                elif neighbor not in visited_in_path and len(path) < max_cycle_size:
+                    has_unvisited_neighbor = True
+                    dfs(neighbor)
+            if not has_unvisited_neighbor and len(path) >= 3:
+                if 4 <= len(path) + 1 <= max_cycle_size:
+                    if start_vertex not in boundary_adjacency.get(current, []):
+                        missing_edge = tuple(sorted([current, start_vertex]))
+                        if missing_edge in edge_to_triangles:
+                            cycles.append(path[:])
+            path.pop()
+            visited_in_path.remove(current)
+
+        dfs(start_vertex)
+        return cycles
+
+    all_cycles: List[List[int]] = []
+    processed_cycles: set[Tuple[int, ...]] = set()
+    for start_v in boundary_vertices:
+        for cycle in find_cycles_dfs(start_v):
+            min_idx = cycle.index(min(cycle))
+            normalized = tuple(cycle[min_idx:] + cycle[:min_idx])
+            reversed_cycle = tuple(reversed(normalized))
+            if normalized not in processed_cycles and reversed_cycle not in processed_cycles:
+                processed_cycles.add(normalized)
+                all_cycles.append(cycle)
+
+    new_triangles: List[Tuple[int, int, int]] = []
+    for loop in all_cycles:
+        v0 = loop[0]
+        for i in range(1, len(loop) - 1):
+            new_triangles.append(tuple(sorted([v0, loop[i], loop[i + 1]])))
+    return new_triangles
+
+
+def extract_mesh_from_verts_normals_edges(
+    vertices: torch.Tensor,
+    vertices_normal: torch.Tensor,
+    edges: torch.Tensor,
+    fill_holes: bool = False,
+    max_cycle_size: int = 10,
+    device: torch.device | None = None,
+) -> Mesh:
+    """Build a Mesh from predicted vertices, normals, and edge pairs."""
+    if device is None:
+        device = vertices.device
+
+    vertices = vertices.to(device, dtype=torch.float32)
+    vertices_normal = vertices_normal.to(device, dtype=torch.float32)
+    edges = edges.to(device, dtype=torch.long)
+
+    num_vertices = vertices.shape[0]
+    if edges.numel() == 0 or num_vertices < 3:
+        raise ValueError(
+            f"Cannot extract mesh: invalid input (edges={edges.numel()}, vertices={num_vertices})"
+        )
+
+    edges = torch.clamp(edges, 0, num_vertices - 1)
+    edges = edges[edges[:, 0] != edges[:, 1]]
+    if edges.numel() == 0:
+        raise ValueError("Cannot extract mesh: no edges after removing self-loops")
+
+    adjacency_list: List[List[int]] = [[] for _ in range(num_vertices)]
+    edge_set: set[Tuple[int, int]] = set()
+    for e in edges.cpu().numpy():
+        v1, v2 = int(e[0]), int(e[1])
+        if (v1, v2) not in edge_set and (v2, v1) not in edge_set:
+            adjacency_list[v1].append(v2)
+            adjacency_list[v2].append(v1)
+            edge_set.add((v1, v2))
+
+    max_neighbors = max(len(adj) for adj in adjacency_list) if adjacency_list else 0
+    if max_neighbors == 0:
+        raise ValueError("Cannot extract mesh: no vertex neighbors in edge graph")
+
+    adj_tensor = torch.full((num_vertices, max_neighbors), -1, dtype=torch.long, device=device)
+    for i, neighbors in enumerate(adjacency_list):
+        if neighbors:
+            adj_tensor[i, : len(neighbors)] = torch.tensor(neighbors, dtype=torch.long, device=device)
+
+    triangles: List[Tuple[int, int, int]] = []
+    for v1, v2 in edge_set:
+        n1_mask = adj_tensor[v1] >= 0
+        n2_mask = adj_tensor[v2] >= 0
+        if not (n1_mask.any() and n2_mask.any()):
+            continue
+        neighbors_v1 = adj_tensor[v1][n1_mask]
+        neighbors_v2 = adj_tensor[v2][n2_mask]
+        matches = neighbors_v1.unsqueeze(1) == neighbors_v2.unsqueeze(0)
+        if not matches.any():
+            continue
+        v1_indices, _ = torch.where(matches)
+        for v3 in neighbors_v1[v1_indices].unique():
+            if v3 != v1 and v3 != v2:
+                triangles.append(tuple(sorted([v1, v2, int(v3.item())])))
+
+    triangles = list(set(triangles))
+    if not triangles:
+        raise ValueError("Cannot extract mesh: no triangles formed from edges")
+
+    if fill_holes:
+        extra = detect_and_triangulate_holes(triangles, max_cycle_size=max_cycle_size)
+        if extra:
+            triangles = list(set(triangles + extra))
+
+    triangle_tensor = torch.tensor(triangles, dtype=torch.long, device=device)
+    triangle_vertices = vertices[triangle_tensor]
+    edge1 = triangle_vertices[:, 1] - triangle_vertices[:, 0]
+    edge2 = triangle_vertices[:, 2] - triangle_vertices[:, 0]
+    face_normals = torch.linalg.cross(edge1, edge2, dim=1)
+    face_normals_norm = torch.linalg.norm(face_normals, dim=1, keepdim=True)
+    fallback = torch.tensor([0.0, 0.0, 1.0], device=device, dtype=torch.float32)
+    face_normals = torch.where(
+        face_normals_norm > 1e-8,
+        face_normals / face_normals_norm,
+        fallback.expand_as(face_normals),
+    )
+
+    triangle_normals = vertices_normal[triangle_tensor]
+    avg_vertex_normals = triangle_normals.mean(dim=1)
+    avg_normals_norm = torch.linalg.norm(avg_vertex_normals, dim=1, keepdim=True)
+    avg_vertex_normals = torch.where(
+        avg_normals_norm > 1e-8,
+        avg_vertex_normals / avg_normals_norm,
+        fallback.expand_as(avg_vertex_normals),
+    )
+
+    faces = triangle_tensor.clone()
+    flip_mask = (face_normals * avg_vertex_normals).sum(dim=1) < 0
+    faces[flip_mask, 1], faces[flip_mask, 2] = faces[flip_mask, 2], faces[flip_mask, 1]
+    faces = torch.clamp(faces, 0, num_vertices - 1)
+
+    mesh = Mesh(verts=vertices, faces=faces, device=device)
+    mesh._v_nrm = vertices_normal
+    return mesh
+
+
+class Mesh:
+    """Mesh class for loading, augmentation, and topology feature extraction."""
+
+    def __init__(self, verts: torch.Tensor, faces: torch.Tensor, device: torch.device | None = None) -> None:
+        """Args:
+            verts: (N, 3) vertex positions.
+            faces: (F, 3) triangle indices.
+            device: target device; defaults to verts.device.
+        """
+        self.device = device if device is not None else verts.device
+        self.verts = verts.to(self.device, dtype=torch.float32)
+        self.faces = faces.to(self.device, dtype=torch.int64)
+        self._v_nrm: Optional[torch.Tensor] = None
+        self._edges: Optional[torch.Tensor] = None
+
+    @classmethod
+    def is_point_cloud_file(cls, filename: str) -> bool:
+        """True for point-cloud formats, including faceless `.ply` point clouds."""
+        ext = os.path.splitext(filename)[1].lower()
+        if ext in POINT_CLOUD_EXTS:
+            return True
+        if ext == ".ply":
+            return _try_read_point_cloud_ply(filename) is not None
+        return False
+
+    @classmethod
+    def load_mesh(
+        cls,
+        filename: str,
+        normalize: bool = False,
+        normalize_by: str = "bbox",
+        obj_scale: float = 2.0,
+        preprocess: bool = True,
+        device: torch.device | None = None,
+    ) -> Mesh:
+        """Load mesh from disk (.stl/.obj/.ply/.glb/.gltf).
+
+        STL files are rotated -90 degrees about X during load.
+        Scene graphs are flattened into a single mesh.
+
+        Args:
+            normalize: whether to call normalize() after load.
+            normalize_by: "bbox" or "bsphere".
+            obj_scale: target extent after normalization (passed as size).
+            preprocess: whether to merge vertices and drop degenerate faces.
+        """
+        _, ext = os.path.splitext(filename)
+        ext = ext.lower()
+        if ext not in (".stl", ".obj", ".ply", ".glb", ".gltf"):
+            raise NotImplementedError(f"Unsupported mesh format: {ext}")
+        if cls.is_point_cloud_file(filename):
+            raise ValueError(
+                f"{filename} is a point cloud; use load_rope_points() or pass the path to "
+                "MeshFlowPipeline.run() instead of load_mesh()."
+            )
+
+        mesh = trimesh.load(filename, force="mesh", process=False, skip_materials=True)
+        if ext == ".stl":
+            mesh.apply_transform(trimesh.transformations.rotation_matrix(np.radians(-90), [1, 0, 0]))
+        if isinstance(mesh, trimesh.Scene):
+            parts = []
+            for _, node in mesh.graph.to_flattened().items():
+                name = node["geometry"]
+                if name in mesh.geometry and isinstance(mesh.geometry[name], trimesh.Trimesh):
+                    parts.append(mesh.geometry[name].apply_transform(node["transform"]))
+            mesh = trimesh.util.concatenate(parts)
+
+        m = cls(
+            torch.tensor(mesh.vertices, dtype=torch.float32),
+            torch.tensor(mesh.faces, dtype=torch.int64),
+            device=device,
+        )
+        if preprocess:
+            m.preprocess()
+        if normalize:
+            m.normalize(normalize_by=normalize_by, size=obj_scale)
+        return m
+
+    @classmethod
+    def load_rope_points(
+        cls,
+        filename: str,
+        n_verts: int,
+        normalize_by: str = "bsphere",
+        obj_scale: float = 2.0,
+        preprocess: bool = True,
+        device: torch.device | None = None,
+    ) -> torch.Tensor:
+        """Load a mesh or point cloud and return fixed-size points for RoPE conditioning."""
+        device = device or torch.device("cpu")
+        ext = os.path.splitext(filename)[1].lower()
+
+        if ext in POINT_CLOUD_EXTS:
+            points = _read_point_cloud_file(filename).to(device)
+        elif ext == ".ply":
+            points = _try_read_point_cloud_ply(filename)
+            if points is None:
+                mesh = cls.load_mesh(
+                    filename,
+                    normalize=True,
+                    normalize_by=normalize_by,
+                    obj_scale=obj_scale,
+                    preprocess=preprocess,
+                    device=device,
+                )
+                surface_pts, _ = mesh.sample_surface(n_verts)
+                return surface_pts
+            points = points.to(device)
+        else:
+            mesh = cls.load_mesh(
+                filename,
+                normalize=True,
+                normalize_by=normalize_by,
+                obj_scale=obj_scale,
+                preprocess=preprocess,
+                device=device,
+            )
+            surface_pts, _ = mesh.sample_surface(n_verts)
+            return surface_pts
+
+        points = _normalize_points(points, normalize_by=normalize_by, size=obj_scale)
+        return _resample_point_cloud(points, n_verts)
+
+    def to_trimesh(self) -> trimesh.Trimesh:
+        """Export to a host-side mesh object (detached from autograd)."""
+        return trimesh.Trimesh(
+            vertices=self.verts.detach().cpu().numpy(),
+            faces=self.faces.detach().cpu().numpy(),
+            process=False,
+        )
+
+    @property
+    def v_nrm(self) -> torch.Tensor:
+        """Per-vertex normals (area-weighted face normal sum), shape (N, 3)."""
+        if self._v_nrm is None:
+            v0, v1, v2 = self.verts[self.faces[:, 0]], self.verts[self.faces[:, 1]], self.verts[self.faces[:, 2]]
+            fn = torch.linalg.cross(v1 - v0, v2 - v0)
+            v_nrm = torch.zeros_like(self.verts)
+            v_nrm.scatter_add_(0, self.faces[:, 0:1].expand(-1, 3), fn)
+            v_nrm.scatter_add_(0, self.faces[:, 1:2].expand(-1, 3), fn)
+            v_nrm.scatter_add_(0, self.faces[:, 2:3].expand(-1, 3), fn)
+            v_nrm = torch.where(
+                (v_nrm * v_nrm).sum(-1, keepdim=True) > 1e-20,
+                v_nrm,
+                torch.tensor([0.0, 0.0, 1.0], device=self.device),
+            )
+            self._v_nrm = F.normalize(v_nrm, dim=1)
+        return self._v_nrm
+
+    @property
+    def edges(self) -> torch.Tensor:
+        """Unique undirected edges, shape (E, 2)."""
+        if self._edges is None:
+            e = torch.cat([self.faces[:, [0, 1]], self.faces[:, [1, 2]], self.faces[:, [2, 0]]], dim=0)
+            self._edges = torch.unique(torch.sort(e, dim=1)[0], dim=0)
+        return self._edges
+
+    def preprocess(self) -> Mesh:
+        """Clean mesh in-place: merge vertices, remove degenerate/unreferenced geometry."""
+        tm = self.to_trimesh()
+        tm.process()
+        tm.merge_vertices(merge_tex=True, merge_norm=True)
+        tm.update_faces(tm.nondegenerate_faces())
+        tm.update_faces(tm.unique_faces())
+        tm.remove_unreferenced_vertices()
+        self.verts = torch.from_numpy(tm.vertices).to(self.verts).float()
+        self.faces = torch.from_numpy(tm.faces).to(self.faces).long()
+        self._v_nrm = self._edges = None
+        return self
+
+    def normalize(self, normalize_by: str = "bsphere", size: float = 2.0) -> Mesh:
+        """Center and scale vertices in-place to fit a canonical volume.
+
+        Args:
+            normalize_by: "bbox" (axis-aligned box) or "bsphere" (bounding sphere).
+            size: target extent along the longest axis / diameter.
+        """
+        if normalize_by == "bbox":
+            bbox_min, bbox_max = self.verts.min(0).values, self.verts.max(0).values
+            self.verts = (self.verts - (bbox_min + bbox_max) / 2) * (size / (bbox_max - bbox_min).max())
+        elif normalize_by == "bsphere":
+            center = (self.verts.min(0).values + self.verts.max(0).values) / 2
+            self.verts = (self.verts - center) * (
+                size / (torch.linalg.norm(self.verts - center, dim=-1).max() * 2)
+            )
+        else:
+            raise NotImplementedError(f"Invalid normalize_by: {normalize_by}")
+        self._v_nrm = self._edges = None
+        return self
+
+    def rotate(self, yaw: float = 0.0, pitch: float = 0.0, roll: float = 0.0) -> Mesh:
+        """Apply ZYX Euler rotation (degrees) and return a new Mesh."""
+        y, p, r = np.radians(yaw), np.radians(pitch), np.radians(roll)
+        cy, sy, cp, sp, cr, sr = np.cos(y), np.sin(y), np.cos(p), np.sin(p), np.cos(r), np.sin(r)
+        R = torch.tensor(
+            [
+                [cy * cr + sy * sp * sr, -cy * sr + sy * sp * cr, sy * cp],
+                [cp * sr, cp * cr, -sp],
+                [-sy * cr + cy * sp * sr, sy * sr + cy * sp * cr, cy * cp],
+            ],
+            dtype=torch.float32,
+            device=self.device,
+        )
+        return Mesh(self.verts @ R.T, self.faces, device=self.device)
+
+    @torch.no_grad()
+    def sample_surface(self, num_points: int, return_face_idx: bool = False):
+        """Uniformly sample points on the mesh surface.
+
+        Returns:
+            points, face_normals; with return_face_idx=True also per-sample face indices.
+        """
+        tm = self.to_trimesh()
+        points, face_idx = tm.sample(num_points, return_index=True)
+        points = torch.from_numpy(points.astype(np.float32)).to(self.device)
+        normals = torch.from_numpy(tm.face_normals[face_idx].astype(np.float32)).to(self.device)
+        if return_face_idx:
+            return points, normals, torch.from_numpy(face_idx).to(self.device, dtype=torch.long)
+        return points, normals
+
+    @torch.no_grad()
+    def build_topology(
+        self,
+        topology_feats_type: Optional[List[str]] = None,
+        n_verts: int = 8192,
+        max_degree: int = 50,
+    ) -> Dict[str, Any]:
+        """Build padded topology tensors for MeshFlow VAE / inference.
+
+        Supported features: verts_normal, neighbor_points, degree, adjacency_matrix.
+
+        Returns:
+            Dict with padded_verts, verts_mask, verts, faces, num_verts, and requested features.
+        """
+        if topology_feats_type is None:
+            topology_feats_type = ["verts_normal", "neighbor_points", "degree", "adjacency_matrix"]
+
+        m = self.preprocess().normalize(size=2.0)
+        num_verts = int(m.verts.shape[0])
+        if num_verts > n_verts:
+            raise ValueError(f"Mesh has {num_verts} verts > n_verts={n_verts}")
+        if num_verts == 0:
+            raise ValueError("Mesh has no vertices")
+
+        device = m.device
+        edges, v_nrm = m.edges, m.v_nrm
+        n_negative_verts = n_verts - num_verts
+
+        negative_verts = torch.zeros(0, 3, dtype=m.verts.dtype, device=device)
+        negative_normals = torch.zeros(0, 3, dtype=v_nrm.dtype, device=device)
+        nearest_verts_idx = None
+        if n_negative_verts > 0:
+            negative_verts, negative_normals, negative_face_idx = m.sample_surface(
+                n_negative_verts, return_face_idx=True
+            )
+            face_vertices = m.faces[negative_face_idx]
+            dist2 = (face_vertices - negative_verts[:, None, :]).pow(2).sum(dim=-1)
+            nearest_choice = dist2.argmin(dim=1)
+            nearest_verts_idx = face_vertices[
+                torch.arange(n_negative_verts, device=device),
+                nearest_choice,
+            ]
+
+        padded = torch.zeros(n_verts, 3, dtype=m.verts.dtype, device=device)
+        padded[:num_verts] = m.verts
+        if n_negative_verts > 0:
+            padded[num_verts:] = negative_verts.to(dtype=m.verts.dtype, device=device)
+
+        verts_mask = (torch.arange(n_verts, device=device) < num_verts).float().unsqueeze(-1)
+        ret: Dict[str, Any] = {
+            "padded_verts": padded,
+            "verts_mask": verts_mask,
+            "num_verts": num_verts,
+            "num_faces": int(m.faces.shape[0]),
+            "num_edges": int(edges.shape[0]),
+            "verts": m.verts,
+            "faces": m.faces,
+        }
+
+        for feat in topology_feats_type:
+            if feat == "verts_normal":
+                out = torch.zeros(n_verts, 3, dtype=v_nrm.dtype, device=device)
+                out[:num_verts] = v_nrm
+                if n_negative_verts > 0:
+                    out[num_verts:] = negative_normals.to(dtype=v_nrm.dtype, device=device)
+            elif feat == "adjacency_matrix":
+                out = torch.zeros(n_verts, n_verts, dtype=torch.bool, device=device)
+                out[edges[:, 0], edges[:, 1]] = True
+                out[edges[:, 1], edges[:, 0]] = True
+                if n_negative_verts > 0:
+                    neg_slice = slice(num_verts, n_verts)
+                    out[neg_slice, :] = out[nearest_verts_idx, :]
+                    out[:, neg_slice] = out[:, nearest_verts_idx]
+            elif feat == "degree":
+                degrees = torch.bincount(edges.flatten(), minlength=num_verts)
+                if (degrees > max_degree).any():
+                    raise ValueError(f"Vertex degree exceeds max_degree ({max_degree}).")
+                out = torch.zeros(n_verts, dtype=torch.long, device=device)
+                out[:num_verts] = degrees
+                if n_negative_verts > 0:
+                    out[num_verts:] = out[nearest_verts_idx]
+                out = out.unsqueeze(-1)
+            elif feat == "neighbor_points":
+                all_e = torch.cat([edges, edges.flip(1)], dim=0)
+                src, dst = all_e[:, 0], all_e[:, 1]
+                order = torch.argsort(src)
+                src, dst = src[order], dst[order]
+                if torch.bincount(src, minlength=num_verts).max() > max_degree:
+                    raise ValueError(f"Vertex degree exceeds max_degree ({max_degree}).")
+                first = torch.searchsorted(src, src)
+                out = torch.zeros(n_verts, max_degree, 3, dtype=m.verts.dtype, device=device)
+                out[src, torch.arange(len(src), device=device) - first] = m.verts[dst]
+                if n_negative_verts > 0:
+                    out[num_verts:] = out[nearest_verts_idx]
+                out = out.view(n_verts, -1)
+            else:
+                raise ValueError(f"Unknown topology_feats_type: {feat}")
+            ret[feat] = out
+        return ret

nodes.py

@@ -0,0 +1,133 @@
+import sys
+import os
+import torch
+import numpy as np
+import trimesh as tm
+from PIL import Image
+import folder_paths
+
+sys.path.append(os.path.dirname(os.path.abspath(__file__)))
+
+from meshflow.pipelines import MeshFlowPipeline
+from meshflow.utils.mesh import Mesh
+
+_PIPELINE_CACHE = {}
+
+def get_pipeline(model_path, device, dtype, compile_models, num_verts):
+    cache_key = (model_path, device, dtype, compile_models, num_verts)
+    if cache_key not in _PIPELINE_CACHE:
+        _PIPELINE_CACHE.clear()
+        pipeline = MeshFlowPipeline.from_pretrained(
+            model_path=model_path,
+            device=device,
+            dtype=dtype,
+            compile_models=compile_models,
+            num_verts=num_verts,
+        )
+        if pipeline._visual_encoder_cfg is not None:
+            hub_dir = pipeline._visual_encoder_cfg.get("hub_dir", "")
+            if not os.path.exists(hub_dir):
+                home_hub_dir = os.path.join(os.path.expanduser("~"), ".cache", "torch", "hub", "facebookresearch_dinov3_main")
+                pipeline._visual_encoder_cfg["hub_dir"] = home_hub_dir
+            
+            local_weights = os.path.join(folder_paths.models_dir, "facebook", "dinov3-vitl16-pretrain-lvd1689m", "dinov3_vitl16_pretrain_lvd1689m-8aa4cbdd.pth")
+            if os.path.exists(local_weights):
+                pipeline._visual_encoder_cfg["hub_weights"] = local_weights
+            else:
+                hub_weights = pipeline._visual_encoder_cfg.get("hub_weights")
+                if hub_weights and not os.path.exists(hub_weights):
+                    pipeline._visual_encoder_cfg["hub_weights"] = None
+        _PIPELINE_CACHE[cache_key] = pipeline
+    return _PIPELINE_CACHE[cache_key]
+
+class MeshFlowRemesh:
+    @classmethod
+    def INPUT_TYPES(cls):
+        return {
+            "required": {
+                "trimesh": ("TRIMESH", {"tooltip": "The input 3D model to be remeshed."}),
+                "model_name": (["meshflow", "meshflow_w_num_verts_control"], {"default": "meshflow", "tooltip": "Select the MeshFlow model to use. 'meshflow' is the standard model. 'meshflow_w_num_verts_control' allows dynamic control of the generated mesh resolution."}),
+                "steps": ("INT", {"default": 28, "min": 1, "max": 1000, "step": 1, "tooltip": "Number of diffusion sampling steps. Higher values can increase detail but take longer."}),
+                "guidance_scale": ("FLOAT", {"default": 2.5, "min": 0.0, "max": 100.0, "step": 0.1, "tooltip": "Classifier-Free Guidance (CFG) scale for visual conditioning. Only effective when reference_image is connected."}),
+                "seed": ("INT", {"default": 42, "min": 0, "max": 0xffffffffffffffff, "tooltip": "Random seed for sampling latents."}),
+                "num_verts": ([1024, 2048, 3072, 4096, 5120, 6144, 7168, 8192], {"default": 4096, "tooltip": "Target vertex resolution. Roughly controls the generated mesh resolution."}),
+                "device": (["cuda", "cpu"], {"default": "cuda", "tooltip": "Computation device to run the model on (cuda or cpu)."}),
+                "dtype": (["fp16", "bf16", "fp32"], {"default": "fp16", "tooltip": "Precision model dtype (fp16, bf16, or fp32)."}),
+                "compile": ("BOOLEAN", {"default": False, "tooltip": "Whether to use torch.compile on CUDA for faster inference."}),
+                "use_rmbg": ("BOOLEAN", {"default": False, "tooltip": "Enable automatic background removal and foreground cropping for the reference image."}),
+                "fill_holes": ("BOOLEAN", {"default": False, "tooltip": "Automatically closes and triangulates boundary loops in the final mesh topology."}),
+            },
+            "optional": {
+                "reference_image": ("IMAGE", {"tooltip": "Optional reference image for image-conditioned generation."}),
+            }
+        }
+
+    RETURN_TYPES = ("TRIMESH",)
+    RETURN_NAMES = ("trimesh",)
+    FUNCTION = "remesh"
+    CATEGORY = "MeshFlow"
+
+    def remesh(self, trimesh, model_name, steps, guidance_scale, seed, num_verts, device, dtype, compile, use_rmbg, fill_holes, reference_image=None):
+        model_path = os.path.join(folder_paths.models_dir, "facebook", "meshflow", model_name)
+        if not os.path.isdir(model_path):
+            raise FileNotFoundError(f"MeshFlow model path not found at {model_path}. Please download and place the config.yaml and model.pth in that directory.")
+
+        pipeline = get_pipeline(model_path, device, dtype, compile, num_verts)
+        pipeline.use_rmbg = use_rmbg
+
+        if reference_image is not None:
+            hub_dir = pipeline._visual_encoder_cfg.get("hub_dir")
+            if hub_dir and not os.path.exists(hub_dir):
+                os.makedirs(os.path.dirname(hub_dir), exist_ok=True)
+                import subprocess
+                subprocess.run(["git", "clone", "https://github.com/facebookresearch/dinov3.git", hub_dir], check=True)
+
+        if isinstance(trimesh, tm.Scene):
+            parts = []
+            for _, node in trimesh.graph.to_flattened().items():
+                name = node["geometry"]
+                if name in trimesh.geometry and isinstance(trimesh.geometry[name], tm.Trimesh):
+                    parts.append(trimesh.geometry[name].copy().apply_transform(node["transform"]))
+            trimesh = tm.util.concatenate(parts)
+
+        mesh_in = Mesh(
+            verts=torch.tensor(trimesh.vertices, dtype=torch.float32),
+            faces=torch.tensor(trimesh.faces, dtype=torch.int64),
+            device=torch.device(device)
+        )
+        mesh_in.preprocess()
+        mesh_in.normalize(normalize_by="bsphere", size=2.0)
+
+        pil_image = None
+        if reference_image is not None:
+            img_tensor = reference_image[0]
+            if img_tensor.shape[-1] == 4:
+                img_tensor = img_tensor[..., :3]
+            img_np = (img_tensor * 255.0).clamp(0, 255).to(torch.uint8).cpu().numpy()
+            pil_image = Image.fromarray(img_np)
+
+        if pil_image is None:
+            guidance_scale = 1.0
+
+        out_mesh, latents = pipeline.run(
+            mesh=mesh_in,
+            image=pil_image,
+            steps=steps,
+            guidance_scale=guidance_scale,
+            seed=seed,
+            num_verts=num_verts,
+            return_latent=True
+        )
+
+        if fill_holes:
+            out_mesh = pipeline.decode_latent(latents, fill_holes=True)
+
+        return (out_mesh.to_trimesh(),)
+
+NODE_CLASS_MAPPINGS = {
+    "MeshFlowRemesh": MeshFlowRemesh
+}
+
+NODE_DISPLAY_NAME_MAPPINGS = {
+    "MeshFlowRemesh": "MeshFlow Remesh"
+}