Spaces:
Sleeping
Sleeping
| """Point embedding utilities.""" | |
| import torch | |
| import torch.nn as nn | |
| class PointCloudEmbedding: | |
| """Generate positional encodings for multi-part point clouds. | |
| Ref: | |
| Nerf-pytorch: https://github.com/yenchenlin/nerf-pytorch/blob/master/run_nerf_helpers.py | |
| """ | |
| def __init__( | |
| self, | |
| include_input: bool = True, | |
| input_dims: int = 3, | |
| max_freq_log2: int = 9, | |
| num_freqs: int = 10, | |
| log_sampling: bool = True, | |
| periodic_fns: list = None, | |
| ): | |
| self.include_input = include_input | |
| self.input_dims = input_dims | |
| self.max_freq_log2 = max_freq_log2 | |
| self.num_freqs = num_freqs | |
| self.log_sampling = log_sampling | |
| self.periodic_fns = periodic_fns or [torch.sin, torch.cos] | |
| self._create_embedding_fn() | |
| def _create_embedding_fn(self): | |
| """Create the embedding function and compute output dimension.""" | |
| embed_fns = [] | |
| d = self.input_dims | |
| out_dim = 0 | |
| if self.include_input: | |
| embed_fns.append(lambda x: x) | |
| out_dim += d | |
| if self.log_sampling: | |
| freq_bands = 2.0 ** torch.linspace(0.0, self.max_freq_log2, steps=self.num_freqs) | |
| else: | |
| freq_bands = torch.linspace(2.0**0.0, 2.0**self.max_freq_log2, steps=self.num_freqs) | |
| for freq in freq_bands: | |
| for p_fn in self.periodic_fns: | |
| embed_fns.append(lambda x, p_fn=p_fn, freq=freq: p_fn(x * freq)) | |
| out_dim += d | |
| self.embed_fns = embed_fns | |
| self.out_dim = out_dim | |
| # with default setting, the output dim is then 60+3 | |
| def embed(self, inputs: torch.Tensor) -> torch.Tensor: | |
| """Embed input tensor using sinusoidal encoding.""" | |
| return torch.cat([fn(inputs) for fn in self.embed_fns], -1) | |
| class PointCloudEncodingManager(nn.Module): | |
| """Generate PointCloudEmbedding from the input. | |
| It includes PointCloudEmbedding for: | |
| - Coordinate of condition PCs | |
| - Coordinate of noise PCs # do we need to do embedding for this? # Or we directly use the original coordinate? | |
| - Normal vector of condition PCs | |
| - Scale of condition PCs | |
| Args: | |
| in_dim: Input feature dimension. | |
| embed_dim: Output embedding dimension. | |
| multires: Multiresolution level for frequency encoding. | |
| """ | |
| def __init__(self, in_dim: int, embed_dim: int, multires: int = 10, scale_emb_on: bool = False, local_feat_concat_on: bool = False, local_feat_dim: int = 0): | |
| super().__init__() | |
| self.in_dim = in_dim | |
| self.embed_dim = embed_dim | |
| self.multires = multires | |
| self.scale_emb_on = scale_emb_on | |
| self.local_feat_concat_on = local_feat_concat_on | |
| self.local_feat_dim = local_feat_dim | |
| # without encoder, the in_dim is 0 | |
| # Coordinate of condition PCs (63 dim) | |
| self.coord_embedding = PointCloudEmbedding( | |
| input_dims=3, | |
| max_freq_log2=multires - 1, | |
| num_freqs=multires, | |
| ) | |
| # Coordinate of noise PCs (63 dim) | |
| self.noise_embedding = PointCloudEmbedding( | |
| input_dims=3, | |
| max_freq_log2=multires - 1, | |
| num_freqs=multires, | |
| ) | |
| # Scale factor of condition PCs (check if this is needed, or do we need to input the part embedding here), 21 dim | |
| self.scale_embedding = PointCloudEmbedding( | |
| input_dims=1, | |
| max_freq_log2=multires - 1, | |
| num_freqs=multires, | |
| ) | |
| # Do not concatenate the view idx embedding to keep the pertubation equivariance and then can handle a flexible number of input views. | |
| # Embedding projection | |
| embed_input_dim = ( | |
| self.in_dim | |
| + self.coord_embedding.out_dim | |
| + self.noise_embedding.out_dim | |
| ) | |
| if self.scale_emb_on: | |
| embed_input_dim += self.scale_embedding.out_dim # 21 | |
| if self.local_feat_concat_on: | |
| embed_input_dim += self.local_feat_dim # 32 | |
| # print(f"in_dim: {self.in_dim}") # 0 | |
| # print(f"noise_embedding.out_dim: {self.noise_embedding.out_dim}") # 3 or 63 | |
| # print(f"coord_embedding.out_dim: {self.coord_embedding.out_dim}") # 63 | |
| # print(f"scale_embedding.out_dim: {self.scale_embedding.out_dim}") # 21 | |
| # print(f"local_feat_dim: {self.local_feat_dim}") # 32 | |
| # print(f"embed_input_dim: {embed_input_dim}") | |
| self.emb_proj = nn.Linear(embed_input_dim, self.embed_dim) | |
| # linearly map to embed_dim (512) | |
| def forward( | |
| self, | |
| x: torch.Tensor, | |
| cond_coord: torch.Tensor, | |
| local_features: torch.Tensor | None, | |
| latent_features: torch.Tensor | None, | |
| scales: torch.Tensor, | |
| ) -> torch.Tensor: | |
| """Generate PointCloudEmbedding from the input. | |
| Args: | |
| x (TP, 3): Noise point coordinates at timestep t. | |
| cond_coord (TP, 3): Point coordinates of condition point cloud. | |
| local_features (TP, f_dim): Local point features of condition point cloud. | |
| latent_features (TP, f2_dim): Latent point features of condition point cloud. | |
| scales (TP, ): Scale factor for the point cloud. | |
| Returns: | |
| Shape embeddings of shape (TP, dim). | |
| """ | |
| TP, _ = x.shape | |
| # Coordinate of noise PCs | |
| x_pos_emb = self.noise_embedding.embed(x) # (TP, dim) | |
| # Coordinate of condition PCs | |
| coord = cond_coord.view(TP, 3) | |
| c_pos_emb = self.coord_embedding.embed(coord) # (TP, dim) | |
| embed = torch.cat([c_pos_emb, x_pos_emb], dim=-1) | |
| # Concatenate with point features (after ptv3) | |
| if latent_features is not None: | |
| feat = latent_features.view(TP, -1) # (TP, dim) | |
| embed = torch.cat([embed, feat], dim=-1) | |
| # Scale of condition PCs | |
| if self.scale_emb_on: | |
| scales = scales.unsqueeze(-1) | |
| scale_emb = self.scale_embedding.embed(scales) # (TP, dim) | |
| embed = torch.cat([embed, scale_emb], dim=-1) | |
| # Local point features | |
| if self.local_feat_concat_on and local_features is not None: | |
| local_feat = local_features.view(TP, -1) # (TP, dim) | |
| embed = torch.cat([embed, local_feat], dim=-1) | |
| return self.emb_proj(embed) # linear project to the desired dimension for DiT | |