From eb57ddf5e8b2a7f3a38a2fcf68e0b2f5832cd17b Mon Sep 17 00:00:00 2001 From: Mike Ull <68799232+Rizzlord@users.noreply.github.com> Date: Sun, 7 Jun 2026 03:52:00 +0200 Subject: [PATCH] Initial commit --- .gitattributes | 2 + .gitignore | 160 +++++ __init__.py | 3 + meshflow/__init__.py | 6 + meshflow/models/__init__.py | 6 + meshflow/models/condition_encoder.py | 137 ++++ meshflow/models/meshflow_dit.py | 887 ++++++++++++++++++++++++ meshflow/models/meshflow_vae.py | 786 +++++++++++++++++++++ meshflow/pipelines/__init__.py | 10 + meshflow/pipelines/meshflow_pipeline.py | 463 +++++++++++++ meshflow/pipelines/utils.py | 226 ++++++ meshflow/utils/__init__.py | 6 + meshflow/utils/config.py | 75 ++ meshflow/utils/dtype.py | 44 ++ meshflow/utils/mesh.py | 626 +++++++++++++++++ nodes.py | 133 ++++ 16 files changed, 3570 insertions(+) create mode 100644 .gitattributes create mode 100644 .gitignore create mode 100644 __init__.py create mode 100644 meshflow/__init__.py create mode 100644 meshflow/models/__init__.py create mode 100644 meshflow/models/condition_encoder.py create mode 100755 meshflow/models/meshflow_dit.py create mode 100644 meshflow/models/meshflow_vae.py create mode 100644 meshflow/pipelines/__init__.py create mode 100644 meshflow/pipelines/meshflow_pipeline.py create mode 100644 meshflow/pipelines/utils.py create mode 100644 meshflow/utils/__init__.py create mode 100644 meshflow/utils/config.py create mode 100644 meshflow/utils/dtype.py create mode 100644 meshflow/utils/mesh.py create mode 100644 nodes.py diff --git a/.gitattributes b/.gitattributes new file mode 100644 index 0000000..dfe0770 --- /dev/null +++ b/.gitattributes @@ -0,0 +1,2 @@ +# Auto detect text files and perform LF normalization +* text=auto diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..68bc17f --- /dev/null +++ b/.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/ diff --git a/__init__.py b/__init__.py new file mode 100644 index 0000000..39a8c6b --- /dev/null +++ b/__init__.py @@ -0,0 +1,3 @@ +from .nodes import NODE_CLASS_MAPPINGS, NODE_DISPLAY_NAME_MAPPINGS + +__all__ = ["NODE_CLASS_MAPPINGS", "NODE_DISPLAY_NAME_MAPPINGS"] diff --git a/meshflow/__init__.py b/meshflow/__init__.py new file mode 100644 index 0000000..fe3d6d0 --- /dev/null +++ b/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. + diff --git a/meshflow/models/__init__.py b/meshflow/models/__init__.py new file mode 100644 index 0000000..fe3d6d0 --- /dev/null +++ b/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. + diff --git a/meshflow/models/condition_encoder.py b/meshflow/models/condition_encoder.py new file mode 100644 index 0000000..5237178 --- /dev/null +++ b/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" + ) diff --git a/meshflow/models/meshflow_dit.py b/meshflow/models/meshflow_dit.py new file mode 100755 index 0000000..6d296f1 --- /dev/null +++ b/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), + ) diff --git a/meshflow/models/meshflow_vae.py b/meshflow/models/meshflow_vae.py new file mode 100644 index 0000000..63dd5df --- /dev/null +++ b/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), + ) diff --git a/meshflow/pipelines/__init__.py b/meshflow/pipelines/__init__.py new file mode 100644 index 0000000..4f55b0f --- /dev/null +++ b/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"] diff --git a/meshflow/pipelines/meshflow_pipeline.py b/meshflow/pipelines/meshflow_pipeline.py new file mode 100644 index 0000000..010b95e --- /dev/null +++ b/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 diff --git a/meshflow/pipelines/utils.py b/meshflow/pipelines/utils.py new file mode 100644 index 0000000..79522b5 --- /dev/null +++ b/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 diff --git a/meshflow/utils/__init__.py b/meshflow/utils/__init__.py new file mode 100644 index 0000000..fe3d6d0 --- /dev/null +++ b/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. + diff --git a/meshflow/utils/config.py b/meshflow/utils/config.py new file mode 100644 index 0000000..0e87be8 --- /dev/null +++ b/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}) diff --git a/meshflow/utils/dtype.py b/meshflow/utils/dtype.py new file mode 100644 index 0000000..eea62f8 --- /dev/null +++ b/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)}" + ) diff --git a/meshflow/utils/mesh.py b/meshflow/utils/mesh.py new file mode 100644 index 0000000..9a7923d --- /dev/null +++ b/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 diff --git a/nodes.py b/nodes.py new file mode 100644 index 0000000..192af0d --- /dev/null +++ b/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" +}