Delete trellis/models

trellis/models/__init__.py

@@ -1,96 +0,0 @@
-import importlib
-
-__attributes = {
-    'SparseStructureEncoder': 'sparse_structure_vae',
-    'SparseStructureDecoder': 'sparse_structure_vae',
-    
-    'SparseStructureFlowModel': 'sparse_structure_flow',
-    
-    'SLatEncoder': 'structured_latent_vae',
-    'SLatGaussianDecoder': 'structured_latent_vae',
-    'SLatRadianceFieldDecoder': 'structured_latent_vae',
-    'SLatMeshDecoder': 'structured_latent_vae',
-    'ElasticSLatEncoder': 'structured_latent_vae',
-    'ElasticSLatGaussianDecoder': 'structured_latent_vae',
-    'ElasticSLatRadianceFieldDecoder': 'structured_latent_vae',
-    'ElasticSLatMeshDecoder': 'structured_latent_vae',
-    
-    'SLatFlowModel': 'structured_latent_flow',
-    'ElasticSLatFlowModel': 'structured_latent_flow',
-}
-
-__submodules = []
-
-__all__ = list(__attributes.keys()) + __submodules
-
-def __getattr__(name):
-    if name not in globals():
-        if name in __attributes:
-            module_name = __attributes[name]
-            module = importlib.import_module(f".{module_name}", __name__)
-            globals()[name] = getattr(module, name)
-        elif name in __submodules:
-            module = importlib.import_module(f".{name}", __name__)
-            globals()[name] = module
-        else:
-            raise AttributeError(f"module {__name__} has no attribute {name}")
-    return globals()[name]
-
-
-def from_pretrained(path: str, **kwargs):
-    """
-    Load a model from a pretrained checkpoint.
-
-    Args:
-        path: The path to the checkpoint. Can be either local path or a Hugging Face model name.
-              NOTE: config file and model file should take the name f'{path}.json' and f'{path}.safetensors' respectively.
-        **kwargs: Additional arguments for the model constructor.
-    """
-    import os
-    import json
-    from safetensors.torch import load_file
-    is_local = os.path.exists(f"{path}.json") and os.path.exists(f"{path}.safetensors")
-
-    if is_local:
-        config_file = f"{path}.json"
-        model_file = f"{path}.safetensors"
-    else:
-        from huggingface_hub import hf_hub_download
-        path_parts = path.split('/')
-        repo_id = f'{path_parts[0]}/{path_parts[1]}'
-        model_name = '/'.join(path_parts[2:])
-        config_file = hf_hub_download(repo_id, f"{model_name}.json")
-        model_file = hf_hub_download(repo_id, f"{model_name}.safetensors")
-
-    with open(config_file, 'r') as f:
-        config = json.load(f)
-    model = __getattr__(config['name'])(**config['args'], **kwargs)
-    model.load_state_dict(load_file(model_file))
-
-    return model
-
-
-# For Pylance
-if __name__ == '__main__':
-    from .sparse_structure_vae import (
-        SparseStructureEncoder, 
-        SparseStructureDecoder,
-    )
-    
-    from .sparse_structure_flow import SparseStructureFlowModel
-    
-    from .structured_latent_vae import (
-        SLatEncoder,
-        SLatGaussianDecoder,
-        SLatRadianceFieldDecoder,
-        SLatMeshDecoder,
-        ElasticSLatEncoder,
-        ElasticSLatGaussianDecoder,
-        ElasticSLatRadianceFieldDecoder,
-        ElasticSLatMeshDecoder,
-    )
-    
-    from .structured_latent_flow import (
-        SLatFlowModel,
-        ElasticSLatFlowModel,
-    )

trellis/models/sparse_elastic_mixin.py

@@ -1,24 +0,0 @@
-from contextlib import contextmanager
-from typing import *
-import math
-from ..modules import sparse as sp
-from ..utils.elastic_utils import ElasticModuleMixin
-
-
-class SparseTransformerElasticMixin(ElasticModuleMixin):
-    def _get_input_size(self, x: sp.SparseTensor, *args, **kwargs):
-        return x.feats.shape[0]
-    
-    @contextmanager
-    def with_mem_ratio(self, mem_ratio=1.0):
-        if mem_ratio == 1.0:
-            yield 1.0
-            return
-        num_blocks = len(self.blocks)
-        num_checkpoint_blocks = min(math.ceil((1 - mem_ratio) * num_blocks) + 1, num_blocks)
-        exact_mem_ratio = 1 - (num_checkpoint_blocks - 1) / num_blocks
-        for i in range(num_blocks):
-            self.blocks[i].use_checkpoint = i < num_checkpoint_blocks
-        yield exact_mem_ratio
-        for i in range(num_blocks):
-            self.blocks[i].use_checkpoint = False

trellis/models/sparse_structure_flow.py

@@ -1,200 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-from ..modules.utils import convert_module_to_f16, convert_module_to_f32
-from ..modules.transformer import AbsolutePositionEmbedder, ModulatedTransformerCrossBlock
-from ..modules.spatial import patchify, unpatchify
-
-
-class TimestepEmbedder(nn.Module):
-    """
-    Embeds scalar timesteps into vector representations.
-    """
-    def __init__(self, hidden_size, frequency_embedding_size=256):
-        super().__init__()
-        self.mlp = nn.Sequential(
-            nn.Linear(frequency_embedding_size, hidden_size, bias=True),
-            nn.SiLU(),
-            nn.Linear(hidden_size, hidden_size, bias=True),
-        )
-        self.frequency_embedding_size = frequency_embedding_size
-
-    @staticmethod
-    def timestep_embedding(t, dim, max_period=10000):
-        """
-        Create sinusoidal timestep embeddings.
-
-        Args:
-            t: a 1-D Tensor of N indices, one per batch element.
-                These may be fractional.
-            dim: the dimension of the output.
-            max_period: controls the minimum frequency of the embeddings.
-
-        Returns:
-            an (N, D) Tensor of positional embeddings.
-        """
-        # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
-        half = dim // 2
-        freqs = torch.exp(
-            -np.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
-        ).to(device=t.device)
-        args = t[:, None].float() * freqs[None]
-        embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
-        if dim % 2:
-            embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
-        return embedding
-
-    def forward(self, t):
-        t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
-        t_emb = self.mlp(t_freq)
-        return t_emb
-
-
-class SparseStructureFlowModel(nn.Module):
-    def __init__(
-        self,
-        resolution: int,
-        in_channels: int,
-        model_channels: int,
-        cond_channels: int,
-        out_channels: int,
-        num_blocks: int,
-        num_heads: Optional[int] = None,
-        num_head_channels: Optional[int] = 64,
-        mlp_ratio: float = 4,
-        patch_size: int = 2,
-        pe_mode: Literal["ape", "rope"] = "ape",
-        use_fp16: bool = False,
-        use_checkpoint: bool = False,
-        share_mod: bool = False,
-        qk_rms_norm: bool = False,
-        qk_rms_norm_cross: bool = False,
-    ):
-        super().__init__()
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.model_channels = model_channels
-        self.cond_channels = cond_channels
-        self.out_channels = out_channels
-        self.num_blocks = num_blocks
-        self.num_heads = num_heads or model_channels // num_head_channels
-        self.mlp_ratio = mlp_ratio
-        self.patch_size = patch_size
-        self.pe_mode = pe_mode
-        self.use_fp16 = use_fp16
-        self.use_checkpoint = use_checkpoint
-        self.share_mod = share_mod
-        self.qk_rms_norm = qk_rms_norm
-        self.qk_rms_norm_cross = qk_rms_norm_cross
-        self.dtype = torch.float16 if use_fp16 else torch.float32
-
-        self.t_embedder = TimestepEmbedder(model_channels)
-        if share_mod:
-            self.adaLN_modulation = nn.Sequential(
-                nn.SiLU(),
-                nn.Linear(model_channels, 6 * model_channels, bias=True)
-            )
-
-        if pe_mode == "ape":
-            pos_embedder = AbsolutePositionEmbedder(model_channels, 3)
-            coords = torch.meshgrid(*[torch.arange(res, device=self.device) for res in [resolution // patch_size] * 3], indexing='ij')
-            coords = torch.stack(coords, dim=-1).reshape(-1, 3)
-            pos_emb = pos_embedder(coords)
-            self.register_buffer("pos_emb", pos_emb)
-
-        self.input_layer = nn.Linear(in_channels * patch_size**3, model_channels)
-            
-        self.blocks = nn.ModuleList([
-            ModulatedTransformerCrossBlock(
-                model_channels,
-                cond_channels,
-                num_heads=self.num_heads,
-                mlp_ratio=self.mlp_ratio,
-                attn_mode='full',
-                use_checkpoint=self.use_checkpoint,
-                use_rope=(pe_mode == "rope"),
-                share_mod=share_mod,
-                qk_rms_norm=self.qk_rms_norm,
-                qk_rms_norm_cross=self.qk_rms_norm_cross,
-            )
-            for _ in range(num_blocks)
-        ])
-
-        self.out_layer = nn.Linear(model_channels, out_channels * patch_size**3)
-
-        self.initialize_weights()
-        if use_fp16:
-            self.convert_to_fp16()
-
-    @property
-    def device(self) -> torch.device:
-        """
-        Return the device of the model.
-        """
-        return next(self.parameters()).device
-
-    def convert_to_fp16(self) -> None:
-        """
-        Convert the torso of the model to float16.
-        """
-        self.blocks.apply(convert_module_to_f16)
-
-    def convert_to_fp32(self) -> None:
-        """
-        Convert the torso of the model to float32.
-        """
-        self.blocks.apply(convert_module_to_f32)
-
-    def initialize_weights(self) -> None:
-        # Initialize transformer layers:
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-        self.apply(_basic_init)
-
-        # Initialize timestep embedding MLP:
-        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
-
-        # Zero-out adaLN modulation layers in DiT blocks:
-        if self.share_mod:
-            nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
-            nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
-        else:
-            for block in self.blocks:
-                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
-                nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
-
-        # Zero-out output layers:
-        nn.init.constant_(self.out_layer.weight, 0)
-        nn.init.constant_(self.out_layer.bias, 0)
-
-    def forward(self, x: torch.Tensor, t: torch.Tensor, cond: torch.Tensor) -> torch.Tensor:
-        assert [*x.shape] == [x.shape[0], self.in_channels, *[self.resolution] * 3], \
-                f"Input shape mismatch, got {x.shape}, expected {[x.shape[0], self.in_channels, *[self.resolution] * 3]}"
-
-        h = patchify(x, self.patch_size)
-        h = h.view(*h.shape[:2], -1).permute(0, 2, 1).contiguous()
-
-        h = self.input_layer(h)
-        h = h + self.pos_emb[None]
-        t_emb = self.t_embedder(t)
-        if self.share_mod:
-            t_emb = self.adaLN_modulation(t_emb)
-        t_emb = t_emb.type(self.dtype)
-        h = h.type(self.dtype)
-        cond = cond.type(self.dtype)
-        for block in self.blocks:
-            h = block(h, t_emb, cond)
-        h = h.type(x.dtype)
-        h = F.layer_norm(h, h.shape[-1:])
-        h = self.out_layer(h)
-
-        h = h.permute(0, 2, 1).view(h.shape[0], h.shape[2], *[self.resolution // self.patch_size] * 3)
-        h = unpatchify(h, self.patch_size).contiguous()
-
-        return h

trellis/models/sparse_structure_vae.py

@@ -1,306 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from ..modules.norm import GroupNorm32, ChannelLayerNorm32
-from ..modules.spatial import pixel_shuffle_3d
-from ..modules.utils import zero_module, convert_module_to_f16, convert_module_to_f32
-
-
-def norm_layer(norm_type: str, *args, **kwargs) -> nn.Module:
-    """
-    Return a normalization layer.
-    """
-    if norm_type == "group":
-        return GroupNorm32(32, *args, **kwargs)
-    elif norm_type == "layer":
-        return ChannelLayerNorm32(*args, **kwargs)
-    else:
-        raise ValueError(f"Invalid norm type {norm_type}")
-
-
-class ResBlock3d(nn.Module):
-    def __init__(
-        self,
-        channels: int,
-        out_channels: Optional[int] = None,
-        norm_type: Literal["group", "layer"] = "layer",
-    ):
-        super().__init__()
-        self.channels = channels
-        self.out_channels = out_channels or channels
-
-        self.norm1 = norm_layer(norm_type, channels)
-        self.norm2 = norm_layer(norm_type, self.out_channels)
-        self.conv1 = nn.Conv3d(channels, self.out_channels, 3, padding=1)
-        self.conv2 = zero_module(nn.Conv3d(self.out_channels, self.out_channels, 3, padding=1))
-        self.skip_connection = nn.Conv3d(channels, self.out_channels, 1) if channels != self.out_channels else nn.Identity()
-    
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        h = self.norm1(x)
-        h = F.silu(h)
-        h = self.conv1(h)
-        h = self.norm2(h)
-        h = F.silu(h)
-        h = self.conv2(h)
-        h = h + self.skip_connection(x)
-        return h
-
-
-class DownsampleBlock3d(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        mode: Literal["conv", "avgpool"] = "conv",
-    ):
-        assert mode in ["conv", "avgpool"], f"Invalid mode {mode}"
-
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-
-        if mode == "conv":
-            self.conv = nn.Conv3d(in_channels, out_channels, 2, stride=2)
-        elif mode == "avgpool":
-            assert in_channels == out_channels, "Pooling mode requires in_channels to be equal to out_channels"
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if hasattr(self, "conv"):
-            return self.conv(x)
-        else:
-            return F.avg_pool3d(x, 2)
-
-
-class UpsampleBlock3d(nn.Module):
-    def __init__(
-        self,
-        in_channels: int,
-        out_channels: int,
-        mode: Literal["conv", "nearest"] = "conv",
-    ):
-        assert mode in ["conv", "nearest"], f"Invalid mode {mode}"
-
-        super().__init__()
-        self.in_channels = in_channels
-        self.out_channels = out_channels
-
-        if mode == "conv":
-            self.conv = nn.Conv3d(in_channels, out_channels*8, 3, padding=1)
-        elif mode == "nearest":
-            assert in_channels == out_channels, "Nearest mode requires in_channels to be equal to out_channels"
-
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        if hasattr(self, "conv"):
-            x = self.conv(x)
-            return pixel_shuffle_3d(x, 2)
-        else:
-            return F.interpolate(x, scale_factor=2, mode="nearest")
-        
-
-class SparseStructureEncoder(nn.Module):
-    """
-    Encoder for Sparse Structure (\mathcal{E}_S in the paper Sec. 3.3).
-    
-    Args:
-        in_channels (int): Channels of the input.
-        latent_channels (int): Channels of the latent representation.
-        num_res_blocks (int): Number of residual blocks at each resolution.
-        channels (List[int]): Channels of the encoder blocks.
-        num_res_blocks_middle (int): Number of residual blocks in the middle.
-        norm_type (Literal["group", "layer"]): Type of normalization layer.
-        use_fp16 (bool): Whether to use FP16.
-    """
-    def __init__(
-        self,
-        in_channels: int,
-        latent_channels: int,
-        num_res_blocks: int,
-        channels: List[int],
-        num_res_blocks_middle: int = 2,
-        norm_type: Literal["group", "layer"] = "layer",
-        use_fp16: bool = False,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.latent_channels = latent_channels
-        self.num_res_blocks = num_res_blocks
-        self.channels = channels
-        self.num_res_blocks_middle = num_res_blocks_middle
-        self.norm_type = norm_type
-        self.use_fp16 = use_fp16
-        self.dtype = torch.float16 if use_fp16 else torch.float32
-
-        self.input_layer = nn.Conv3d(in_channels, channels[0], 3, padding=1)
-
-        self.blocks = nn.ModuleList([])
-        for i, ch in enumerate(channels):
-            self.blocks.extend([
-                ResBlock3d(ch, ch)
-                for _ in range(num_res_blocks)
-            ])
-            if i < len(channels) - 1:
-                self.blocks.append(
-                    DownsampleBlock3d(ch, channels[i+1])
-                )
-        
-        self.middle_block = nn.Sequential(*[
-            ResBlock3d(channels[-1], channels[-1])
-            for _ in range(num_res_blocks_middle)
-        ])
-
-        self.out_layer = nn.Sequential(
-            norm_layer(norm_type, channels[-1]),
-            nn.SiLU(),
-            nn.Conv3d(channels[-1], latent_channels*2, 3, padding=1)
-        )
-
-        if use_fp16:
-            self.convert_to_fp16()
-
-    @property
-    def device(self) -> torch.device:
-        """
-        Return the device of the model.
-        """
-        return next(self.parameters()).device
-
-    def convert_to_fp16(self) -> None:
-        """
-        Convert the torso of the model to float16.
-        """
-        self.use_fp16 = True
-        self.dtype = torch.float16
-        self.blocks.apply(convert_module_to_f16)
-        self.middle_block.apply(convert_module_to_f16)
-
-    def convert_to_fp32(self) -> None:
-        """
-        Convert the torso of the model to float32.
-        """
-        self.use_fp16 = False
-        self.dtype = torch.float32
-        self.blocks.apply(convert_module_to_f32)
-        self.middle_block.apply(convert_module_to_f32)
-
-    def forward(self, x: torch.Tensor, sample_posterior: bool = False, return_raw: bool = False) -> torch.Tensor:
-        h = self.input_layer(x)
-        h = h.type(self.dtype)
-
-        for block in self.blocks:
-            h = block(h)
-        h = self.middle_block(h)
-
-        h = h.type(x.dtype)
-        h = self.out_layer(h)
-
-        mean, logvar = h.chunk(2, dim=1)
-
-        if sample_posterior:
-            std = torch.exp(0.5 * logvar)
-            z = mean + std * torch.randn_like(std)
-        else:
-            z = mean
-            
-        if return_raw:
-            return z, mean, logvar
-        return z
-        
-
-class SparseStructureDecoder(nn.Module):
-    """
-    Decoder for Sparse Structure (\mathcal{D}_S in the paper Sec. 3.3).
-    
-    Args:
-        out_channels (int): Channels of the output.
-        latent_channels (int): Channels of the latent representation.
-        num_res_blocks (int): Number of residual blocks at each resolution.
-        channels (List[int]): Channels of the decoder blocks.
-        num_res_blocks_middle (int): Number of residual blocks in the middle.
-        norm_type (Literal["group", "layer"]): Type of normalization layer.
-        use_fp16 (bool): Whether to use FP16.
-    """ 
-    def __init__(
-        self,
-        out_channels: int,
-        latent_channels: int,
-        num_res_blocks: int,
-        channels: List[int],
-        num_res_blocks_middle: int = 2,
-        norm_type: Literal["group", "layer"] = "layer",
-        use_fp16: bool = False,
-    ):
-        super().__init__()
-        self.out_channels = out_channels
-        self.latent_channels = latent_channels
-        self.num_res_blocks = num_res_blocks
-        self.channels = channels
-        self.num_res_blocks_middle = num_res_blocks_middle
-        self.norm_type = norm_type
-        self.use_fp16 = use_fp16
-        self.dtype = torch.float16 if use_fp16 else torch.float32
-
-        self.input_layer = nn.Conv3d(latent_channels, channels[0], 3, padding=1)
-
-        self.middle_block = nn.Sequential(*[
-            ResBlock3d(channels[0], channels[0])
-            for _ in range(num_res_blocks_middle)
-        ])
-
-        self.blocks = nn.ModuleList([])
-        for i, ch in enumerate(channels):
-            self.blocks.extend([
-                ResBlock3d(ch, ch)
-                for _ in range(num_res_blocks)
-            ])
-            if i < len(channels) - 1:
-                self.blocks.append(
-                    UpsampleBlock3d(ch, channels[i+1])
-                )
-
-        self.out_layer = nn.Sequential(
-            norm_layer(norm_type, channels[-1]),
-            nn.SiLU(),
-            nn.Conv3d(channels[-1], out_channels, 3, padding=1)
-        )
-
-        if use_fp16:
-            self.convert_to_fp16()
-
-    @property
-    def device(self) -> torch.device:
-        """
-        Return the device of the model.
-        """
-        return next(self.parameters()).device
-    
-    def convert_to_fp16(self) -> None:
-        """
-        Convert the torso of the model to float16.
-        """
-        self.use_fp16 = True
-        self.dtype = torch.float16
-        self.blocks.apply(convert_module_to_f16)
-        self.middle_block.apply(convert_module_to_f16)
-
-    def convert_to_fp32(self) -> None:
-        """
-        Convert the torso of the model to float32.
-        """
-        self.use_fp16 = False
-        self.dtype = torch.float32
-        self.blocks.apply(convert_module_to_f32)
-        self.middle_block.apply(convert_module_to_f32)
-    
-    def forward(self, x: torch.Tensor) -> torch.Tensor:
-        h = self.input_layer(x)
-        
-        h = h.type(self.dtype)
-                
-        h = self.middle_block(h)
-        for block in self.blocks:
-            h = block(h)
-
-        h = h.type(x.dtype)
-        h = self.out_layer(h)
-        return h

trellis/models/structured_latent_flow.py

@@ -1,276 +0,0 @@
-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.transformer import AbsolutePositionEmbedder
-from ..modules.norm import LayerNorm32
-from ..modules import sparse as sp
-from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock
-from .sparse_structure_flow import TimestepEmbedder
-from .sparse_elastic_mixin import SparseTransformerElasticMixin
-
-
-class SparseResBlock3d(nn.Module):
-    def __init__(
-        self,
-        channels: int,
-        emb_channels: int,
-        out_channels: Optional[int] = None,
-        downsample: bool = False,
-        upsample: bool = False,
-    ):
-        super().__init__()
-        self.channels = channels
-        self.emb_channels = emb_channels
-        self.out_channels = out_channels or channels
-        self.downsample = downsample
-        self.upsample = upsample
-        
-        assert not (downsample and upsample), "Cannot downsample and upsample at the same time"
-
-        self.norm1 = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
-        self.norm2 = LayerNorm32(self.out_channels, elementwise_affine=False, eps=1e-6)
-        self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3)
-        self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
-        self.emb_layers = nn.Sequential(
-            nn.SiLU(),
-            nn.Linear(emb_channels, 2 * self.out_channels, bias=True),
-        )
-        self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
-        self.updown = None
-        if self.downsample:
-            self.updown = sp.SparseDownsample(2)
-        elif self.upsample:
-            self.updown = sp.SparseUpsample(2)
-
-    def _updown(self, x: sp.SparseTensor) -> sp.SparseTensor:
-        if self.updown is not None:
-            x = self.updown(x)
-        return x
-
-    def forward(self, x: sp.SparseTensor, emb: torch.Tensor) -> sp.SparseTensor:
-        emb_out = self.emb_layers(emb).type(x.dtype)
-        scale, shift = torch.chunk(emb_out, 2, dim=1)
-
-        x = self._updown(x)
-        h = x.replace(self.norm1(x.feats))
-        h = h.replace(F.silu(h.feats))
-        h = self.conv1(h)
-        h = h.replace(self.norm2(h.feats)) * (1 + scale) + shift
-        h = h.replace(F.silu(h.feats))
-        h = self.conv2(h)
-        h = h + self.skip_connection(x)
-
-        return h
-    
-
-class SLatFlowModel(nn.Module):
-    def __init__(
-        self,
-        resolution: int,
-        in_channels: int,
-        model_channels: int,
-        cond_channels: int,
-        out_channels: int,
-        num_blocks: int,
-        num_heads: Optional[int] = None,
-        num_head_channels: Optional[int] = 64,
-        mlp_ratio: float = 4,
-        patch_size: int = 2,
-        num_io_res_blocks: int = 2,
-        io_block_channels: List[int] = None,
-        pe_mode: Literal["ape", "rope"] = "ape",
-        use_fp16: bool = False,
-        use_checkpoint: bool = False,
-        use_skip_connection: bool = True,
-        share_mod: bool = False,
-        qk_rms_norm: bool = False,
-        qk_rms_norm_cross: bool = False,
-    ):
-        super().__init__()
-        self.resolution = resolution
-        self.in_channels = in_channels
-        self.model_channels = model_channels
-        self.cond_channels = cond_channels
-        self.out_channels = out_channels
-        self.num_blocks = num_blocks
-        self.num_heads = num_heads or model_channels // num_head_channels
-        self.mlp_ratio = mlp_ratio
-        self.patch_size = patch_size
-        self.num_io_res_blocks = num_io_res_blocks
-        self.io_block_channels = io_block_channels
-        self.pe_mode = pe_mode
-        self.use_fp16 = use_fp16
-        self.use_checkpoint = use_checkpoint
-        self.use_skip_connection = use_skip_connection
-        self.share_mod = share_mod
-        self.qk_rms_norm = qk_rms_norm
-        self.qk_rms_norm_cross = qk_rms_norm_cross
-        self.dtype = torch.float16 if use_fp16 else torch.float32
-
-        if self.io_block_channels is not None:
-            assert int(np.log2(patch_size)) == np.log2(patch_size), "Patch size must be a power of 2"
-            assert np.log2(patch_size) == len(io_block_channels), "Number of IO ResBlocks must match the number of stages"
-
-        self.t_embedder = TimestepEmbedder(model_channels)
-        if share_mod:
-            self.adaLN_modulation = nn.Sequential(
-                nn.SiLU(),
-                nn.Linear(model_channels, 6 * model_channels, bias=True)
-            )
-
-        if pe_mode == "ape":
-            self.pos_embedder = AbsolutePositionEmbedder(model_channels)
-
-        self.input_layer = sp.SparseLinear(in_channels, model_channels if io_block_channels is None else io_block_channels[0])
-        
-        self.input_blocks = nn.ModuleList([])
-        if io_block_channels is not None:
-            for chs, next_chs in zip(io_block_channels, io_block_channels[1:] + [model_channels]):
-                self.input_blocks.extend([
-                    SparseResBlock3d(
-                        chs,
-                        model_channels,
-                        out_channels=chs,
-                    )
-                    for _ in range(num_io_res_blocks-1)
-                ])
-                self.input_blocks.append(
-                    SparseResBlock3d(
-                        chs,
-                        model_channels,
-                        out_channels=next_chs,
-                        downsample=True,
-                    )
-                )
-            
-        self.blocks = nn.ModuleList([
-            ModulatedSparseTransformerCrossBlock(
-                model_channels,
-                cond_channels,
-                num_heads=self.num_heads,
-                mlp_ratio=self.mlp_ratio,
-                attn_mode='full',
-                use_checkpoint=self.use_checkpoint,
-                use_rope=(pe_mode == "rope"),
-                share_mod=self.share_mod,
-                qk_rms_norm=self.qk_rms_norm,
-                qk_rms_norm_cross=self.qk_rms_norm_cross,
-            )
-            for _ in range(num_blocks)
-        ])
-
-        self.out_blocks = nn.ModuleList([])
-        if io_block_channels is not None:
-            for chs, prev_chs in zip(reversed(io_block_channels), [model_channels] + list(reversed(io_block_channels[1:]))):
-                self.out_blocks.append(
-                    SparseResBlock3d(
-                        prev_chs * 2 if self.use_skip_connection else prev_chs,
-                        model_channels,
-                        out_channels=chs,
-                        upsample=True,
-                    )
-                )
-                self.out_blocks.extend([
-                    SparseResBlock3d(
-                        chs * 2 if self.use_skip_connection else chs,
-                        model_channels,
-                        out_channels=chs,
-                    )
-                    for _ in range(num_io_res_blocks-1)
-                ])
-            
-        self.out_layer = sp.SparseLinear(model_channels if io_block_channels is None else io_block_channels[0], out_channels)
-
-        self.initialize_weights()
-        if use_fp16:
-            self.convert_to_fp16()
-
-    @property
-    def device(self) -> torch.device:
-        """
-        Return the device of the model.
-        """
-        return next(self.parameters()).device
-
-    def convert_to_fp16(self) -> None:
-        """
-        Convert the torso of the model to float16.
-        """
-        self.input_blocks.apply(convert_module_to_f16)
-        self.blocks.apply(convert_module_to_f16)
-        self.out_blocks.apply(convert_module_to_f16)
-
-    def convert_to_fp32(self) -> None:
-        """
-        Convert the torso of the model to float32.
-        """
-        self.input_blocks.apply(convert_module_to_f32)
-        self.blocks.apply(convert_module_to_f32)
-        self.out_blocks.apply(convert_module_to_f32)
-
-    def initialize_weights(self) -> None:
-        # Initialize transformer layers:
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-        self.apply(_basic_init)
-
-        # Initialize timestep embedding MLP:
-        nn.init.normal_(self.t_embedder.mlp[0].weight, std=0.02)
-        nn.init.normal_(self.t_embedder.mlp[2].weight, std=0.02)
-
-        # Zero-out adaLN modulation layers in DiT blocks:
-        if self.share_mod:
-            nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
-            nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
-        else:
-            for block in self.blocks:
-                nn.init.constant_(block.adaLN_modulation[-1].weight, 0)
-                nn.init.constant_(block.adaLN_modulation[-1].bias, 0)
-
-        # Zero-out output layers:
-        nn.init.constant_(self.out_layer.weight, 0)
-        nn.init.constant_(self.out_layer.bias, 0)
-
-    def forward(self, x: sp.SparseTensor, t: torch.Tensor, cond: torch.Tensor) -> sp.SparseTensor:
-        h = self.input_layer(x).type(self.dtype)
-        t_emb = self.t_embedder(t)
-        if self.share_mod:
-            t_emb = self.adaLN_modulation(t_emb)
-        t_emb = t_emb.type(self.dtype)
-        cond = cond.type(self.dtype)
-
-        skips = []
-        # pack with input blocks
-        for block in self.input_blocks:
-            h = block(h, t_emb)
-            skips.append(h.feats)
-        
-        if self.pe_mode == "ape":
-            h = h + self.pos_embedder(h.coords[:, 1:]).type(self.dtype)
-        for block in self.blocks:
-            h = block(h, t_emb, cond)
-
-        # unpack with output blocks
-        for block, skip in zip(self.out_blocks, reversed(skips)):
-            if self.use_skip_connection:
-                h = block(h.replace(torch.cat([h.feats, skip], dim=1)), t_emb)
-            else:
-                h = block(h, t_emb)
-
-        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
-        h = self.out_layer(h.type(x.dtype))
-        return h
-    
-
-class ElasticSLatFlowModel(SparseTransformerElasticMixin, SLatFlowModel):
-    """
-    SLat Flow Model with elastic memory management.
-    Used for training with low VRAM.
-    """
-    pass

trellis/models/structured_latent_vae/__init__.py

@@ -1,4 +0,0 @@
-from .encoder import SLatEncoder
-from .decoder_gs import SLatGaussianDecoder
-from .decoder_rf import SLatRadianceFieldDecoder
-from .decoder_mesh import SLatMeshDecoder

trellis/models/structured_latent_vae/base.py

@@ -1,117 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-from ...modules.utils import convert_module_to_f16, convert_module_to_f32
-from ...modules import sparse as sp
-from ...modules.transformer import AbsolutePositionEmbedder
-from ...modules.sparse.transformer import SparseTransformerBlock
-
-
-def block_attn_config(self):
-    """
-    Return the attention configuration of the model.
-    """
-    for i in range(self.num_blocks):
-        if self.attn_mode == "shift_window":
-            yield "serialized", self.window_size, 0, (16 * (i % 2),) * 3, sp.SerializeMode.Z_ORDER
-        elif self.attn_mode == "shift_sequence":
-            yield "serialized", self.window_size, self.window_size // 2 * (i % 2), (0, 0, 0), sp.SerializeMode.Z_ORDER
-        elif self.attn_mode == "shift_order":
-            yield "serialized", self.window_size, 0, (0, 0, 0), sp.SerializeModes[i % 4]
-        elif self.attn_mode == "full":
-            yield "full", None, None, None, None
-        elif self.attn_mode == "swin":
-            yield "windowed", self.window_size, None, self.window_size // 2 * (i % 2), None
-
-
-class SparseTransformerBase(nn.Module):
-    """
-    Sparse Transformer without output layers.
-    Serve as the base class for encoder and decoder.
-    """
-    def __init__(
-        self,
-        in_channels: int,
-        model_channels: int,
-        num_blocks: int,
-        num_heads: Optional[int] = None,
-        num_head_channels: Optional[int] = 64,
-        mlp_ratio: float = 4.0,
-        attn_mode: Literal["full", "shift_window", "shift_sequence", "shift_order", "swin"] = "full",
-        window_size: Optional[int] = None,
-        pe_mode: Literal["ape", "rope"] = "ape",
-        use_fp16: bool = False,
-        use_checkpoint: bool = False,
-        qk_rms_norm: bool = False,
-    ):
-        super().__init__()
-        self.in_channels = in_channels
-        self.model_channels = model_channels
-        self.num_blocks = num_blocks
-        self.window_size = window_size
-        self.num_heads = num_heads or model_channels // num_head_channels
-        self.mlp_ratio = mlp_ratio
-        self.attn_mode = attn_mode
-        self.pe_mode = pe_mode
-        self.use_fp16 = use_fp16
-        self.use_checkpoint = use_checkpoint
-        self.qk_rms_norm = qk_rms_norm
-        self.dtype = torch.float16 if use_fp16 else torch.float32
-
-        if pe_mode == "ape":
-            self.pos_embedder = AbsolutePositionEmbedder(model_channels)
-
-        self.input_layer = sp.SparseLinear(in_channels, model_channels)
-        self.blocks = nn.ModuleList([
-            SparseTransformerBlock(
-                model_channels,
-                num_heads=self.num_heads,
-                mlp_ratio=self.mlp_ratio,
-                attn_mode=attn_mode,
-                window_size=window_size,
-                shift_sequence=shift_sequence,
-                shift_window=shift_window,
-                serialize_mode=serialize_mode,
-                use_checkpoint=self.use_checkpoint,
-                use_rope=(pe_mode == "rope"),
-                qk_rms_norm=self.qk_rms_norm,
-            )
-            for attn_mode, window_size, shift_sequence, shift_window, serialize_mode in block_attn_config(self)
-        ])
-
-    @property
-    def device(self) -> torch.device:
-        """
-        Return the device of the model.
-        """
-        return next(self.parameters()).device
-
-    def convert_to_fp16(self) -> None:
-        """
-        Convert the torso of the model to float16.
-        """
-        self.blocks.apply(convert_module_to_f16)
-
-    def convert_to_fp32(self) -> None:
-        """
-        Convert the torso of the model to float32.
-        """
-        self.blocks.apply(convert_module_to_f32)
-
-    def initialize_weights(self) -> None:
-        # Initialize transformer layers:
-        def _basic_init(module):
-            if isinstance(module, nn.Linear):
-                torch.nn.init.xavier_uniform_(module.weight)
-                if module.bias is not None:
-                    nn.init.constant_(module.bias, 0)
-        self.apply(_basic_init)
-
-    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
-        h = self.input_layer(x)
-        if self.pe_mode == "ape":
-            h = h + self.pos_embedder(x.coords[:, 1:])
-        h = h.type(self.dtype)
-        for block in self.blocks:
-            h = block(h)
-        return h

trellis/models/structured_latent_vae/decoder_gs.py

@@ -1,122 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from ...modules import sparse as sp
-from ...utils.random_utils import hammersley_sequence
-from .base import SparseTransformerBase
-from ...representations import Gaussian
-
-
-class SLatGaussianDecoder(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._calc_layout()
-        self.out_layer = sp.SparseLinear(model_channels, self.out_channels)
-        self._build_perturbation()
-
-        self.initialize_weights()
-        if use_fp16:
-            self.convert_to_fp16()
-
-    def initialize_weights(self) -> None:
-        super().initialize_weights()
-        # Zero-out output layers:
-        nn.init.constant_(self.out_layer.weight, 0)
-        nn.init.constant_(self.out_layer.bias, 0)
-
-    def _build_perturbation(self) -> None:
-        perturbation = [hammersley_sequence(3, i, self.rep_config['num_gaussians']) for i in range(self.rep_config['num_gaussians'])]
-        perturbation = torch.tensor(perturbation).float() * 2 - 1
-        perturbation = perturbation / self.rep_config['voxel_size']
-        perturbation = torch.atanh(perturbation).to(self.device)
-        self.register_buffer('offset_perturbation', perturbation)
-
-    def _calc_layout(self) -> None:
-        self.layout = {
-            '_xyz' : {'shape': (self.rep_config['num_gaussians'], 3), 'size': self.rep_config['num_gaussians'] * 3},
-            '_features_dc' : {'shape': (self.rep_config['num_gaussians'], 1, 3), 'size': self.rep_config['num_gaussians'] * 3},
-            '_scaling' : {'shape': (self.rep_config['num_gaussians'], 3), 'size': self.rep_config['num_gaussians'] * 3},
-            '_rotation' : {'shape': (self.rep_config['num_gaussians'], 4), 'size': self.rep_config['num_gaussians'] * 4},
-            '_opacity' : {'shape': (self.rep_config['num_gaussians'], 1), 'size': self.rep_config['num_gaussians']},
-        }
-        start = 0
-        for k, v in self.layout.items():
-            v['range'] = (start, start + v['size'])
-            start += v['size']
-        self.out_channels = start
-    
-    def to_representation(self, x: sp.SparseTensor) -> List[Gaussian]:
-        """
-        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]):
-            representation = Gaussian(
-                sh_degree=0,
-                aabb=[-0.5, -0.5, -0.5, 1.0, 1.0, 1.0],
-                mininum_kernel_size = self.rep_config['3d_filter_kernel_size'],
-                scaling_bias = self.rep_config['scaling_bias'],
-                opacity_bias = self.rep_config['opacity_bias'],
-                scaling_activation = self.rep_config['scaling_activation']
-            )
-            xyz = (x.coords[x.layout[i]][:, 1:].float() + 0.5) / self.resolution
-            for k, v in self.layout.items():
-                if k == '_xyz':
-                    offset = x.feats[x.layout[i]][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape'])
-                    offset = offset * self.rep_config['lr'][k]
-                    if self.rep_config['perturb_offset']:
-                        offset = offset + self.offset_perturbation
-                    offset = torch.tanh(offset) / self.resolution * 0.5 * self.rep_config['voxel_size']
-                    _xyz = xyz.unsqueeze(1) + offset
-                    setattr(representation, k, _xyz.flatten(0, 1))
-                else:
-                    feats = x.feats[x.layout[i]][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape']).flatten(0, 1)
-                    feats = feats * self.rep_config['lr'][k]
-                    setattr(representation, k, feats)
-            ret.append(representation)
-        return ret
-
-    def forward(self, x: sp.SparseTensor) -> List[Gaussian]:
-        h = super().forward(x)
-        h = h.type(x.dtype)
-        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
-        h = self.out_layer(h)
-        return self.to_representation(h)

trellis/models/structured_latent_vae/decoder_mesh.py

@@ -1,167 +0,0 @@
-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()
-        # Zero-out output layers:
-        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)

trellis/models/structured_latent_vae/decoder_rf.py

@@ -1,104 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-import numpy as np
-from ...modules import sparse as sp
-from .base import SparseTransformerBase
-from ...representations import Strivec
-
-
-class SLatRadianceFieldDecoder(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._calc_layout()
-        self.out_layer = sp.SparseLinear(model_channels, self.out_channels)
-
-        self.initialize_weights()
-        if use_fp16:
-            self.convert_to_fp16()
-
-    def initialize_weights(self) -> None:
-        super().initialize_weights()
-        # Zero-out output layers:
-        nn.init.constant_(self.out_layer.weight, 0)
-        nn.init.constant_(self.out_layer.bias, 0)
-
-    def _calc_layout(self) -> None:
-        self.layout = {
-            'trivec': {'shape': (self.rep_config['rank'], 3, self.rep_config['dim']), 'size': self.rep_config['rank'] * 3 * self.rep_config['dim']},
-            'density': {'shape': (self.rep_config['rank'],), 'size': self.rep_config['rank']},
-            'features_dc': {'shape': (self.rep_config['rank'], 1, 3), 'size': self.rep_config['rank'] * 3},
-        }
-        start = 0
-        for k, v in self.layout.items():
-            v['range'] = (start, start + v['size'])
-            start += v['size']
-        self.out_channels = start    
-    
-    def to_representation(self, x: sp.SparseTensor) -> List[Strivec]:
-        """
-        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]):
-            representation = Strivec(
-                sh_degree=0,
-                resolution=self.resolution,
-                aabb=[-0.5, -0.5, -0.5, 1, 1, 1],
-                rank=self.rep_config['rank'],
-                dim=self.rep_config['dim'],
-                device='cuda',
-            )
-            representation.density_shift = 0.0
-            representation.position = (x.coords[x.layout[i]][:, 1:].float() + 0.5) / self.resolution
-            representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(self.resolution)), dtype=torch.uint8, device='cuda')
-            for k, v in self.layout.items():
-                setattr(representation, k, x.feats[x.layout[i]][:, v['range'][0]:v['range'][1]].reshape(-1, *v['shape']))
-            representation.trivec = representation.trivec + 1
-            ret.append(representation)
-        return ret
-
-    def forward(self, x: sp.SparseTensor) -> List[Strivec]:
-        h = super().forward(x)
-        h = h.type(x.dtype)
-        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
-        h = self.out_layer(h)
-        return self.to_representation(h)

trellis/models/structured_latent_vae/encoder.py

@@ -1,72 +0,0 @@
-from typing import *
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
-from ...modules import sparse as sp
-from .base import SparseTransformerBase
-
-
-class SLatEncoder(SparseTransformerBase):
-    def __init__(
-        self,
-        resolution: int,
-        in_channels: 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,
-    ):
-        super().__init__(
-            in_channels=in_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.out_layer = sp.SparseLinear(model_channels, 2 * latent_channels)
-
-        self.initialize_weights()
-        if use_fp16:
-            self.convert_to_fp16()
-
-    def initialize_weights(self) -> None:
-        super().initialize_weights()
-        # Zero-out output layers:
-        nn.init.constant_(self.out_layer.weight, 0)
-        nn.init.constant_(self.out_layer.bias, 0)
-
-    def forward(self, x: sp.SparseTensor, sample_posterior=True, return_raw=False):
-        h = super().forward(x)
-        h = h.type(x.dtype)
-        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
-        h = self.out_layer(h)
-        
-        # Sample from the posterior distribution
-        mean, logvar = h.feats.chunk(2, dim=-1)
-        if sample_posterior:
-            std = torch.exp(0.5 * logvar)
-            z = mean + std * torch.randn_like(std)
-        else:
-            z = mean
-        z = h.replace(z)
-            
-        if return_raw:
-            return z, mean, logvar
-        else:
-            return z