souging/TRELLIS_TextTo3D

1import torch2import torch.nn as nn3from .. import SparseTensor4from .. import DEBUG5from . import SPCONV_ALGO6 7class SparseConv3d(nn.Module):8    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, padding=None, bias=True, indice_key=None):9        super(SparseConv3d, self).__init__()10        if 'spconv' not in globals():11            import spconv.pytorch as spconv12        algo = None13        if SPCONV_ALGO == 'native':14            algo = spconv.ConvAlgo.Native15        elif SPCONV_ALGO == 'implicit_gemm':16            algo = spconv.ConvAlgo.MaskImplicitGemm17        if stride == 1 and (padding is None):18            self.conv = spconv.SubMConv3d(in_channels, out_channels, kernel_size, dilation=dilation, bias=bias, indice_key=indice_key, algo=algo)19        else:20            self.conv = spconv.SparseConv3d(in_channels, out_channels, kernel_size, stride=stride, dilation=dilation, padding=padding, bias=bias, indice_key=indice_key, algo=algo)21        self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride)22        self.padding = padding23 24    def forward(self, x: SparseTensor) -> SparseTensor:25        spatial_changed = any(s != 1 for s in self.stride) or (self.padding is not None)26        new_data = self.conv(x.data)27        new_shape = [x.shape[0], self.conv.out_channels]28        new_layout = None if spatial_changed else x.layout29 30        if spatial_changed and (x.shape[0] != 1):31            # spconv was non-1 stride will break the contiguous of the output tensor, sort by the coords32            fwd = new_data.indices[:, 0].argsort()33            bwd = torch.zeros_like(fwd).scatter_(0, fwd, torch.arange(fwd.shape[0], device=fwd.device))34            sorted_feats = new_data.features[fwd]35            sorted_coords = new_data.indices[fwd]36            unsorted_data = new_data37            new_data = spconv.SparseConvTensor(sorted_feats, sorted_coords, unsorted_data.spatial_shape, unsorted_data.batch_size)  # type: ignore38 39        out = SparseTensor(40            new_data, shape=torch.Size(new_shape), layout=new_layout,41            scale=tuple([s * stride for s, stride in zip(x._scale, self.stride)]),42            spatial_cache=x._spatial_cache,43        )44 45        if spatial_changed and (x.shape[0] != 1):46            out.register_spatial_cache(f'conv_{self.stride}_unsorted_data', unsorted_data)47            out.register_spatial_cache(f'conv_{self.stride}_sort_bwd', bwd)48 49        return out50 51 52class SparseInverseConv3d(nn.Module):53    def __init__(self, in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True, indice_key=None):54        super(SparseInverseConv3d, self).__init__()55        if 'spconv' not in globals():56            import spconv.pytorch as spconv57        self.conv = spconv.SparseInverseConv3d(in_channels, out_channels, kernel_size, bias=bias, indice_key=indice_key)58        self.stride = tuple(stride) if isinstance(stride, (list, tuple)) else (stride, stride, stride)59 60    def forward(self, x: SparseTensor) -> SparseTensor:61        spatial_changed = any(s != 1 for s in self.stride)62        if spatial_changed:63            # recover the original spconv order64            data = x.get_spatial_cache(f'conv_{self.stride}_unsorted_data')65            bwd = x.get_spatial_cache(f'conv_{self.stride}_sort_bwd')66            data = data.replace_feature(x.feats[bwd])67            if DEBUG:68                assert torch.equal(data.indices, x.coords[bwd]), 'Recover the original order failed'69        else:70            data = x.data71 72        new_data = self.conv(data)73        new_shape = [x.shape[0], self.conv.out_channels]74        new_layout = None if spatial_changed else x.layout75        out = SparseTensor(76            new_data, shape=torch.Size(new_shape), layout=new_layout,77            scale=tuple([s // stride for s, stride in zip(x._scale, self.stride)]),78            spatial_cache=x._spatial_cache,79        )80        return out81