Remove PyTorch3D runtime dependency

anigen/models/structured_latent_vae/anigen_decoder.py

@@ -8,7 +8,7 @@ from ...modules import sparse as sp
 from ...representations import MeshExtractResult
 from ...representations.mesh import AniGenSparseFeatures2Mesh, AniGenSklFeatures2Skeleton
 from ..sparse_elastic_mixin import SparseTransformerElasticMixin
-from pytorch3d.ops import knn_points
+from ...utils.pytorch3d_compat import knn_points
 from .anigen_base import AniGenSparseTransformerBase, FreqPositionalEmbedder
 from .skin_models import SKIN_MODEL_DICT
 import torch.nn.functional as F

anigen/models/structured_latent_vae/anigen_encoder.py

@@ -5,7 +5,7 @@ import torch.nn.functional as F
 from ...modules import sparse as sp
 from ..sparse_elastic_mixin import SparseTransformerElasticMixin
 from .anigen_base import AniGenSparseTransformerBase, FreqPositionalEmbedder
-from pytorch3d.ops import knn_points
+from ...utils.pytorch3d_compat import knn_points
 from .skin_models import SkinEncoder
     
 

anigen/representations/mesh/cube2mesh_skeleton.py

@@ -5,7 +5,7 @@ from ...modules.sparse import SparseTensor
 from easydict import EasyDict as edict
 from .utils_cube import *
 from .flexicubes.flexicubes import FlexiCubes
-from pytorch3d.ops import knn_points
+from ...utils.pytorch3d_compat import knn_points
 
 
 class AniGenMeshExtractResult:

anigen/representations/skeleton/grouping.py

@@ -1,6 +1,6 @@
 import torch
 import numpy as np
-from pytorch3d.ops import ball_query, knn_points
+from ...utils.pytorch3d_compat import ball_query, knn_points
 
 
 def disjoint_set_unioin_find(N, pairs):
@@ -178,4 +178,3 @@ GROUPING_STRATEGIES = {
     "threshold": threshold_grouping,
     "mean_shift": mean_shift_grouping,
 }
-

anigen/trainers/flow_matching/anigen_sparse_flow_matching.py

@@ -19,7 +19,7 @@ from .mixins.text_conditioned import TextConditionedMixin
 from .mixins.image_conditioned import ImageConditionedMixin
 from ...representations import MeshExtractResult
 from ...utils.skin_utils import get_transform
-from pytorch3d.ops import knn_points
+from ...utils.pytorch3d_compat import knn_points
 from ...renderers import MeshRenderer
 from ...utils.data_utils import recursive_to_device
 import copy

anigen/trainers/vae/anigen_skin_ae.py

@@ -19,7 +19,7 @@ from ...renderers import OctreeRenderer
 
 from ...modules.sparse import SparseTensor
 from ...utils.loss_utils import l1_loss, smooth_l1_loss, l2_loss, ssim, lpips
-from pytorch3d.ops import knn_points
+from ...utils.pytorch3d_compat import knn_points
 import torch.nn.functional as F
 
 

anigen/trainers/vae/anigen_slat_mesh_vae.py

@@ -19,7 +19,7 @@ from ...renderers import OctreeRenderer
 
 from ...modules.sparse import SparseTensor
 from ...utils.loss_utils import l1_loss, smooth_l1_loss, l2_loss, ssim, lpips
-from pytorch3d.ops import knn_points
+from ...utils.pytorch3d_compat import knn_points
 import torch.nn.functional as F
 
 

anigen/utils/export_utils.py

@@ -472,10 +472,10 @@ def transfer_vertex_colors_nearest(src_vertices, src_colors, dst_vertices):
     dst_v = torch.from_numpy(dst_vertices).to(device=device, dtype=torch.float32)
     src_c = torch.from_numpy(src_colors).to(device=device, dtype=torch.float32)
 
-    # Prefer pytorch3d for KNN if available.
+    # Prefer the local PyTorch3D-compatible helper for KNN.
     idx = None
     try:
-        from pytorch3d.ops import knn_points
+        from .pytorch3d_compat import knn_points
         _, nn_idx, _ = knn_points(dst_v[None], src_v[None], K=1, return_nn=False)
         idx = nn_idx[0, :, 0]
     except Exception:

anigen/utils/pytorch3d_compat.py

@@ -0,0 +1,107 @@
+from collections import namedtuple
+
+import torch
+
+
+KNNResult = namedtuple("KNNResult", ["dists", "idx", "knn"])
+BallQueryResult = namedtuple("BallQueryResult", ["dists", "idx", "knn"])
+
+
+def knn_points(p1, p2, lengths1=None, lengths2=None, K=1, norm=2, return_nn=False, return_sorted=True):
+    if p1.dim() != 3 or p2.dim() != 3:
+        raise ValueError("p1 and p2 must have shape (N, P, D)")
+    if p1.shape[0] != p2.shape[0] or p1.shape[2] != p2.shape[2]:
+        raise ValueError("p1 and p2 must have matching batch and point dimensions")
+
+    batch, p1_count, _ = p1.shape
+    p2_count = p2.shape[1]
+    k = min(K, p2_count)
+    if k <= 0:
+        empty_dists = p1.new_empty((batch, p1_count, 0))
+        empty_idx = torch.empty((batch, p1_count, 0), dtype=torch.long, device=p1.device)
+        empty_nn = p2.new_empty((batch, p1_count, 0, p2.shape[2])) if return_nn else None
+        return KNNResult(empty_dists, empty_idx, empty_nn)
+
+    dists = torch.cdist(p1.float(), p2.float(), p=norm)
+    if norm == 2:
+        dists = dists.square()
+
+    if lengths2 is not None:
+        arange = torch.arange(p2_count, device=p2.device).view(1, 1, p2_count)
+        valid = arange < lengths2.to(device=p2.device).view(batch, 1, 1)
+        dists = dists.masked_fill(~valid, torch.inf)
+
+    dists_k, idx = torch.topk(dists, k=k, dim=-1, largest=False, sorted=return_sorted)
+
+    if K > k:
+        pad = K - k
+        dists_k = torch.cat([dists_k, dists_k.new_full((batch, p1_count, pad), torch.inf)], dim=-1)
+        idx = torch.cat([idx, idx.new_full((batch, p1_count, pad), -1)], dim=-1)
+
+    if lengths1 is not None:
+        arange = torch.arange(p1_count, device=p1.device).view(1, p1_count, 1)
+        valid = arange < lengths1.to(device=p1.device).view(batch, 1, 1)
+        dists_k = dists_k.masked_fill(~valid, torch.inf)
+        idx = idx.masked_fill(~valid, -1)
+
+    knn = None
+    if return_nn:
+        safe_idx = idx.clamp_min(0)
+        gather_idx = safe_idx.unsqueeze(-1).expand(-1, -1, -1, p2.shape[2])
+        points = p2.unsqueeze(1).expand(-1, p1_count, -1, -1)
+        knn = torch.gather(points, 2, gather_idx)
+        knn = knn.masked_fill(idx.unsqueeze(-1) < 0, 0)
+
+    return KNNResult(dists_k.to(dtype=p1.dtype), idx, knn)
+
+
+def ball_query(p1, p2, lengths1=None, lengths2=None, K=1, radius=0.2, return_nn=False):
+    if p1.dim() != 3 or p2.dim() != 3:
+        raise ValueError("p1 and p2 must have shape (N, P, D)")
+    if p1.shape[0] != p2.shape[0] or p1.shape[2] != p2.shape[2]:
+        raise ValueError("p1 and p2 must have matching batch and point dimensions")
+
+    batch, p1_count, _ = p1.shape
+    p2_count = p2.shape[1]
+    k = max(int(K), 0)
+    if k == 0:
+        empty_dists = p1.new_empty((batch, p1_count, 0))
+        empty_idx = torch.empty((batch, p1_count, 0), dtype=torch.long, device=p1.device)
+        empty_nn = p2.new_empty((batch, p1_count, 0, p2.shape[2])) if return_nn else None
+        return BallQueryResult(empty_dists, empty_idx, empty_nn)
+
+    dists = torch.cdist(p1.float(), p2.float(), p=2).square()
+    max_dist = float(radius) * float(radius)
+    valid = dists <= max_dist
+
+    if lengths2 is not None:
+        arange = torch.arange(p2_count, device=p2.device).view(1, 1, p2_count)
+        valid = valid & (arange < lengths2.to(device=p2.device).view(batch, 1, 1))
+
+    masked = dists.masked_fill(~valid, torch.inf)
+    take = min(k, p2_count)
+    dists_k, idx = torch.topk(masked, k=take, dim=-1, largest=False, sorted=True)
+    invalid = torch.isinf(dists_k)
+    idx = idx.masked_fill(invalid, -1)
+    dists_k = dists_k.masked_fill(invalid, 0)
+
+    if k > take:
+        pad = k - take
+        dists_k = torch.cat([dists_k, dists_k.new_zeros((batch, p1_count, pad))], dim=-1)
+        idx = torch.cat([idx, idx.new_full((batch, p1_count, pad), -1)], dim=-1)
+
+    if lengths1 is not None:
+        arange = torch.arange(p1_count, device=p1.device).view(1, p1_count, 1)
+        valid_p1 = arange < lengths1.to(device=p1.device).view(batch, 1, 1)
+        dists_k = dists_k.masked_fill(~valid_p1, 0)
+        idx = idx.masked_fill(~valid_p1, -1)
+
+    knn = None
+    if return_nn:
+        safe_idx = idx.clamp_min(0)
+        gather_idx = safe_idx.unsqueeze(-1).expand(-1, -1, -1, p2.shape[2])
+        points = p2.unsqueeze(1).expand(-1, p1_count, -1, -1)
+        knn = torch.gather(points, 2, gather_idx)
+        knn = knn.masked_fill(idx.unsqueeze(-1) < 0, 0)
+
+    return BallQueryResult(dists_k.to(dtype=p1.dtype), idx, knn)