|
from typing import *
|
|
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
import numpy as np
|
|
from ...modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32
|
|
from ...modules import sparse as sp
|
|
from .base import SparseTransformerBase
|
|
from ...representations import MeshExtractResult
|
|
from ...representations.mesh import SparseFeatures2Mesh
|
|
|
|
|
|
class SparseSubdivideBlock3d(nn.Module):
|
|
"""
|
|
A 3D subdivide block that can subdivide the sparse tensor.
|
|
|
|
Args:
|
|
channels: channels in the inputs and outputs.
|
|
out_channels: if specified, the number of output channels.
|
|
num_groups: the number of groups for the group norm.
|
|
"""
|
|
def __init__(
|
|
self,
|
|
channels: int,
|
|
resolution: int,
|
|
out_channels: Optional[int] = None,
|
|
num_groups: int = 32
|
|
):
|
|
super().__init__()
|
|
self.channels = channels
|
|
self.resolution = resolution
|
|
self.out_resolution = resolution * 2
|
|
self.out_channels = out_channels or channels
|
|
|
|
self.act_layers = nn.Sequential(
|
|
sp.SparseGroupNorm32(num_groups, channels),
|
|
sp.SparseSiLU()
|
|
)
|
|
|
|
self.sub = sp.SparseSubdivide()
|
|
|
|
self.out_layers = nn.Sequential(
|
|
sp.SparseConv3d(channels, self.out_channels, 3, indice_key=f"res_{self.out_resolution}"),
|
|
sp.SparseGroupNorm32(num_groups, self.out_channels),
|
|
sp.SparseSiLU(),
|
|
zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3, indice_key=f"res_{self.out_resolution}")),
|
|
)
|
|
|
|
if self.out_channels == channels:
|
|
self.skip_connection = nn.Identity()
|
|
else:
|
|
self.skip_connection = sp.SparseConv3d(channels, self.out_channels, 1, indice_key=f"res_{self.out_resolution}")
|
|
|
|
def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
|
|
"""
|
|
Apply the block to a Tensor, conditioned on a timestep embedding.
|
|
|
|
Args:
|
|
x: an [N x C x ...] Tensor of features.
|
|
Returns:
|
|
an [N x C x ...] Tensor of outputs.
|
|
"""
|
|
h = self.act_layers(x)
|
|
h = self.sub(h)
|
|
x = self.sub(x)
|
|
h = self.out_layers(h)
|
|
h = h + self.skip_connection(x)
|
|
return h
|
|
|
|
|
|
class SLatMeshDecoder(SparseTransformerBase):
|
|
def __init__(
|
|
self,
|
|
resolution: int,
|
|
model_channels: int,
|
|
latent_channels: int,
|
|
num_blocks: int,
|
|
num_heads: Optional[int] = None,
|
|
num_head_channels: Optional[int] = 64,
|
|
mlp_ratio: float = 4,
|
|
attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "swin",
|
|
window_size: int = 8,
|
|
pe_mode: Literal["ape", "rope"] = "ape",
|
|
use_fp16: bool = False,
|
|
use_checkpoint: bool = False,
|
|
qk_rms_norm: bool = False,
|
|
representation_config: dict = None,
|
|
):
|
|
super().__init__(
|
|
in_channels=latent_channels,
|
|
model_channels=model_channels,
|
|
num_blocks=num_blocks,
|
|
num_heads=num_heads,
|
|
num_head_channels=num_head_channels,
|
|
mlp_ratio=mlp_ratio,
|
|
attn_mode=attn_mode,
|
|
window_size=window_size,
|
|
pe_mode=pe_mode,
|
|
use_fp16=use_fp16,
|
|
use_checkpoint=use_checkpoint,
|
|
qk_rms_norm=qk_rms_norm,
|
|
)
|
|
self.resolution = resolution
|
|
self.rep_config = representation_config
|
|
self.mesh_extractor = SparseFeatures2Mesh(res=self.resolution*4, use_color=self.rep_config.get('use_color', False))
|
|
self.out_channels = self.mesh_extractor.feats_channels
|
|
self.upsample = nn.ModuleList([
|
|
SparseSubdivideBlock3d(
|
|
channels=model_channels,
|
|
resolution=resolution,
|
|
out_channels=model_channels // 4
|
|
),
|
|
SparseSubdivideBlock3d(
|
|
channels=model_channels // 4,
|
|
resolution=resolution * 2,
|
|
out_channels=model_channels // 8
|
|
)
|
|
])
|
|
self.out_layer = sp.SparseLinear(model_channels // 8, self.out_channels)
|
|
|
|
self.initialize_weights()
|
|
if use_fp16:
|
|
self.convert_to_fp16()
|
|
|
|
def initialize_weights(self) -> None:
|
|
super().initialize_weights()
|
|
|
|
nn.init.constant_(self.out_layer.weight, 0)
|
|
nn.init.constant_(self.out_layer.bias, 0)
|
|
|
|
def convert_to_fp16(self) -> None:
|
|
"""
|
|
Convert the torso of the model to float16.
|
|
"""
|
|
super().convert_to_fp16()
|
|
self.upsample.apply(convert_module_to_f16)
|
|
|
|
def convert_to_fp32(self) -> None:
|
|
"""
|
|
Convert the torso of the model to float32.
|
|
"""
|
|
super().convert_to_fp32()
|
|
self.upsample.apply(convert_module_to_f32)
|
|
|
|
def to_representation(self, x: sp.SparseTensor) -> List[MeshExtractResult]:
|
|
"""
|
|
Convert a batch of network outputs to 3D representations.
|
|
|
|
Args:
|
|
x: The [N x * x C] sparse tensor output by the network.
|
|
|
|
Returns:
|
|
list of representations
|
|
"""
|
|
ret = []
|
|
for i in range(x.shape[0]):
|
|
mesh = self.mesh_extractor(x[i], training=self.training)
|
|
ret.append(mesh)
|
|
return ret
|
|
|
|
def forward(self, x: sp.SparseTensor) -> List[MeshExtractResult]:
|
|
h = super().forward(x)
|
|
for block in self.upsample:
|
|
h = block(h)
|
|
h = h.type(x.dtype)
|
|
h = self.out_layer(h)
|
|
return self.to_representation(h)
|
|
|
