Add trellis2

trellis2/__init__.py

@@ -0,0 +1,6 @@
+from . import models
+from . import modules
+from . import pipelines
+from . import renderers
+from . import representations
+from . import utils

trellis2/datasets/__init__.py

@@ -0,0 +1,46 @@
+import importlib
+
+__attributes = {
+    'FlexiDualGridDataset': 'flexi_dual_grid',
+    'SparseVoxelPbrDataset':'sparse_voxel_pbr',
+    
+    'SparseStructureLatent': 'sparse_structure_latent',
+    'TextConditionedSparseStructureLatent': 'sparse_structure_latent',
+    'ImageConditionedSparseStructureLatent': 'sparse_structure_latent',
+    
+    'SLat': 'structured_latent',
+    'ImageConditionedSLat': 'structured_latent',
+    'SLatShape': 'structured_latent_shape',
+    'ImageConditionedSLatShape': 'structured_latent_shape',
+    'SLatPbr': 'structured_latent_svpbr',
+    'ImageConditionedSLatPbr': 'structured_latent_svpbr',
+}
+
+__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]
+
+
+# For Pylance
+if __name__ == '__main__':    
+    from .flexi_dual_grid import FlexiDualGridDataset
+    from .sparse_voxel_pbr import SparseVoxelPbrDataset
+    
+    from .sparse_structure_latent import SparseStructureLatent, ImageConditionedSparseStructureLatent
+    from .structured_latent import SLat, ImageConditionedSLat
+    from .structured_latent_shape import SLatShape, ImageConditionedSLatShape
+    from .structured_latent_svpbr import SLatPbr, ImageConditionedSLatPbr
+    

trellis2/datasets/components.py

@@ -0,0 +1,192 @@
+from typing import *
+import json
+from abc import abstractmethod
+import os
+import json
+import torch
+import numpy as np
+import pandas as pd
+from PIL import Image
+from torch.utils.data import Dataset
+
+
+class StandardDatasetBase(Dataset):
+    """
+    Base class for standard datasets.
+
+    Args:
+        roots (str): paths to the dataset
+    """
+
+    def __init__(self,
+        roots: str,
+    ):
+        super().__init__()
+        try:
+            self.roots = json.loads(roots)
+            root_type = 'obj'
+        except:
+            self.roots = roots.split(',')
+            root_type = 'list'
+        self.instances = []
+        self.metadata = pd.DataFrame()
+        
+        self._stats = {}
+        if root_type == 'obj':
+            for key, root in self.roots.items():
+                self._stats[key] = {}
+                metadata = pd.DataFrame(columns=['sha256']).set_index('sha256')
+                for _, r in root.items():
+                    metadata = metadata.combine_first(pd.read_csv(os.path.join(r, 'metadata.csv')).set_index('sha256'))
+                self._stats[key]['Total'] = len(metadata)
+                metadata, stats = self.filter_metadata(metadata)
+                self._stats[key].update(stats)
+                self.instances.extend([(root, sha256) for sha256 in metadata.index.values])
+                self.metadata = pd.concat([self.metadata, metadata])
+        else:
+            for root in self.roots:
+                key = os.path.basename(root)
+                self._stats[key] = {}
+                metadata = pd.read_csv(os.path.join(root, 'metadata.csv'))
+                self._stats[key]['Total'] = len(metadata)
+                metadata, stats = self.filter_metadata(metadata)
+                self._stats[key].update(stats)
+                self.instances.extend([(root, sha256) for sha256 in metadata['sha256'].values])
+                metadata.set_index('sha256', inplace=True)
+                self.metadata = pd.concat([self.metadata, metadata])
+            
+    @abstractmethod
+    def filter_metadata(self, metadata: pd.DataFrame) -> Tuple[pd.DataFrame, Dict[str, int]]:
+        pass
+    
+    @abstractmethod
+    def get_instance(self, root, instance: str) -> Dict[str, Any]:
+        pass
+        
+    def __len__(self):
+        return len(self.instances)
+
+    def __getitem__(self, index) -> Dict[str, Any]:
+        try:
+            root, instance = self.instances[index]
+            return self.get_instance(root, instance)
+        except Exception as e:
+            print(f'Error loading {instance}: {e}')
+            return self.__getitem__(np.random.randint(0, len(self)))
+        
+    def __str__(self):
+        lines = []
+        lines.append(self.__class__.__name__)
+        lines.append(f'  - Total instances: {len(self)}')
+        lines.append(f'  - Sources:')
+        for key, stats in self._stats.items():
+            lines.append(f'    - {key}:')
+            for k, v in stats.items():
+                lines.append(f'      - {k}: {v}')
+        return '\n'.join(lines)
+
+
+class ImageConditionedMixin:
+    def __init__(self, roots, *, image_size=518, **kwargs):
+        self.image_size = image_size
+        super().__init__(roots, **kwargs)
+    
+    def filter_metadata(self, metadata):
+        metadata, stats = super().filter_metadata(metadata)
+        metadata = metadata[metadata['cond_rendered'].notna()]
+        stats['Cond rendered'] = len(metadata)
+        return metadata, stats
+    
+    def get_instance(self, root, instance):
+        pack = super().get_instance(root, instance)
+       
+        image_root = os.path.join(root['render_cond'], instance)
+        with open(os.path.join(image_root, 'transforms.json')) as f:
+            metadata = json.load(f)
+        n_views = len(metadata['frames'])
+        view = np.random.randint(n_views)
+        metadata = metadata['frames'][view]
+
+        image_path = os.path.join(image_root, metadata['file_path'])
+        image = Image.open(image_path)
+
+        alpha = np.array(image.getchannel(3))
+        bbox = np.array(alpha).nonzero()
+        bbox = [bbox[1].min(), bbox[0].min(), bbox[1].max(), bbox[0].max()]
+        center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
+        hsize = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
+        aug_hsize = hsize
+        aug_center_offset = [0, 0]
+        aug_center = [center[0] + aug_center_offset[0], center[1] + aug_center_offset[1]]
+        aug_bbox = [int(aug_center[0] - aug_hsize), int(aug_center[1] - aug_hsize), int(aug_center[0] + aug_hsize), int(aug_center[1] + aug_hsize)]
+        image = image.crop(aug_bbox)
+
+        image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+        alpha = image.getchannel(3)
+        image = image.convert('RGB')
+        image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
+        alpha = torch.tensor(np.array(alpha)).float() / 255.0
+        image = image * alpha.unsqueeze(0)
+        pack['cond'] = image
+       
+        return pack
+
+
+class MultiImageConditionedMixin:
+    def __init__(self, roots, *, image_size=518, max_image_cond_view = 4, **kwargs):
+        self.image_size = image_size
+        self.max_image_cond_view = max_image_cond_view
+        super().__init__(roots, **kwargs)
+
+    def filter_metadata(self, metadata):
+        metadata, stats = super().filter_metadata(metadata)
+        metadata = metadata[metadata['cond_rendered'].notna()]
+        stats['Cond rendered'] = len(metadata)
+        return metadata, stats
+    
+    def get_instance(self, root, instance):
+        pack = super().get_instance(root, instance)
+       
+        image_root = os.path.join(root['render_cond'], instance)
+        with open(os.path.join(image_root, 'transforms.json')) as f:
+            metadata = json.load(f)
+
+        n_views = len(metadata['frames'])
+        n_sample_views = np.random.randint(1, self.max_image_cond_view+1)
+
+        assert n_views >= n_sample_views, f'Not enough views to sample {n_sample_views} unique images.'
+
+        sampled_views = np.random.choice(n_views, size=n_sample_views, replace=False)
+
+        cond_images = []
+        for v in sampled_views:
+            frame_info = metadata['frames'][v]
+            image_path = os.path.join(image_root, frame_info['file_path'])
+            image = Image.open(image_path)
+
+            alpha = np.array(image.getchannel(3))
+            bbox = np.array(alpha).nonzero()
+            bbox = [bbox[1].min(), bbox[0].min(), bbox[1].max(), bbox[0].max()]
+            center = [(bbox[0] + bbox[2]) / 2, (bbox[1] + bbox[3]) / 2]
+            hsize = max(bbox[2] - bbox[0], bbox[3] - bbox[1]) / 2
+            aug_hsize = hsize
+            aug_center = center
+            aug_bbox = [
+                int(aug_center[0] - aug_hsize),
+                int(aug_center[1] - aug_hsize),
+                int(aug_center[0] + aug_hsize),
+                int(aug_center[1] + aug_hsize),
+            ]
+
+            img = image.crop(aug_bbox)
+            img = img.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+            alpha = img.getchannel(3)
+            img = img.convert('RGB')
+            img = torch.tensor(np.array(img)).permute(2, 0, 1).float() / 255.0
+            alpha = torch.tensor(np.array(alpha)).float() / 255.0
+            img = img * alpha.unsqueeze(0)
+
+            cond_images.append(img)
+
+        pack['cond'] = [torch.stack(cond_images, dim=0)]  # (V,3,H,W)
+        return pack

trellis2/datasets/flexi_dual_grid.py

@@ -0,0 +1,173 @@
+import os
+import numpy as np
+import pickle
+import torch
+import utils3d
+from .components import StandardDatasetBase
+from ..modules import sparse as sp
+from ..renderers import MeshRenderer
+from ..representations import Mesh
+from ..utils.data_utils import load_balanced_group_indices
+import o_voxel
+
+
+class FlexiDualGridVisMixin:
+    @torch.no_grad()
+    def visualize_sample(self, x: dict):
+        mesh = x['mesh']
+        
+        renderer = MeshRenderer({'near': 1, 'far': 3})
+        renderer.rendering_options.resolution = 512
+        renderer.rendering_options.ssaa = 4
+        
+        # Build camera
+        yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
+        yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
+        yaws = [y + yaws_offset for y in yaws]
+        pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
+
+        exts = []
+        ints = []
+        for yaw, pitch in zip(yaws, pitch):
+            orig = torch.tensor([
+                np.sin(yaw) * np.cos(pitch),
+                np.cos(yaw) * np.cos(pitch),
+                np.sin(pitch),
+            ]).float().cuda() * 2
+            fov = torch.deg2rad(torch.tensor(30)).cuda()
+            extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
+            intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
+            exts.append(extrinsics)
+            ints.append(intrinsics)
+        
+        # Build each representation
+        images = []
+        for m in mesh:
+            image = torch.zeros(3, 1024, 1024).cuda()
+            tile = [2, 2]
+            for j, (ext, intr) in enumerate(zip(exts, ints)):
+                image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = \
+                    renderer.render(m.cuda(), ext, intr)['normal']
+            images.append(image)
+        images = torch.stack(images)
+        
+        return images
+    
+
+class FlexiDualGridDataset(FlexiDualGridVisMixin, StandardDatasetBase):
+    """
+    Flexible Dual Grid Dataset
+    
+    Args:
+        roots (str): path to the dataset
+        resolution (int): resolution of the voxel grid
+        min_aesthetic_score (float): minimum aesthetic score of the instances to be included in the dataset
+    """
+
+    def __init__(
+        self,
+        roots,
+        resolution: int = 1024,
+        max_active_voxels: int = 1000000,
+        max_num_faces: int = None,
+        min_aesthetic_score: float = 5.0,
+    ):
+        self.resolution = resolution
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_active_voxels = max_active_voxels
+        self.max_num_faces = max_num_faces
+        self.value_range = (0, 1)
+
+        super().__init__(roots)
+        
+        self.loads = [self.metadata.loc[sha256, f'dual_grid_size'] for _, sha256 in self.instances]
+        
+    def __str__(self):
+        lines = [
+            super().__str__(),
+            f'  - Resolution: {self.resolution}',
+        ]
+        return '\n'.join(lines)
+        
+    def filter_metadata(self, metadata):
+        stats = {}
+        metadata = metadata[metadata[f'dual_grid_converted'] == True]
+        stats['Dual Grid Converted'] = len(metadata)
+        if self.min_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+            stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        metadata = metadata[metadata[f'dual_grid_size'] <= self.max_active_voxels]
+        stats[f'Active Voxels <= {self.max_active_voxels}'] = len(metadata)
+        if self.max_num_faces is not None:
+            metadata = metadata[metadata['num_faces'] <= self.max_num_faces]
+            stats[f'Faces <= {self.max_num_faces}'] = len(metadata)
+        return metadata, stats
+    
+    def read_mesh(self, root, instance):
+        with open(os.path.join(root, f'{instance}.pickle'), 'rb') as f:
+            dump = pickle.load(f)
+        start = 0
+        vertices = []
+        faces = []
+        for obj in dump['objects']:
+            if obj['vertices'].size == 0 or obj['faces'].size == 0:
+                continue
+            vertices.append(obj['vertices'])
+            faces.append(obj['faces'] + start)
+            start += len(obj['vertices'])
+        vertices = torch.from_numpy(np.concatenate(vertices, axis=0)).float()
+        faces = torch.from_numpy(np.concatenate(faces, axis=0)).long()
+        vertices_min = vertices.min(dim=0)[0]
+        vertices_max = vertices.max(dim=0)[0]
+        center = (vertices_min + vertices_max) / 2
+        scale = 0.99999 / (vertices_max - vertices_min).max()
+        vertices = (vertices - center) * scale
+        assert torch.all(vertices >= -0.5) and torch.all(vertices <= 0.5), 'vertices out of range'
+        return {'mesh': [Mesh(vertices=vertices, faces=faces)]}
+    
+    def read_dual_grid(self, root, instance):
+        coords, attr = o_voxel.io.read_vxz(os.path.join(root, f'{instance}.vxz'), num_threads=4)
+        vertices = sp.SparseTensor(
+            (attr['vertices'] / 255.0).float(),
+            torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
+        )
+        intersected = vertices.replace(torch.cat([
+            attr['intersected'] % 2,
+            attr['intersected'] // 2 % 2,
+            attr['intersected'] // 4 % 2,
+        ], dim=-1).bool())
+        return {'vertices': vertices, 'intersected': intersected}
+
+    def get_instance(self, root, instance):
+        mesh = self.read_mesh(root['mesh_dump'], instance)
+        dual_grid = self.read_dual_grid(root['dual_grid'], instance)
+        return {**mesh, **dual_grid}
+    
+    @staticmethod
+    def collate_fn(batch, split_size=None):
+        if split_size is None:
+            group_idx = [list(range(len(batch)))]
+        else:
+            group_idx = load_balanced_group_indices([b['vertices'].feats.shape[0] for b in batch], split_size)
+        packs = []
+        for group in group_idx:
+            sub_batch = [batch[i] for i in group]
+            pack = {}
+
+            keys = [k for k in sub_batch[0].keys()]
+            for k in keys:
+                if isinstance(sub_batch[0][k], torch.Tensor):
+                    pack[k] = torch.stack([b[k] for b in sub_batch])
+                elif isinstance(sub_batch[0][k], sp.SparseTensor):
+                    pack[k] = sp.sparse_cat([b[k] for b in sub_batch], dim=0)
+                elif isinstance(sub_batch[0][k], list):
+                    pack[k] = sum([b[k] for b in sub_batch], [])
+                else:
+                    pack[k] = [b[k] for b in sub_batch]
+            
+            packs.append(pack)
+        
+        if split_size is None:
+            return packs[0]
+        return packs
+    

trellis2/datasets/sparse_structure_latent.py

@@ -0,0 +1,160 @@
+import os
+import json
+from typing import *
+import numpy as np
+import torch
+from ..representations import Voxel
+from ..renderers import VoxelRenderer
+from .components import StandardDatasetBase, ImageConditionedMixin
+from .. import models
+from ..utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics
+
+
+class SparseStructureLatentVisMixin:
+    def __init__(
+        self,
+        *args,
+        pretrained_ss_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16.json',
+        ss_dec_path: Optional[str] = None,
+        ss_dec_ckpt: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.ss_dec = None
+        self.pretrained_ss_dec = pretrained_ss_dec
+        self.ss_dec_path = ss_dec_path
+        self.ss_dec_ckpt = ss_dec_ckpt
+        
+    def _loading_ss_dec(self):
+        if self.ss_dec is not None:
+            return
+        if self.ss_dec_path is not None:
+            cfg = json.load(open(os.path.join(self.ss_dec_path, 'config.json'), 'r'))
+            decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
+            ckpt_path = os.path.join(self.ss_dec_path, 'ckpts', f'decoder_{self.ss_dec_ckpt}.pt')
+            decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+        else:
+            decoder = models.from_pretrained(self.pretrained_ss_dec)
+        self.ss_dec = decoder.cuda().eval()
+
+    def _delete_ss_dec(self):
+        del self.ss_dec
+        self.ss_dec = None
+
+    @torch.no_grad()
+    def decode_latent(self, z, batch_size=4):
+        self._loading_ss_dec()
+        ss = []
+        if self.normalization:
+            z = z * self.std.to(z.device) + self.mean.to(z.device)
+        for i in range(0, z.shape[0], batch_size):
+            ss.append(self.ss_dec(z[i:i+batch_size]))
+        ss = torch.cat(ss, dim=0)
+        self._delete_ss_dec()
+        return ss
+
+    @torch.no_grad()
+    def visualize_sample(self, x_0: Union[torch.Tensor, dict]):
+        x_0 = x_0 if isinstance(x_0, torch.Tensor) else x_0['x_0']
+        x_0 = self.decode_latent(x_0.cuda())
+        
+        renderer = VoxelRenderer()
+        renderer.rendering_options.resolution = 512
+        renderer.rendering_options.ssaa = 4
+        
+        # build camera
+        yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
+        yaw_offset = -16 / 180 * np.pi
+        yaw = [y + yaw_offset for y in yaw]
+        pitch = [20 / 180 * np.pi for _ in range(4)]
+        exts, ints = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, 2, 30)
+
+        images = []
+        
+        # Build each representation
+        x_0 = x_0.cuda()
+        for i in range(x_0.shape[0]):
+            coords = torch.nonzero(x_0[i, 0] > 0, as_tuple=False)
+            resolution = x_0.shape[-1]
+            color = coords / resolution
+            rep = Voxel(
+                origin=[-0.5, -0.5, -0.5],
+                voxel_size=1/resolution,
+                coords=coords,
+                attrs=color,
+                layout={
+                    'color': slice(0, 3),
+                }
+            )
+            image = torch.zeros(3, 1024, 1024).cuda()
+            tile = [2, 2]
+            for j, (ext, intr) in enumerate(zip(exts, ints)):
+                res = renderer.render(rep, ext, intr, colors_overwrite=color)
+                image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
+            images.append(image)
+            
+        return torch.stack(images)
+
+
+class SparseStructureLatent(SparseStructureLatentVisMixin, StandardDatasetBase):
+    """
+    Sparse structure latent dataset
+    
+    Args:
+        roots (str): path to the dataset
+        min_aesthetic_score (float): minimum aesthetic score
+        normalization (dict): normalization stats
+        pretrained_ss_dec (str): name of the pretrained sparse structure decoder
+        ss_dec_path (str): path to the sparse structure decoder, if given, will override the pretrained_ss_dec
+        ss_dec_ckpt (str): name of the sparse structure decoder checkpoint
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        min_aesthetic_score: float = 5.0,
+        normalization: Optional[dict] = None,
+        pretrained_ss_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/ss_dec_conv3d_16l8_fp16',
+        ss_dec_path: Optional[str] = None,
+        ss_dec_ckpt: Optional[str] = None,
+    ):
+        self.min_aesthetic_score = min_aesthetic_score
+        self.normalization = normalization
+        self.value_range = (0, 1)
+        
+        super().__init__(
+            roots,
+            pretrained_ss_dec=pretrained_ss_dec,
+            ss_dec_path=ss_dec_path,
+            ss_dec_ckpt=ss_dec_ckpt,
+        )
+        
+        if self.normalization is not None:
+            self.mean = torch.tensor(self.normalization['mean']).reshape(-1, 1, 1, 1)
+            self.std = torch.tensor(self.normalization['std']).reshape(-1, 1, 1, 1)
+  
+    def filter_metadata(self, metadata):
+        stats = {}
+        metadata = metadata[metadata['ss_latent_encoded'] == True]
+        stats['With latent'] = len(metadata)
+        metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+        stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        return metadata, stats
+                
+    def get_instance(self, root, instance):
+        latent = np.load(os.path.join(root['ss_latent'], f'{instance}.npz'))
+        z = torch.tensor(latent['z']).float()
+        if self.normalization is not None:
+            z = (z - self.mean) / self.std
+
+        pack = {
+            'x_0': z,
+        }
+        return pack
+
+
+class ImageConditionedSparseStructureLatent(ImageConditionedMixin, SparseStructureLatent):
+    """
+    Image-conditioned sparse structure dataset
+    """
+    pass
+    

trellis2/datasets/sparse_voxel_pbr.py

@@ -0,0 +1,298 @@
+import os
+import io
+from typing import Union
+import numpy as np
+import pickle
+import torch
+from PIL import Image
+import o_voxel
+import utils3d
+from .components import StandardDatasetBase
+from ..modules import sparse as sp
+from ..renderers import VoxelRenderer
+from ..representations import Voxel
+from ..representations.mesh import MeshWithPbrMaterial, TextureFilterMode, TextureWrapMode, AlphaMode, PbrMaterial, Texture
+
+from ..utils.data_utils import load_balanced_group_indices
+
+
+def is_power_of_two(n: int) -> bool:
+    return n > 0 and (n & (n - 1)) == 0
+
+
+def nearest_power_of_two(n: int) -> int:
+    if n < 1:
+        raise ValueError("n must be >= 1")
+    if is_power_of_two(n):
+        return n
+    lower = 2 ** (n.bit_length() - 1)
+    upper = 2 ** n.bit_length()
+    if n - lower < upper - n:
+        return lower
+    else:
+        return upper
+    
+
+class SparseVoxelPbrVisMixin:
+    @torch.no_grad()
+    def visualize_sample(self, x: Union[sp.SparseTensor, dict]):
+        x = x if isinstance(x, sp.SparseTensor) else x['x']
+        
+        renderer = VoxelRenderer()
+        renderer.rendering_options.resolution = 512
+        renderer.rendering_options.ssaa = 4
+        
+        # Build camera
+        yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
+        yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
+        yaws = [y + yaws_offset for y in yaws]
+        pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
+
+        exts = []
+        ints = []
+        for yaw, pitch in zip(yaws, pitch):
+            orig = torch.tensor([
+                np.sin(yaw) * np.cos(pitch),
+                np.cos(yaw) * np.cos(pitch),
+                np.sin(pitch),
+            ]).float().cuda() * 2
+            fov = torch.deg2rad(torch.tensor(30)).cuda()
+            extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
+            intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
+            exts.append(extrinsics)
+            ints.append(intrinsics)
+
+        images = {k: [] for k in self.layout}
+        
+        # Build each representation
+        x = x.cuda()
+        for i in range(x.shape[0]):
+            rep = Voxel(
+                origin=[-0.5, -0.5, -0.5],
+                voxel_size=1/self.resolution,
+                coords=x[i].coords[:, 1:].contiguous(),
+                attrs=None,
+                layout={
+                    'color': slice(0, 3),
+                }
+            )
+            for k in self.layout:
+                image = torch.zeros(3, 1024, 1024).cuda()
+                tile = [2, 2]
+                for j, (ext, intr) in enumerate(zip(exts, ints)):
+                    attr = x[i].feats[:, self.layout[k]].expand(-1, 3)
+                    res = renderer.render(rep, ext, intr, colors_overwrite=attr)
+                    image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
+                images[k].append(image)
+        
+        for k in self.layout:
+            images[k] = torch.stack(images[k])
+            
+        return images
+
+
+class SparseVoxelPbrDataset(SparseVoxelPbrVisMixin, StandardDatasetBase):
+    """
+    Sparse Voxel PBR dataset.
+    
+    Args:
+        roots (str): path to the dataset
+        resolution (int): resolution of the voxel grid
+        min_aesthetic_score (float): minimum aesthetic score of the instances to be included in the dataset
+    """
+
+    def __init__(
+        self,
+        roots,
+        resolution: int = 1024,
+        max_active_voxels: int = 1000000,
+        max_num_faces: int = None,
+        min_aesthetic_score: float = 5.0,
+        attrs: list[str] = ['base_color', 'metallic', 'roughness', 'emissive', 'alpha'],
+        with_mesh: bool = True,
+    ):
+        self.resolution = resolution
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_active_voxels = max_active_voxels
+        self.max_num_faces = max_num_faces
+        self.with_mesh = with_mesh
+        self.value_range = (-1, 1)
+        self.channels = {
+            'base_color': 3,
+            'metallic': 1,
+            'roughness': 1,
+            'emissive': 3,
+            'alpha': 1,
+        }
+        self.layout = {}
+        start = 0
+        for attr in attrs:
+            self.layout[attr] = slice(start, start + self.channels[attr])
+            start += self.channels[attr]
+
+        super().__init__(roots)
+        
+        self.loads = [self.metadata.loc[sha256, f'num_pbr_voxels'] for _, sha256 in self.instances]
+        
+    def __str__(self):
+        lines = [
+            super().__str__(),
+            f'  - Resolution: {self.resolution}',
+            f'  - Attributes: {list(self.layout.keys())}',
+        ]
+        return '\n'.join(lines)
+        
+    def filter_metadata(self, metadata):
+        stats = {}
+        metadata = metadata[metadata['pbr_voxelized'] == True]
+        stats['PBR Voxelized'] = len(metadata)
+        if self.min_aesthetic_score is not None:
+            metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+            stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        metadata = metadata[metadata['num_pbr_voxels'] <= self.max_active_voxels]
+        stats[f'Active voxels <= {self.max_active_voxels}'] = len(metadata)
+        if self.max_num_faces is not None:
+            metadata = metadata[metadata['num_faces'] <= self.max_num_faces]
+            stats[f'Faces <= {self.max_num_faces}'] = len(metadata)
+        return metadata, stats
+
+    @staticmethod
+    def _texture_from_dump(pack) -> Texture:
+        png_bytes = pack['image']
+        image = Image.open(io.BytesIO(png_bytes))
+        if image.width != image.height or not is_power_of_two(image.width):
+            size = nearest_power_of_two(max(image.width, image.height))
+            image = image.resize((size, size), Image.LANCZOS)
+        texture = torch.tensor(np.array(image) / 255.0, dtype=torch.float32).reshape(image.height, image.width, -1)
+        filter_mode = {
+            'Linear': TextureFilterMode.LINEAR,
+            'Closest': TextureFilterMode.CLOSEST,
+            'Cubic': TextureFilterMode.LINEAR,
+            'Smart': TextureFilterMode.LINEAR,
+        }[pack['interpolation']]
+        wrap_mode = {
+            'REPEAT': TextureWrapMode.REPEAT,
+            'EXTEND': TextureWrapMode.CLAMP_TO_EDGE,
+            'CLIP': TextureWrapMode.CLAMP_TO_EDGE,
+            'MIRROR': TextureWrapMode.MIRRORED_REPEAT,
+        }[pack['extension']]
+        return Texture(texture, filter_mode=filter_mode, wrap_mode=wrap_mode)
+
+    def read_mesh_with_texture(self, root, instance):
+        with open(os.path.join(root, f'{instance}.pickle'), 'rb') as f:
+            dump = pickle.load(f)
+            
+        # Fix dump alpha map
+        for mat in dump['materials']:
+            if mat['alphaTexture'] is not None and mat['alphaMode'] == 'OPAQUE':
+                mat['alphaMode'] = 'BLEND'
+
+        # process material
+        materials = []
+        for mat in dump['materials']:
+            materials.append(PbrMaterial(
+                base_color_texture=self._texture_from_dump(mat['baseColorTexture']) if mat['baseColorTexture'] is not None else None,
+                base_color_factor=mat['baseColorFactor'],
+                metallic_texture=self._texture_from_dump(mat['metallicTexture']) if mat['metallicTexture'] is not None else None,
+                metallic_factor=mat['metallicFactor'],
+                roughness_texture=self._texture_from_dump(mat['roughnessTexture']) if mat['roughnessTexture'] is not None else None,
+                roughness_factor=mat['roughnessFactor'],
+                alpha_texture=self._texture_from_dump(mat['alphaTexture']) if mat['alphaTexture'] is not None else None,
+                alpha_factor=mat['alphaFactor'],
+                alpha_mode={
+                    'OPAQUE': AlphaMode.OPAQUE,
+                    'MASK': AlphaMode.MASK,
+                    'BLEND': AlphaMode.BLEND,
+                }[mat['alphaMode']],
+                alpha_cutoff=mat['alphaCutoff'],
+            ))
+        materials.append(PbrMaterial(
+            base_color_factor=[0.8, 0.8, 0.8],
+            alpha_factor=1.0,
+            metallic_factor=0.0,
+            roughness_factor=0.5,
+            alpha_mode=AlphaMode.OPAQUE,
+            alpha_cutoff=0.5,
+        ))  # append default material
+
+        # process mesh
+        start = 0
+        vertices = []
+        faces = []
+        material_ids = []
+        uv_coords = []
+        for obj in dump['objects']:
+            if obj['vertices'].size == 0 or obj['faces'].size == 0:
+                continue
+            vertices.append(obj['vertices'])
+            faces.append(obj['faces'] + start)
+            obj['mat_ids'][obj['mat_ids'] == -1] = len(materials) - 1
+            material_ids.append(obj['mat_ids'])
+            uv_coords.append(obj['uvs'] if obj['uvs'] is not None else np.zeros((obj['faces'].shape[0], 3, 2), dtype=np.float32))
+            start += len(obj['vertices'])
+        
+        vertices = torch.from_numpy(np.concatenate(vertices, axis=0)).float()
+        faces = torch.from_numpy(np.concatenate(faces, axis=0)).long()
+        material_ids = torch.from_numpy(np.concatenate(material_ids, axis=0)).long()
+        uv_coords = torch.from_numpy(np.concatenate(uv_coords, axis=0)).float()
+        
+        # Normalize vertices
+        vertices_min = vertices.min(dim=0)[0]
+        vertices_max = vertices.max(dim=0)[0]
+        center = (vertices_min + vertices_max) / 2
+        scale = 0.99999 / (vertices_max - vertices_min).max()
+        vertices = (vertices - center) * scale
+        assert torch.all(vertices >= -0.5) and torch.all(vertices <= 0.5), 'vertices out of range'
+        
+        return {'mesh': [MeshWithPbrMaterial(
+            vertices=vertices,
+            faces=faces,
+            material_ids=material_ids,
+            uv_coords=uv_coords,
+            materials=materials,
+        )]}
+
+    def read_pbr_voxel(self, root, instance):
+        coords, attr = o_voxel.io.read_vxz(os.path.join(root, f'{instance}.vxz'), num_threads=4)
+        feats = torch.concat([attr[k] for k in self.layout], dim=-1) / 255.0 * 2 - 1
+        x = sp.SparseTensor(
+            feats.float(),
+            torch.cat([torch.zeros_like(coords[:, 0:1]), coords], dim=-1),
+        )
+        return {'x': x}
+    
+    def get_instance(self, root, instance):
+        if self.with_mesh:
+            mesh = self.read_mesh_with_texture(root['pbr_dump'], instance)
+            pbr_voxel = self.read_pbr_voxel(root['pbr_voxel'], instance)
+            return {**mesh, **pbr_voxel}
+        else:
+            return self.read_pbr_voxel(root['pbr_voxel'], instance)
+    
+    @staticmethod
+    def collate_fn(batch, split_size=None):
+        if split_size is None:
+            group_idx = [list(range(len(batch)))]
+        else:
+            group_idx = load_balanced_group_indices([b['x'].feats.shape[0] for b in batch], split_size)
+        packs = []
+        for group in group_idx:
+            sub_batch = [batch[i] for i in group]
+            pack = {}
+
+            keys = [k for k in sub_batch[0].keys()]
+            for k in keys:
+                if isinstance(sub_batch[0][k], torch.Tensor):
+                    pack[k] = torch.stack([b[k] for b in sub_batch])
+                elif isinstance(sub_batch[0][k], sp.SparseTensor):
+                    pack[k] = sp.sparse_cat([b[k] for b in sub_batch], dim=0)
+                elif isinstance(sub_batch[0][k], list):
+                    pack[k] = sum([b[k] for b in sub_batch], [])
+                else:
+                    pack[k] = [b[k] for b in sub_batch]
+            
+            packs.append(pack)
+        
+        if split_size is None:
+            return packs[0]
+        return packs

trellis2/datasets/structured_latent.py

@@ -0,0 +1,210 @@
+import json
+import os
+from typing import *
+import numpy as np
+import torch
+import utils3d.torch
+from .components import StandardDatasetBase, ImageConditionedMixin
+from ..modules.sparse.basic import SparseTensor
+from .. import models
+from ..utils.render_utils import get_renderer
+from ..utils.data_utils import load_balanced_group_indices
+
+
+class SLatVisMixin:
+    def __init__(
+        self,
+        *args,
+        pretrained_slat_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16',
+        slat_dec_path: Optional[str] = None,
+        slat_dec_ckpt: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.slat_dec = None
+        self.pretrained_slat_dec = pretrained_slat_dec
+        self.slat_dec_path = slat_dec_path
+        self.slat_dec_ckpt = slat_dec_ckpt
+        
+    def _loading_slat_dec(self):
+        if self.slat_dec is not None:
+            return
+        if self.slat_dec_path is not None:
+            cfg = json.load(open(os.path.join(self.slat_dec_path, 'config.json'), 'r'))
+            decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
+            ckpt_path = os.path.join(self.slat_dec_path, 'ckpts', f'decoder_{self.slat_dec_ckpt}.pt')
+            decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+        else:
+            decoder = models.from_pretrained(self.pretrained_slat_dec)
+        self.slat_dec = decoder.cuda().eval()
+
+    def _delete_slat_dec(self):
+        del self.slat_dec
+        self.slat_dec = None
+
+    @torch.no_grad()
+    def decode_latent(self, z, batch_size=4):
+        self._loading_slat_dec()
+        reps = []
+        if self.normalization is not None:
+            z = z * self.std.to(z.device) + self.mean.to(z.device)
+        for i in range(0, z.shape[0], batch_size):
+            reps.append(self.slat_dec(z[i:i+batch_size]))
+        reps = sum(reps, [])
+        self._delete_slat_dec()
+        return reps
+
+    @torch.no_grad()
+    def visualize_sample(self, x_0: Union[SparseTensor, dict]):
+        x_0 = x_0 if isinstance(x_0, SparseTensor) else x_0['x_0']
+        reps = self.decode_latent(x_0.cuda())
+        
+        # Build camera
+        yaws = [0, np.pi / 2, np.pi, 3 * np.pi / 2]
+        yaws_offset = np.random.uniform(-np.pi / 4, np.pi / 4)
+        yaws = [y + yaws_offset for y in yaws]
+        pitch = [np.random.uniform(-np.pi / 4, np.pi / 4) for _ in range(4)]
+
+        exts = []
+        ints = []
+        for yaw, pitch in zip(yaws, pitch):
+            orig = torch.tensor([
+                np.sin(yaw) * np.cos(pitch),
+                np.cos(yaw) * np.cos(pitch),
+                np.sin(pitch),
+            ]).float().cuda() * 2
+            fov = torch.deg2rad(torch.tensor(40)).cuda()
+            extrinsics = utils3d.torch.extrinsics_look_at(orig, torch.tensor([0, 0, 0]).float().cuda(), torch.tensor([0, 0, 1]).float().cuda())
+            intrinsics = utils3d.torch.intrinsics_from_fov_xy(fov, fov)
+            exts.append(extrinsics)
+            ints.append(intrinsics)
+
+        renderer = get_renderer(reps[0])
+        images = []
+        for representation in reps:
+            image = torch.zeros(3, 1024, 1024).cuda()
+            tile = [2, 2]
+            for j, (ext, intr) in enumerate(zip(exts, ints)):
+                res = renderer.render(representation, ext, intr)
+                image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['color']
+            images.append(image)
+        images = torch.stack(images)
+            
+        return images
+    
+    
+class SLat(SLatVisMixin, StandardDatasetBase):
+    """
+    structured latent V2 dataset
+    
+    Args:
+        roots (str): path to the dataset
+        min_aesthetic_score (float): minimum aesthetic score
+        max_tokens (int): maximum number of tokens
+        latent_key (str): key of the latent to be used
+        normalization (dict): normalization stats
+        pretrained_slat_dec (str): name of the pretrained slat decoder
+        slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
+        slat_dec_ckpt (str): name of the slat decoder checkpoint
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        min_aesthetic_score: float = 5.0,
+        max_tokens: int = 32768,
+        latent_key: str = 'shape_latent',
+        normalization: Optional[dict] = None,
+        pretrained_slat_dec: str = 'JeffreyXiang/TRELLIS-image-large/ckpts/slat_dec_gs_swin8_B_64l8gs32_fp16',
+        slat_dec_path: Optional[str] = None,
+        slat_dec_ckpt: Optional[str] = None,
+    ):
+        self.normalization = normalization
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_tokens = max_tokens
+        self.latent_key = latent_key
+        self.value_range = (0, 1)
+        
+        super().__init__(
+            roots,
+            pretrained_slat_dec=pretrained_slat_dec,
+            slat_dec_path=slat_dec_path,
+            slat_dec_ckpt=slat_dec_ckpt,
+        )
+
+        self.loads = [self.metadata.loc[sha256, f'{latent_key}_tokens'] for _, sha256 in self.instances]
+        
+        if self.normalization is not None:
+            self.mean = torch.tensor(self.normalization['mean']).reshape(1, -1)
+            self.std = torch.tensor(self.normalization['std']).reshape(1, -1)
+      
+    def filter_metadata(self, metadata):
+        stats = {}
+        metadata = metadata[metadata[f'{self.latent_key}_encoded'] == True]
+        stats['With latent'] = len(metadata)
+        metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+        stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        metadata = metadata[metadata[f'{self.latent_key}_tokens'] <= self.max_tokens]
+        stats[f'Num tokens <= {self.max_tokens}'] = len(metadata)
+        return metadata, stats
+
+    def get_instance(self, root, instance):
+        data = np.load(os.path.join(root[self.latent_key], f'{instance}.npz'))
+        coords = torch.tensor(data['coords']).int()
+        feats = torch.tensor(data['feats']).float()
+        if self.normalization is not None:
+            feats = (feats - self.mean) / self.std
+        return {
+            'coords': coords,
+            'feats': feats,
+        }
+        
+    @staticmethod
+    def collate_fn(batch, split_size=None):
+        if split_size is None:
+            group_idx = [list(range(len(batch)))]
+        else:
+            group_idx = load_balanced_group_indices([b['coords'].shape[0] for b in batch], split_size)
+        packs = []
+        for group in group_idx:
+            sub_batch = [batch[i] for i in group]
+            pack = {}
+            coords = []
+            feats = []
+            layout = []
+            start = 0
+            for i, b in enumerate(sub_batch):
+                coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
+                feats.append(b['feats'])
+                layout.append(slice(start, start + b['coords'].shape[0]))
+                start += b['coords'].shape[0]
+            coords = torch.cat(coords)
+            feats = torch.cat(feats)
+            pack['x_0'] = SparseTensor(
+                coords=coords,
+                feats=feats,
+            )
+            pack['x_0']._shape = torch.Size([len(group), *sub_batch[0]['feats'].shape[1:]])
+            pack['x_0'].register_spatial_cache('layout', layout)
+            
+            # collate other data
+            keys = [k for k in sub_batch[0].keys() if k not in ['coords', 'feats']]
+            for k in keys:
+                if isinstance(sub_batch[0][k], torch.Tensor):
+                    pack[k] = torch.stack([b[k] for b in sub_batch])
+                elif isinstance(sub_batch[0][k], list):
+                    pack[k] = sum([b[k] for b in sub_batch], [])
+                else:
+                    pack[k] = [b[k] for b in sub_batch]
+                    
+            packs.append(pack)
+          
+        if split_size is None:
+            return packs[0]
+        return packs
+
+
+class ImageConditionedSLat(ImageConditionedMixin, SLat):
+    """
+    Image conditioned structured latent dataset
+    """
+    pass

trellis2/datasets/structured_latent_shape.py

@@ -0,0 +1,96 @@
+import os
+import json
+from typing import *
+import numpy as np
+import torch
+from .. import models
+from .components import ImageConditionedMixin
+from ..modules.sparse import SparseTensor
+from .structured_latent import SLatVisMixin, SLat
+from ..utils.render_utils import get_renderer, yaw_pitch_r_fov_to_extrinsics_intrinsics
+
+
+class SLatShapeVisMixin(SLatVisMixin):
+    def _loading_slat_dec(self):
+        if self.slat_dec is not None:
+            return
+        if self.slat_dec_path is not None:
+            cfg = json.load(open(os.path.join(self.slat_dec_path, 'config.json'), 'r'))
+            decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
+            ckpt_path = os.path.join(self.slat_dec_path, 'ckpts', f'decoder_{self.slat_dec_ckpt}.pt')
+            decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+        else:
+            decoder = models.from_pretrained(self.pretrained_slat_dec)
+        decoder.set_resolution(self.resolution)
+        self.slat_dec = decoder.cuda().eval()
+
+    @torch.no_grad()
+    def visualize_sample(self, x_0: Union[SparseTensor, dict]):
+        x_0 = x_0 if isinstance(x_0, SparseTensor) else x_0['x_0']
+        reps = self.decode_latent(x_0.cuda())
+        
+        # build camera
+        yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
+        yaw_offset = -16 / 180 * np.pi
+        yaw = [y + yaw_offset for y in yaw]
+        pitch = [20 / 180 * np.pi for _ in range(4)]
+        exts, ints = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, 2, 30)
+        
+        # render
+        renderer = get_renderer(reps[0])
+        images = []
+        for representation in reps:
+            image = torch.zeros(3, 1024, 1024).cuda()
+            tile = [2, 2]
+            for j, (ext, intr) in enumerate(zip(exts, ints)):
+                res = renderer.render(representation, ext, intr)
+                image[:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = res['normal']
+            images.append(image)
+        images = torch.stack(images)
+        return images
+    
+    
+class SLatShape(SLatShapeVisMixin, SLat):
+    """
+    structured latent for shape generation
+    
+    Args:
+        roots (str): path to the dataset
+        resolution (int): resolution of the shape
+        min_aesthetic_score (float): minimum aesthetic score
+        max_tokens (int): maximum number of tokens
+        latent_key (str): key of the latent to be used
+        normalization (dict): normalization stats
+        pretrained_slat_dec (str): name of the pretrained slat decoder
+        slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
+        slat_dec_ckpt (str): name of the slat decoder checkpoint
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        resolution: int,
+        min_aesthetic_score: float = 5.0,
+        max_tokens: int = 32768,
+        normalization: Optional[dict] = None,
+        pretrained_slat_dec: str = 'microsoft/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
+        slat_dec_path: Optional[str] = None,
+        slat_dec_ckpt: Optional[str] = None,
+    ):
+        super().__init__(
+            roots,
+            min_aesthetic_score=min_aesthetic_score,
+            max_tokens=max_tokens,
+            latent_key='shape_latent',
+            normalization=normalization,
+            pretrained_slat_dec=pretrained_slat_dec,
+            slat_dec_path=slat_dec_path,
+            slat_dec_ckpt=slat_dec_ckpt,
+        )
+        self.resolution = resolution
+
+
+class ImageConditionedSLatShape(ImageConditionedMixin, SLatShape):
+    """
+    Image conditioned structured latent for shape generation
+    """
+    pass

trellis2/datasets/structured_latent_svpbr.py

@@ -0,0 +1,290 @@
+import os
+os.environ['OPENCV_IO_ENABLE_OPENEXR'] = '1'
+import json
+from typing import *
+import numpy as np
+import torch
+import cv2
+from .. import models
+from .components import StandardDatasetBase, ImageConditionedMixin
+from ..modules.sparse import SparseTensor, sparse_cat
+from ..representations import MeshWithVoxel
+from ..renderers import PbrMeshRenderer, EnvMap
+from ..utils.data_utils import load_balanced_group_indices
+from ..utils.render_utils import yaw_pitch_r_fov_to_extrinsics_intrinsics
+
+
+class SLatPbrVisMixin:
+    def __init__(
+        self,
+        *args,
+        pretrained_pbr_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16',
+        pbr_slat_dec_path: Optional[str] = None,
+        pbr_slat_dec_ckpt: Optional[str] = None,
+        pretrained_shape_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
+        shape_slat_dec_path: Optional[str] = None,
+        shape_slat_dec_ckpt: Optional[str] = None,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.pbr_slat_dec = None
+        self.pretrained_pbr_slat_dec = pretrained_pbr_slat_dec
+        self.pbr_slat_dec_path = pbr_slat_dec_path
+        self.pbr_slat_dec_ckpt = pbr_slat_dec_ckpt
+        self.shape_slat_dec = None
+        self.pretrained_shape_slat_dec = pretrained_shape_slat_dec
+        self.shape_slat_dec_path = shape_slat_dec_path
+        self.shape_slat_dec_ckpt = shape_slat_dec_ckpt
+        
+    def _loading_slat_dec(self):
+        if self.pbr_slat_dec is not None and self.shape_slat_dec is not None:
+            return
+        if self.pbr_slat_dec_path is not None:
+            cfg = json.load(open(os.path.join(self.pbr_slat_dec_path, 'config.json'), 'r'))
+            decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
+            ckpt_path = os.path.join(self.pbr_slat_dec_path, 'ckpts', f'decoder_{self.pbr_slat_dec_ckpt}.pt')
+            decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+        else:
+            decoder = models.from_pretrained(self.pretrained_pbr_slat_dec)
+        self.pbr_slat_dec = decoder.cuda().eval()
+
+        if self.shape_slat_dec_path is not None:
+            cfg = json.load(open(os.path.join(self.shape_slat_dec_path, 'config.json'), 'r'))
+            decoder = getattr(models, cfg['models']['decoder']['name'])(**cfg['models']['decoder']['args'])
+            ckpt_path = os.path.join(self.shape_slat_dec_path, 'ckpts', f'decoder_{self.shape_slat_dec_ckpt}.pt')
+            decoder.load_state_dict(torch.load(ckpt_path, map_location='cpu', weights_only=True))
+        else:
+            decoder = models.from_pretrained(self.pretrained_shape_slat_dec)
+        decoder.set_resolution(self.resolution)
+        self.shape_slat_dec = decoder.cuda().eval()
+
+    def _delete_slat_dec(self):
+        del self.pbr_slat_dec
+        self.pbr_slat_dec = None
+        del self.shape_slat_dec
+        self.shape_slat_dec = None
+        
+    @torch.no_grad()
+    def decode_latent(self, z, shape_z, batch_size=4):
+        self._loading_slat_dec()
+        reps = []
+        if self.shape_slat_normalization is not None:
+            shape_z = shape_z * self.shape_slat_std.to(z.device) + self.shape_slat_mean.to(z.device)
+        if self.pbr_slat_normalization is not None:
+            z = z * self.pbr_slat_std.to(z.device) + self.pbr_slat_mean.to(z.device)
+        for i in range(0, z.shape[0], batch_size):
+            mesh, subs = self.shape_slat_dec(shape_z[i:i+batch_size], return_subs=True)
+            vox = self.pbr_slat_dec(z[i:i+batch_size], guide_subs=subs) * 0.5 + 0.5
+            reps.extend([
+                MeshWithVoxel(
+                    m.vertices, m.faces,
+                    origin = [-0.5, -0.5, -0.5],
+                    voxel_size = 1 / self.resolution,
+                    coords = v.coords[:, 1:],
+                    attrs = v.feats,
+                    voxel_shape = torch.Size([*v.shape, *v.spatial_shape]),
+                    layout = self.layout,
+                )
+                for m, v in zip(mesh, vox)
+            ])
+        self._delete_slat_dec()
+        return reps
+    
+    @torch.no_grad()
+    def visualize_sample(self, sample: dict):
+        shape_z = sample['concat_cond'].cuda()
+        z = sample['x_0'].cuda()
+        reps = self.decode_latent(z, shape_z)
+        
+        # build camera
+        yaw = [0, np.pi/2, np.pi, 3*np.pi/2]
+        yaw_offset = -16 / 180 * np.pi
+        yaw = [y + yaw_offset for y in yaw]
+        pitch = [20 / 180 * np.pi for _ in range(4)]
+        exts, ints = yaw_pitch_r_fov_to_extrinsics_intrinsics(yaw, pitch, 2, 30)
+        
+        # render
+        renderer = PbrMeshRenderer()
+        renderer.rendering_options.resolution = 512
+        renderer.rendering_options.near = 1
+        renderer.rendering_options.far = 100
+        renderer.rendering_options.ssaa = 2
+        renderer.rendering_options.peel_layers = 8
+        envmap = EnvMap(torch.tensor(
+            cv2.cvtColor(cv2.imread('assets/hdri/forest.exr', cv2.IMREAD_UNCHANGED), cv2.COLOR_BGR2RGB),
+            dtype=torch.float32, device='cuda'
+        ))
+        
+        images = {}
+        for representation in reps:
+            image = {}
+            tile = [2, 2]
+            for j, (ext, intr) in enumerate(zip(exts, ints)):
+                res = renderer.render(representation, ext, intr, envmap=envmap)
+                for k, v in res.items():
+                    if k not in images:
+                        images[k] = []
+                    if k not in image:
+                        image[k] = torch.zeros(3, 1024, 1024).cuda()  
+                    image[k][:, 512 * (j // tile[1]):512 * (j // tile[1] + 1), 512 * (j % tile[1]):512 * (j % tile[1] + 1)] = v
+            for k in images.keys():
+                images[k].append(image[k])
+        for k in images.keys():
+            images[k] = torch.stack(images[k], dim=0)
+        return images
+    
+    
+class SLatPbr(SLatPbrVisMixin, StandardDatasetBase):
+    """
+    structured latent for sparse voxel pbr dataset
+    
+    Args:
+        roots (str): path to the dataset
+        latent_key (str): key of the latent to be used
+        min_aesthetic_score (float): minimum aesthetic score
+        normalization (dict): normalization stats
+        resolution (int): resolution of decoded sparse voxel
+        attrs (list): attributes to be decoded
+        pretained_slat_dec (str): name of the pretrained slat decoder
+        slat_dec_path (str): path to the slat decoder, if given, will override the pretrained_slat_dec
+        slat_dec_ckpt (str): name of the slat decoder checkpoint
+    """
+    def __init__(self,
+        roots: str,
+        *,
+        resolution: int,
+        min_aesthetic_score: float = 5.0,
+        max_tokens: int = 32768,
+        full_pbr: bool = False,
+        pbr_slat_normalization: Optional[dict] = None,
+        shape_slat_normalization: Optional[dict] = None,
+        attrs: list[str] = ['base_color', 'metallic', 'roughness', 'emissive', 'alpha'],
+        pretrained_pbr_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/tex_dec_next_dc_f16c32_fp16',
+        pbr_slat_dec_path: Optional[str] = None,
+        pbr_slat_dec_ckpt: Optional[str] = None,
+        pretrained_shape_slat_dec: str = 'JeffreyXiang/TRELLIS.2-4B/ckpts/shape_dec_next_dc_f16c32_fp16',
+        shape_slat_dec_path: Optional[str] = None,
+        shape_slat_dec_ckpt: Optional[str] = None,
+        **kwargs
+    ):  
+        self.resolution = resolution
+        self.pbr_slat_normalization = pbr_slat_normalization
+        self.shape_slat_normalization = shape_slat_normalization
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_tokens = max_tokens
+        self.full_pbr = full_pbr
+        self.value_range = (0, 1)
+        
+        super().__init__(
+            roots,
+            pretrained_pbr_slat_dec=pretrained_pbr_slat_dec,
+            pbr_slat_dec_path=pbr_slat_dec_path,
+            pbr_slat_dec_ckpt=pbr_slat_dec_ckpt,
+            pretrained_shape_slat_dec=pretrained_shape_slat_dec,
+            shape_slat_dec_path=shape_slat_dec_path,
+            shape_slat_dec_ckpt=shape_slat_dec_ckpt,
+            **kwargs
+        )
+        
+        self.loads = [self.metadata.loc[sha256, 'pbr_latent_tokens'] for _, sha256 in self.instances]
+        
+        if self.pbr_slat_normalization is not None:
+            self.pbr_slat_mean = torch.tensor(self.pbr_slat_normalization['mean']).reshape(1, -1)
+            self.pbr_slat_std = torch.tensor(self.pbr_slat_normalization['std']).reshape(1, -1)
+        
+        if self.shape_slat_normalization is not None:
+            self.shape_slat_mean = torch.tensor(self.shape_slat_normalization['mean']).reshape(1, -1)
+            self.shape_slat_std = torch.tensor(self.shape_slat_normalization['std']).reshape(1, -1)
+        
+        self.attrs = attrs
+        self.channels = {
+            'base_color': 3,
+            'metallic': 1,
+            'roughness': 1,
+            'emissive': 3,
+            'alpha': 1,
+        }
+        self.layout = {}
+        start = 0
+        for attr in attrs:
+            self.layout[attr] = slice(start, start + self.channels[attr])
+            start += self.channels[attr]
+            
+    def filter_metadata(self, metadata):
+        stats = {}
+        metadata = metadata[metadata['pbr_latent_encoded'] == True]
+        stats['With PBR latent'] = len(metadata)
+        metadata = metadata[metadata['shape_latent_encoded'] == True]
+        stats['With shape latent'] = len(metadata)
+        metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+        stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        metadata = metadata[metadata['pbr_latent_tokens'] <= self.max_tokens]
+        stats[f'Num tokens <= {self.max_tokens}'] = len(metadata)
+        if self.full_pbr:
+            metadata = metadata[metadata['num_basecolor_tex'] > 0]
+            metadata = metadata[metadata['num_metallic_tex'] > 0]
+            metadata = metadata[metadata['num_roughness_tex'] > 0]
+            stats['Full PBR'] = len(metadata)
+        return metadata, stats
+    
+    def get_instance(self, root, instance):
+        # PBR latent
+        data = np.load(os.path.join(root['pbr_latent'], f'{instance}.npz'))
+        coords = torch.tensor(data['coords']).int()
+        coords = torch.cat([torch.zeros_like(coords)[:, :1], coords], dim=1)
+        feats = torch.tensor(data['feats']).float()
+        if self.pbr_slat_normalization is not None:
+            feats = (feats - self.pbr_slat_mean) / self.pbr_slat_std
+        pbr_z = SparseTensor(feats, coords)
+        
+        # Shape latent
+        data = np.load(os.path.join(root['shape_latent'], f'{instance}.npz'))
+        coords = torch.tensor(data['coords']).int()
+        coords = torch.cat([torch.zeros_like(coords)[:, :1], coords], dim=1)
+        feats = torch.tensor(data['feats']).float()
+        if self.shape_slat_normalization is not None:
+            feats = (feats - self.shape_slat_mean) / self.shape_slat_std
+        shape_z = SparseTensor(feats, coords)
+        
+        assert torch.equal(shape_z.coords, pbr_z.coords), \
+            f"Shape latent and PBR latent have different coordinates: {shape_z.coords.shape} vs {pbr_z.coords.shape}"
+            
+        return {
+            'x_0': pbr_z,
+            'concat_cond': shape_z,
+        }
+        
+    @staticmethod
+    def collate_fn(batch, split_size=None):
+        if split_size is None:
+            group_idx = [list(range(len(batch)))]
+        else:
+            group_idx = load_balanced_group_indices([b['x_0'].feats.shape[0] for b in batch], split_size)
+        packs = []
+        for group in group_idx:
+            sub_batch = [batch[i] for i in group]
+            pack = {}
+
+            keys = [k for k in sub_batch[0].keys()]
+            for k in keys:
+                if isinstance(sub_batch[0][k], torch.Tensor):
+                    pack[k] = torch.stack([b[k] for b in sub_batch])
+                elif isinstance(sub_batch[0][k], SparseTensor):
+                    pack[k] = sparse_cat([b[k] for b in sub_batch], dim=0)
+                elif isinstance(sub_batch[0][k], list):
+                    pack[k] = sum([b[k] for b in sub_batch], [])
+                else:
+                    pack[k] = [b[k] for b in sub_batch]
+            
+            packs.append(pack)
+        
+        if split_size is None:
+            return packs[0]
+        return packs
+
+
+class ImageConditionedSLatPbr(ImageConditionedMixin, SLatPbr):
+    """
+    Image conditioned structured latent dataset
+    """
+    pass

trellis2/models/__init__.py

@@ -0,0 +1,78 @@
+import importlib
+
+__attributes = {
+    # Sparse Structure
+    'SparseStructureEncoder': 'sparse_structure_vae',
+    'SparseStructureDecoder': 'sparse_structure_vae',
+    'SparseStructureFlowModel': 'sparse_structure_flow',
+    
+    # SLat Generation
+    'SLatFlowModel': 'structured_latent_flow',
+    'ElasticSLatFlowModel': 'structured_latent_flow',
+    
+    # SC-VAEs
+    'SparseUnetVaeEncoder': 'sc_vaes.sparse_unet_vae',
+    'SparseUnetVaeDecoder': 'sc_vaes.sparse_unet_vae',
+    'FlexiDualGridVaeEncoder': 'sc_vaes.fdg_vae',
+    'FlexiDualGridVaeDecoder': 'sc_vaes.fdg_vae'
+}
+
+__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), strict=False)
+
+    return model
+
+
+# For Pylance
+if __name__ == '__main__':
+    from .sparse_structure_vae import SparseStructureEncoder, SparseStructureDecoder
+    from .sparse_structure_flow import SparseStructureFlowModel
+    from .structured_latent_flow import SLatFlowModel, ElasticSLatFlowModel
+        
+    from .sc_vaes.sparse_unet_vae import SparseUnetVaeEncoder, SparseUnetVaeDecoder
+    from .sc_vaes.fdg_vae import FlexiDualGridVaeEncoder, FlexiDualGridVaeDecoder

trellis2/models/sc_vaes/fdg_vae.py

@@ -0,0 +1,110 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from ...modules import sparse as sp
+from .sparse_unet_vae import (
+    SparseResBlock3d,
+    SparseConvNeXtBlock3d,
+    
+    SparseResBlockDownsample3d,
+    SparseResBlockUpsample3d,
+    SparseResBlockS2C3d,
+    SparseResBlockC2S3d,
+)
+from .sparse_unet_vae import (
+    SparseUnetVaeEncoder,
+    SparseUnetVaeDecoder,
+)
+from ...representations import Mesh
+from o_voxel.convert import flexible_dual_grid_to_mesh
+
+
+class FlexiDualGridVaeEncoder(SparseUnetVaeEncoder):
+    def __init__(
+        self,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        down_block_type: List[str],
+        block_args: List[Dict[str, Any]],
+        use_fp16: bool = False,
+    ):
+        super().__init__(
+            6,
+            model_channels,
+            latent_channels,
+            num_blocks,
+            block_type,
+            down_block_type,
+            block_args,
+            use_fp16,
+        )
+        
+    def forward(self, vertices: sp.SparseTensor, intersected: sp.SparseTensor, sample_posterior=False, return_raw=False):
+        x = vertices.replace(torch.cat([
+            vertices.feats - 0.5,
+            intersected.feats.float() - 0.5,
+        ], dim=1))
+        return super().forward(x, sample_posterior, return_raw)
+    
+    
+class FlexiDualGridVaeDecoder(SparseUnetVaeDecoder):
+    def __init__(
+        self,
+        resolution: int,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        up_block_type: List[str],
+        block_args: List[Dict[str, Any]],
+        voxel_margin: float = 0.5,
+        use_fp16: bool = False,
+    ):
+        self.resolution = resolution
+        self.voxel_margin = voxel_margin
+        
+        super().__init__(
+            7,
+            model_channels,
+            latent_channels,
+            num_blocks,
+            block_type,
+            up_block_type,
+            block_args,
+            use_fp16,
+        )
+
+    def set_resolution(self, resolution: int) -> None:
+        self.resolution = resolution
+        
+    def forward(self, x: sp.SparseTensor, gt_intersected: sp.SparseTensor = None, **kwargs):
+        decoded = super().forward(x, **kwargs)
+        if self.training:
+            h, subs_gt, subs = decoded
+            vertices = h.replace((1 + 2 * self.voxel_margin) * F.sigmoid(h.feats[..., 0:3]) - self.voxel_margin)
+            intersected_logits = h.replace(h.feats[..., 3:6])
+            quad_lerp = h.replace(F.softplus(h.feats[..., 6:7]))
+            mesh = [Mesh(*flexible_dual_grid_to_mesh(
+                v.coords[:, 1:], v.feats, i.feats, q.feats,
+                aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+                grid_size=self.resolution,
+                train=True
+            )) for v, i, q in zip(vertices, gt_intersected, quad_lerp)]
+            return mesh, vertices, intersected_logits, subs_gt, subs
+        else:
+            out_list = list(decoded) if isinstance(decoded, tuple) else [decoded]
+            h = out_list[0]
+            vertices = h.replace((1 + 2 * self.voxel_margin) * F.sigmoid(h.feats[..., 0:3]) - self.voxel_margin)
+            intersected = h.replace(h.feats[..., 3:6] > 0)
+            quad_lerp = h.replace(F.softplus(h.feats[..., 6:7]))
+            mesh = [Mesh(*flexible_dual_grid_to_mesh(
+                v.coords[:, 1:], v.feats, i.feats, q.feats,
+                aabb=[[-0.5, -0.5, -0.5], [0.5, 0.5, 0.5]],
+                grid_size=self.resolution,
+                train=False
+            )) for v, i, q in zip(vertices, intersected, quad_lerp)]
+            out_list[0] = mesh
+            return out_list[0] if len(out_list) == 1 else tuple(out_list)

trellis2/models/sc_vaes/sparse_unet_vae.py

@@ -0,0 +1,522 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.checkpoint
+from ...modules.utils import convert_module_to_f16, convert_module_to_f32, zero_module
+from ...modules import sparse as sp
+from ...modules.norm import LayerNorm32
+
+
+class SparseResBlock3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        downsample: bool = False,
+        upsample: bool = False,
+        resample_mode: Literal['nearest', 'spatial2channel'] = 'nearest',
+        use_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.downsample = downsample
+        self.upsample = upsample
+        self.resample_mode = resample_mode
+        self.use_checkpoint = use_checkpoint
+        
+        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)
+        if resample_mode == 'nearest':
+            self.conv1 = sp.SparseConv3d(channels, self.out_channels, 3)
+        elif resample_mode =='spatial2channel' and not self.downsample:
+            self.conv1 = sp.SparseConv3d(channels, self.out_channels * 8, 3)
+        elif resample_mode =='spatial2channel' and self.downsample:
+            self.conv1 = sp.SparseConv3d(channels, self.out_channels // 8, 3)
+        self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
+        if resample_mode == 'nearest':
+            self.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
+        elif resample_mode =='spatial2channel' and self.downsample:
+            self.skip_connection = lambda x: x.replace(x.feats.reshape(x.feats.shape[0], out_channels, channels * 8 // out_channels).mean(dim=-1))
+        elif resample_mode =='spatial2channel' and not self.downsample:
+            self.skip_connection = lambda x: x.replace(x.feats.repeat_interleave(out_channels // (channels // 8), dim=1))
+        self.updown = None
+        if self.downsample:
+            if resample_mode == 'nearest':
+                self.updown = sp.SparseDownsample(2)
+            elif resample_mode =='spatial2channel':
+                self.updown = sp.SparseSpatial2Channel(2)
+        elif self.upsample:
+            self.to_subdiv = sp.SparseLinear(channels, 8)
+            if resample_mode == 'nearest':
+                self.updown = sp.SparseUpsample(2)
+            elif resample_mode =='spatial2channel':
+                self.updown = sp.SparseChannel2Spatial(2)
+
+    def _updown(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
+        if self.downsample:
+            x = self.updown(x)
+        elif self.upsample:
+            x = self.updown(x, subdiv.replace(subdiv.feats > 0))
+        return x
+
+    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        subdiv = None
+        if self.upsample:
+            subdiv = self.to_subdiv(x)
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        if self.resample_mode == 'spatial2channel':
+            h = self.conv1(h)
+        h = self._updown(h, subdiv)
+        x = self._updown(x, subdiv)
+        if self.resample_mode == 'nearest':
+            h = self.conv1(h)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        if self.upsample:
+            return h, subdiv
+        return h
+    
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseResBlockDownsample3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        use_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_checkpoint = use_checkpoint
+        
+        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.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
+        self.updown = sp.SparseDownsample(2)
+
+    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.updown(h)
+        x = self.updown(x)
+        h = self.conv1(h)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        return h
+    
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseResBlockUpsample3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        use_checkpoint: bool = False,
+        pred_subdiv: bool = True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_checkpoint = use_checkpoint
+        self.pred_subdiv = pred_subdiv
+        
+        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.skip_connection = sp.SparseLinear(channels, self.out_channels) if channels != self.out_channels else nn.Identity()
+        if self.pred_subdiv:
+            self.to_subdiv = sp.SparseLinear(channels, 8)
+        self.updown = sp.SparseUpsample(2)
+
+    def _forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
+        if self.pred_subdiv:
+            subdiv = self.to_subdiv(x)
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        subdiv_binarized = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None
+        h = self.updown(h, subdiv_binarized)
+        x = self.updown(x, subdiv_binarized)
+        h = self.conv1(h)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        if self.pred_subdiv:
+            return h, subdiv
+        else:
+            return h
+    
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseResBlockS2C3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        use_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_checkpoint = use_checkpoint
+        
+        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 // 8, 3)
+        self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
+        self.skip_connection = lambda x: x.replace(x.feats.reshape(x.feats.shape[0], out_channels, channels * 8 // out_channels).mean(dim=-1))
+        self.updown = sp.SparseSpatial2Channel(2)
+
+    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv1(h)
+        h = self.updown(h)
+        x = self.updown(x)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        return h
+    
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseResBlockC2S3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        out_channels: Optional[int] = None,
+        use_checkpoint: bool = False,
+        pred_subdiv: bool = True,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.out_channels = out_channels or channels
+        self.use_checkpoint = use_checkpoint
+        self.pred_subdiv = pred_subdiv
+        
+        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 * 8, 3)
+        self.conv2 = zero_module(sp.SparseConv3d(self.out_channels, self.out_channels, 3))
+        self.skip_connection = lambda x: x.replace(x.feats.repeat_interleave(out_channels // (channels // 8), dim=1))
+        if pred_subdiv:
+            self.to_subdiv = sp.SparseLinear(channels, 8)
+        self.updown = sp.SparseChannel2Spatial(2)
+
+    def _forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
+        if self.pred_subdiv:
+            subdiv = self.to_subdiv(x)
+        h = x.replace(self.norm1(x.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv1(h)
+        subdiv_binarized = subdiv.replace(subdiv.feats > 0) if subdiv is not None else None
+        h = self.updown(h, subdiv_binarized)
+        x = self.updown(x, subdiv_binarized)
+        h = h.replace(self.norm2(h.feats))
+        h = h.replace(F.silu(h.feats))
+        h = self.conv2(h)
+        h = h + self.skip_connection(x)
+        if self.pred_subdiv:
+            return h, subdiv
+        else:
+            return h
+    
+    def forward(self, x: sp.SparseTensor, subdiv: sp.SparseTensor = None) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, subdiv, use_reentrant=False)
+        else:
+            return self._forward(x, subdiv)
+        
+    
+class SparseConvNeXtBlock3d(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        mlp_ratio: float = 4.0,
+        use_checkpoint: bool = False,
+    ):
+        super().__init__()
+        self.channels = channels
+        self.use_checkpoint = use_checkpoint
+        
+        self.norm = LayerNorm32(channels, elementwise_affine=True, eps=1e-6)
+        self.conv = sp.SparseConv3d(channels, channels, 3)
+        self.mlp = nn.Sequential(
+            nn.Linear(channels, int(channels * mlp_ratio)),
+            nn.SiLU(),
+            zero_module(nn.Linear(int(channels * mlp_ratio), channels)),
+        )
+
+    def _forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        h = self.conv(x)
+        h = h.replace(self.norm(h.feats))
+        h = h.replace(self.mlp(h.feats))
+        return h + x
+    
+    def forward(self, x: sp.SparseTensor) -> sp.SparseTensor:
+        if self.use_checkpoint:
+            return torch.utils.checkpoint.checkpoint(self._forward, x, use_reentrant=False)
+        else:
+            return self._forward(x)
+
+
+class SparseUnetVaeEncoder(nn.Module):
+    """
+    Sparse Swin Transformer Unet VAE model.
+    """
+    def __init__(
+        self,
+        in_channels: int,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        down_block_type: List[str],
+        block_args: List[Dict[str, Any]],
+        use_fp16: bool = False,
+    ):
+        super().__init__()
+        self.in_channels = in_channels
+        self.model_channels = model_channels
+        self.num_blocks = num_blocks
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+
+        self.input_layer = sp.SparseLinear(in_channels, model_channels[0])
+        self.to_latent = sp.SparseLinear(model_channels[-1], 2 * latent_channels)
+        
+        self.blocks = nn.ModuleList([])
+        for i in range(len(num_blocks)):
+            self.blocks.append(nn.ModuleList([]))
+            for j in range(num_blocks[i]):
+                self.blocks[-1].append(
+                    globals()[block_type[i]](
+                        model_channels[i],
+                        **block_args[i],
+                    )
+                )
+            if i < len(num_blocks) - 1:
+                self.blocks[-1].append(
+                    globals()[down_block_type[i]](
+                        model_channels[i],
+                        model_channels[i+1],
+                        **block_args[i],
+                    )
+                )
+                
+        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)
+
+    def forward(self, x: sp.SparseTensor, sample_posterior=False, return_raw=False):
+        h = self.input_layer(x)
+        h = h.type(self.dtype)
+        for i, res in enumerate(self.blocks):
+            for j, block in enumerate(res):
+                h = block(h)
+        h = h.type(x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.to_latent(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
+    
+    
+class SparseUnetVaeDecoder(nn.Module):
+    """
+    Sparse Swin Transformer Unet VAE model.
+    """
+    def __init__(
+        self,
+        out_channels: int,
+        model_channels: List[int],
+        latent_channels: int,
+        num_blocks: List[int],
+        block_type: List[str],
+        up_block_type: List[str],
+        block_args: List[Dict[str, Any]],
+        use_fp16: bool = False,
+        pred_subdiv: bool = True,
+    ):
+        super().__init__()
+        self.out_channels = out_channels
+        self.model_channels = model_channels
+        self.num_blocks = num_blocks
+        self.use_fp16 = use_fp16
+        self.pred_subdiv = pred_subdiv
+        self.dtype = torch.float16 if use_fp16 else torch.float32
+        self.low_vram = False
+        
+        self.output_layer = sp.SparseLinear(model_channels[-1], out_channels)
+        self.from_latent = sp.SparseLinear(latent_channels, model_channels[0])
+        
+        self.blocks = nn.ModuleList([])
+        for i in range(len(num_blocks)):
+            self.blocks.append(nn.ModuleList([]))
+            for j in range(num_blocks[i]):
+                self.blocks[-1].append(
+                    globals()[block_type[i]](
+                        model_channels[i],
+                        **block_args[i],
+                    )
+                )
+            if i < len(num_blocks) - 1:
+                self.blocks[-1].append(
+                    globals()[up_block_type[i]](
+                        model_channels[i],
+                        model_channels[i+1],
+                        pred_subdiv=pred_subdiv,
+                        **block_args[i],
+                    )
+                )
+                    
+        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)
+
+    def forward(self, x: sp.SparseTensor, guide_subs: Optional[List[sp.SparseTensor]] = None, return_subs: bool = False) -> sp.SparseTensor:
+        assert guide_subs is None or self.pred_subdiv == False, "Only decoders with pred_subdiv=False can be used with guide_subs"
+        assert return_subs == False or self.pred_subdiv == True, "Only decoders with pred_subdiv=True can be used with return_subs"
+        
+        h = self.from_latent(x)
+        h = h.type(self.dtype)
+        subs_gt = []
+        subs = []
+        for i, res in enumerate(self.blocks):
+            for j, block in enumerate(res):
+                if i < len(self.blocks) - 1 and j == len(res) - 1:
+                    if self.pred_subdiv:
+                        if self.training:
+                            subs_gt.append(h.get_spatial_cache('subdivision'))
+                        h, sub = block(h)
+                        subs.append(sub)
+                    else:
+                        h = block(h, subdiv=guide_subs[i] if guide_subs is not None else None)
+                else:
+                    h = block(h)
+        h = h.type(x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.output_layer(h)
+        if self.training and self.pred_subdiv:
+            return h, subs_gt, subs
+        else:
+            if return_subs:
+                return h, subs
+            else:
+                return h
+    
+    def upsample(self, x: sp.SparseTensor, upsample_times: int) -> torch.Tensor:
+        assert self.pred_subdiv == True, "Only decoders with pred_subdiv=True can be used with upsampling"
+        
+        h = self.from_latent(x)
+        h = h.type(self.dtype)
+        for i, res in enumerate(self.blocks):
+            if i == upsample_times:
+                return h.coords
+            for j, block in enumerate(res):
+                if i < len(self.blocks) - 1 and j == len(res) - 1:
+                    h, sub = block(h)
+                else:
+                    h = block(h)
+       

trellis2/models/sparse_elastic_mixin.py

@@ -0,0 +1,24 @@
+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

trellis2/models/sparse_structure_flow.py

@@ -0,0 +1,247 @@
+from typing import *
+from functools import partial
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ..modules.utils import convert_module_to, manual_cast, str_to_dtype
+from ..modules.transformer import AbsolutePositionEmbedder, ModulatedTransformerCrossBlock
+from ..modules.attention import RotaryPositionEmbedder
+
+
+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,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        dtype: str = 'float32',
+        use_checkpoint: bool = False,
+        share_mod: bool = False,
+        initialization: str = 'vanilla',
+        qk_rms_norm: bool = False,
+        qk_rms_norm_cross: bool = False,
+        **kwargs
+    ):
+        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.pe_mode = pe_mode
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.initialization = initialization
+        self.qk_rms_norm = qk_rms_norm
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+        self.dtype = str_to_dtype(dtype)
+
+        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] * 3], indexing='ij')
+            coords = torch.stack(coords, dim=-1).reshape(-1, 3)
+            pos_emb = pos_embedder(coords)
+            self.register_buffer("pos_emb", pos_emb)
+        elif pe_mode == "rope":
+            pos_embedder = RotaryPositionEmbedder(self.model_channels // self.num_heads, 3)
+            coords = torch.meshgrid(*[torch.arange(res, device=self.device) for res in [resolution] * 3], indexing='ij')
+            coords = torch.stack(coords, dim=-1).reshape(-1, 3)
+            rope_phases = pos_embedder(coords)
+            self.register_buffer("rope_phases", rope_phases)
+            
+        if pe_mode != "rope":
+            self.rope_phases = None
+
+        self.input_layer = nn.Linear(in_channels, 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"),
+                rope_freq=rope_freq,
+                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)
+
+        self.initialize_weights()
+        self.convert_to(self.dtype)
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to(self, dtype: torch.dtype) -> None:
+        """
+        Convert the torso of the model to the specified dtype.
+        """
+        self.dtype = dtype
+        self.blocks.apply(partial(convert_module_to, dtype=dtype))
+
+    def initialize_weights(self) -> None:
+        if self.initialization == 'vanilla':
+            # 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)
+            
+        elif self.initialization == 'scaled':
+            # Initialize transformer layers:
+            def _basic_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.normal_(module.weight, std=np.sqrt(2.0 / (5.0 * self.model_channels)))
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            self.apply(_basic_init)
+            
+            # Scaled init for to_out and ffn2
+            def _scaled_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.normal_(module.weight, std=1.0 / np.sqrt(5 * self.num_blocks * self.model_channels))
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            for block in self.blocks:
+                block.self_attn.to_out.apply(_scaled_init)
+                block.cross_attn.to_out.apply(_scaled_init)
+                block.mlp.mlp[2].apply(_scaled_init)
+            
+            # Initialize input layer to make the initial representation have variance 1
+            nn.init.normal_(self.input_layer.weight, std=1.0 / np.sqrt(self.in_channels))
+            nn.init.zeros_(self.input_layer.bias)
+            
+            # 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 = x.view(*x.shape[:2], -1).permute(0, 2, 1).contiguous()
+
+        h = self.input_layer(h)
+        if self.pe_mode == "ape":
+            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 = manual_cast(t_emb, self.dtype)
+        h = manual_cast(h, self.dtype)
+        cond = manual_cast(cond, self.dtype)
+        for block in self.blocks:
+            h = block(h, t_emb, cond, self.rope_phases)
+        h = manual_cast(h, 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] * 3).contiguous()
+
+        return h

trellis2/models/sparse_structure_vae.py

@@ -0,0 +1,306 @@
+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

trellis2/models/structured_latent_flow.py

@@ -0,0 +1,207 @@
+from typing import *
+from functools import partial
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from ..modules.utils import convert_module_to, manual_cast, str_to_dtype
+from ..modules.transformer import AbsolutePositionEmbedder
+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 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,
+        pe_mode: Literal["ape", "rope"] = "ape",
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        dtype: str = 'float32',
+        use_checkpoint: bool = False,
+        share_mod: bool = False,
+        initialization: str = 'vanilla',
+        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.pe_mode = pe_mode
+        self.use_checkpoint = use_checkpoint
+        self.share_mod = share_mod
+        self.initialization = initialization
+        self.qk_rms_norm = qk_rms_norm
+        self.qk_rms_norm_cross = qk_rms_norm_cross
+        self.dtype = str_to_dtype(dtype)
+
+        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)
+            
+        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"),
+                rope_freq=rope_freq,
+                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_layer = sp.SparseLinear(model_channels, out_channels)
+
+        self.initialize_weights()
+        self.convert_to(self.dtype)
+
+    @property
+    def device(self) -> torch.device:
+        """
+        Return the device of the model.
+        """
+        return next(self.parameters()).device
+
+    def convert_to(self, dtype: torch.dtype) -> None:
+        """
+        Convert the torso of the model to the specified dtype.
+        """
+        self.dtype = dtype
+        self.blocks.apply(partial(convert_module_to, dtype=dtype))
+
+    def initialize_weights(self) -> None:
+        if self.initialization == 'vanilla':
+            # 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)
+            
+        elif self.initialization == 'scaled':
+            # Initialize transformer layers:
+            def _basic_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.normal_(module.weight, std=np.sqrt(2.0 / (5.0 * self.model_channels)))
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            self.apply(_basic_init)
+            
+            # Scaled init for to_out and ffn2
+            def _scaled_init(module):
+                if isinstance(module, nn.Linear):
+                    torch.nn.init.normal_(module.weight, std=1.0 / np.sqrt(5 * self.num_blocks * self.model_channels))
+                    if module.bias is not None:
+                        nn.init.constant_(module.bias, 0)
+            for block in self.blocks:
+                block.self_attn.to_out.apply(_scaled_init)
+                block.cross_attn.to_out.apply(_scaled_init)
+                block.mlp.mlp[2].apply(_scaled_init)
+            
+            # Initialize input layer to make the initial representation have variance 1
+            nn.init.normal_(self.input_layer.weight, std=1.0 / np.sqrt(self.in_channels))
+            nn.init.zeros_(self.input_layer.bias)
+            
+            # 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: Union[torch.Tensor, List[torch.Tensor]],
+        concat_cond: Optional[sp.SparseTensor] = None,
+        **kwargs
+    ) -> sp.SparseTensor:
+        if concat_cond is not None:
+            x = sp.sparse_cat([x, concat_cond], dim=-1)
+        if isinstance(cond, list):
+            cond = sp.VarLenTensor.from_tensor_list(cond)
+
+        h = self.input_layer(x)
+        h = manual_cast(h, self.dtype)
+        t_emb = self.t_embedder(t)
+        if self.share_mod:
+            t_emb = self.adaLN_modulation(t_emb)
+        t_emb = manual_cast(t_emb, self.dtype)
+        cond = manual_cast(cond, self.dtype)
+
+        if self.pe_mode == "ape":
+            pe = self.pos_embedder(h.coords[:, 1:])
+            h = h + manual_cast(pe, self.dtype)
+        for block in self.blocks:
+            h = block(h, t_emb, cond)
+
+        h = manual_cast(h, x.dtype)
+        h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
+        h = self.out_layer(h)
+        return h
+
+
+class ElasticSLatFlowModel(SparseTransformerElasticMixin, SLatFlowModel):
+    """
+    SLat Flow Model with elastic memory management.
+    Used for training with low VRAM.
+    """
+    pass

trellis2/modules/attention/__init__.py

@@ -0,0 +1,3 @@
+from .full_attn import *
+from .modules import *
+from .rope import *

trellis2/modules/attention/config.py

@@ -0,0 +1,34 @@
+from typing import *
+import sys
+
+# Default to 'sdpa' (PyTorch's built-in) on Windows since flash_attn isn't available
+BACKEND = 'sdpa' if sys.platform == 'win32' else 'flash_attn'
+DEBUG = False
+
+def __from_env():
+    import os
+
+    global BACKEND
+    global DEBUG
+
+    env_attn_backend = os.environ.get('ATTN_BACKEND')
+    env_attn_debug = os.environ.get('ATTN_DEBUG')
+
+    if env_attn_backend is not None and env_attn_backend in ['xformers', 'flash_attn', 'flash_attn_3', 'sdpa', 'naive']:
+        BACKEND = env_attn_backend
+    if env_attn_debug is not None:
+        DEBUG = env_attn_debug == '1'
+
+    print(f"[ATTENTION] Using backend: {BACKEND}")
+        
+
+__from_env()
+    
+
+def set_backend(backend: Literal['xformers', 'flash_attn']):
+    global BACKEND
+    BACKEND = backend
+
+def set_debug(debug: bool):
+    global DEBUG
+    DEBUG = debug

trellis2/modules/attention/full_attn.py

@@ -0,0 +1,145 @@
+from typing import *
+import torch
+import math
+from . import config
+
+
+__all__ = [
+    'scaled_dot_product_attention',
+]
+
+
+def _naive_sdpa(q, k, v):
+    """
+    Naive implementation of scaled dot product attention.
+    """
+    q = q.permute(0, 2, 1, 3)   # [N, H, L, C]
+    k = k.permute(0, 2, 1, 3)   # [N, H, L, C]
+    v = v.permute(0, 2, 1, 3)   # [N, H, L, C]
+    scale_factor = 1 / math.sqrt(q.size(-1))
+    attn_weight = q @ k.transpose(-2, -1) * scale_factor
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+    out = attn_weight @ v
+    out = out.permute(0, 2, 1, 3)   # [N, L, H, C]
+    return out
+
+
+@overload
+def scaled_dot_product_attention(qkv: torch.Tensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention.
+
+    Args:
+        qkv (torch.Tensor): A [N, L, 3, H, C] tensor containing Qs, Ks, and Vs.
+    """
+    ...
+
+@overload
+def scaled_dot_product_attention(q: torch.Tensor, kv: torch.Tensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention.
+
+    Args:
+        q (torch.Tensor): A [N, L, H, C] tensor containing Qs.
+        kv (torch.Tensor): A [N, L, 2, H, C] tensor containing Ks and Vs.
+    """
+    ...
+
+@overload
+def scaled_dot_product_attention(q: torch.Tensor, k: torch.Tensor, v: torch.Tensor) -> torch.Tensor:
+    """
+    Apply scaled dot product attention.
+
+    Args:
+        q (torch.Tensor): A [N, L, H, Ci] tensor containing Qs.
+        k (torch.Tensor): A [N, L, H, Ci] tensor containing Ks.
+        v (torch.Tensor): A [N, L, H, Co] tensor containing Vs.
+
+    Note:
+        k and v are assumed to have the same coordinate map.
+    """
+    ...
+
+def scaled_dot_product_attention(*args, **kwargs):
+    arg_names_dict = {
+        1: ['qkv'],
+        2: ['q', 'kv'],
+        3: ['q', 'k', 'v']
+    }
+    num_all_args = len(args) + len(kwargs)
+    assert num_all_args in arg_names_dict, f"Invalid number of arguments, got {num_all_args}, expected 1, 2, or 3"
+    for key in arg_names_dict[num_all_args][len(args):]:
+        assert key in kwargs, f"Missing argument {key}"
+
+    if num_all_args == 1:
+        qkv = args[0] if len(args) > 0 else kwargs['qkv']
+        assert len(qkv.shape) == 5 and qkv.shape[2] == 3, f"Invalid shape for qkv, got {qkv.shape}, expected [N, L, 3, H, C]"
+        device = qkv.device
+
+    elif num_all_args == 2:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        kv = args[1] if len(args) > 1 else kwargs['kv']
+        assert q.shape[0] == kv.shape[0], f"Batch size mismatch, got {q.shape[0]} and {kv.shape[0]}"
+        assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, C]"
+        assert len(kv.shape) == 5, f"Invalid shape for kv, got {kv.shape}, expected [N, L, 2, H, C]"
+        device = q.device
+
+    elif num_all_args == 3:
+        q = args[0] if len(args) > 0 else kwargs['q']
+        k = args[1] if len(args) > 1 else kwargs['k']
+        v = args[2] if len(args) > 2 else kwargs['v']
+        assert q.shape[0] == k.shape[0] == v.shape[0], f"Batch size mismatch, got {q.shape[0]}, {k.shape[0]}, and {v.shape[0]}"
+        assert len(q.shape) == 4, f"Invalid shape for q, got {q.shape}, expected [N, L, H, Ci]"
+        assert len(k.shape) == 4, f"Invalid shape for k, got {k.shape}, expected [N, L, H, Ci]"
+        assert len(v.shape) == 4, f"Invalid shape for v, got {v.shape}, expected [N, L, H, Co]"
+        device = q.device    
+
+    if config.BACKEND == 'xformers':
+        if 'xops' not in globals():
+            import xformers.ops as xops
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=2)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+        out = xops.memory_efficient_attention(q, k, v)
+    elif config.BACKEND == 'flash_attn':
+        if 'flash_attn' not in globals():
+            import flash_attn
+        if num_all_args == 1:
+            out = flash_attn.flash_attn_qkvpacked_func(qkv)
+        elif num_all_args == 2:
+            out = flash_attn.flash_attn_kvpacked_func(q, kv)
+        elif num_all_args == 3:
+            out = flash_attn.flash_attn_func(q, k, v)
+    elif config.BACKEND == 'flash_attn_3':
+        if 'flash_attn_3' not in globals():
+            import flash_attn_interface as flash_attn_3
+            if num_all_args == 1:
+                out = flash_attn_3.flash_attn_qkvpacked_func(qkv)
+            elif num_all_args == 2:
+                k, v = kv.unbind(dim=2)
+                out = flash_attn_3.flash_attn_func(q, k, v)
+            elif num_all_args == 3:
+                out = flash_attn_3.flash_attn_func(q, k, v)
+    elif config.BACKEND == 'sdpa':
+        if 'sdpa' not in globals():
+            from torch.nn.functional import scaled_dot_product_attention as sdpa
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=2)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+        q = q.permute(0, 2, 1, 3)   # [N, H, L, C]
+        k = k.permute(0, 2, 1, 3)   # [N, H, L, C]
+        v = v.permute(0, 2, 1, 3)   # [N, H, L, C]
+        out = sdpa(q, k, v)         # [N, H, L, C]
+        out = out.permute(0, 2, 1, 3)   # [N, L, H, C]
+    elif config.BACKEND == 'naive':
+        if num_all_args == 1:
+            q, k, v = qkv.unbind(dim=2)
+        elif num_all_args == 2:
+            k, v = kv.unbind(dim=2)
+        out = _naive_sdpa(q, k, v)
+    else:
+        raise ValueError(f"Unknown attention module: {config.BACKEND}")
+    
+    return out

trellis2/modules/attention/modules.py

@@ -0,0 +1,102 @@
+from typing import *
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from .full_attn import scaled_dot_product_attention
+from .rope import RotaryPositionEmbedder
+
+
+class MultiHeadRMSNorm(nn.Module):
+    def __init__(self, dim: int, heads: int):
+        super().__init__()
+        self.scale = dim ** 0.5
+        self.gamma = nn.Parameter(torch.ones(heads, dim))
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        return (F.normalize(x.float(), dim = -1) * self.gamma * self.scale).to(x.dtype)
+    
+
+class MultiHeadAttention(nn.Module):
+    def __init__(
+        self,
+        channels: int,
+        num_heads: int,
+        ctx_channels: Optional[int]=None,
+        type: Literal["self", "cross"] = "self",
+        attn_mode: Literal["full", "windowed"] = "full",
+        window_size: Optional[int] = None,
+        shift_window: Optional[Tuple[int, int, int]] = None,
+        qkv_bias: bool = True,
+        use_rope: bool = False,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0),
+        qk_rms_norm: bool = False,
+    ):
+        super().__init__()
+        assert channels % num_heads == 0
+        assert type in ["self", "cross"], f"Invalid attention type: {type}"
+        assert attn_mode in ["full", "windowed"], f"Invalid attention mode: {attn_mode}"
+        assert type == "self" or attn_mode == "full", "Cross-attention only supports full attention"
+        
+        if attn_mode == "windowed":
+            raise NotImplementedError("Windowed attention is not yet implemented")
+        
+        self.channels = channels
+        self.head_dim = channels // num_heads
+        self.ctx_channels = ctx_channels if ctx_channels is not None else channels
+        self.num_heads = num_heads
+        self._type = type
+        self.attn_mode = attn_mode
+        self.window_size = window_size
+        self.shift_window = shift_window
+        self.use_rope = use_rope
+        self.qk_rms_norm = qk_rms_norm
+
+        if self._type == "self":
+            self.to_qkv = nn.Linear(channels, channels * 3, bias=qkv_bias)
+        else:
+            self.to_q = nn.Linear(channels, channels, bias=qkv_bias)
+            self.to_kv = nn.Linear(self.ctx_channels, channels * 2, bias=qkv_bias)
+            
+        if self.qk_rms_norm:
+            self.q_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+            self.k_rms_norm = MultiHeadRMSNorm(self.head_dim, num_heads)
+            
+        self.to_out = nn.Linear(channels, channels)
+    
+    def forward(self, x: torch.Tensor, context: Optional[torch.Tensor] = None, phases: Optional[torch.Tensor] = None) -> torch.Tensor:
+        B, L, C = x.shape
+        if self._type == "self":
+            qkv = self.to_qkv(x)
+            qkv = qkv.reshape(B, L, 3, self.num_heads, -1)
+            
+            if self.attn_mode == "full":
+                if self.qk_rms_norm or self.use_rope:
+                    q, k, v = qkv.unbind(dim=2)
+                    if self.qk_rms_norm:
+                        q = self.q_rms_norm(q)
+                        k = self.k_rms_norm(k)
+                    if self.use_rope:
+                        assert phases is not None, "Phases must be provided for RoPE"
+                        q = RotaryPositionEmbedder.apply_rotary_embedding(q, phases)
+                        k = RotaryPositionEmbedder.apply_rotary_embedding(k, phases)
+                    h = scaled_dot_product_attention(q, k, v)
+                else:
+                    h = scaled_dot_product_attention(qkv)
+            elif self.attn_mode == "windowed":
+                raise NotImplementedError("Windowed attention is not yet implemented")
+        else:
+            Lkv = context.shape[1]
+            q = self.to_q(x)
+            kv = self.to_kv(context)
+            q = q.reshape(B, L, self.num_heads, -1)
+            kv = kv.reshape(B, Lkv, 2, self.num_heads, -1)
+            if self.qk_rms_norm:
+                q = self.q_rms_norm(q)
+                k, v = kv.unbind(dim=2)
+                k = self.k_rms_norm(k)
+                h = scaled_dot_product_attention(q, k, v)
+            else:
+                h = scaled_dot_product_attention(q, kv)
+        h = h.reshape(B, L, -1)
+        h = self.to_out(h)
+        return h

trellis2/modules/attention/rope.py

@@ -0,0 +1,48 @@
+from typing import *
+import torch
+import torch.nn as nn
+
+
+class RotaryPositionEmbedder(nn.Module):
+    def __init__(
+        self, 
+        head_dim: int,
+        dim: int = 3,
+        rope_freq: Tuple[float, float] = (1.0, 10000.0)
+    ):
+        super().__init__()
+        assert head_dim % 2 == 0, "Head dim must be divisible by 2"
+        self.head_dim = head_dim
+        self.dim = dim
+        self.rope_freq = rope_freq
+        self.freq_dim = head_dim // 2 // dim
+        self.freqs = torch.arange(self.freq_dim, dtype=torch.float32) / self.freq_dim
+        self.freqs = rope_freq[0] / (rope_freq[1] ** (self.freqs))
+        
+    def _get_phases(self, indices: torch.Tensor) -> torch.Tensor:
+        self.freqs = self.freqs.to(indices.device)
+        phases = torch.outer(indices, self.freqs)
+        phases = torch.polar(torch.ones_like(phases), phases)
+        return phases
+    
+    @staticmethod
+    def apply_rotary_embedding(x: torch.Tensor, phases: torch.Tensor) -> torch.Tensor:
+        x_complex = torch.view_as_complex(x.float().reshape(*x.shape[:-1], -1, 2))
+        x_rotated = x_complex * phases.unsqueeze(-2)
+        x_embed = torch.view_as_real(x_rotated).reshape(*x_rotated.shape[:-1], -1).to(x.dtype)
+        return x_embed
+        
+    def forward(self, indices: torch.Tensor) -> torch.Tensor:
+        """
+        Args:
+            indices (torch.Tensor): [..., N, C] tensor of spatial positions
+        """
+        assert indices.shape[-1] == self.dim, f"Last dim of indices must be {self.dim}"
+        phases = self._get_phases(indices.reshape(-1)).reshape(*indices.shape[:-1], -1)
+        if phases.shape[-1] < self.head_dim // 2:
+                padn = self.head_dim // 2 - phases.shape[-1]
+                phases = torch.cat([phases, torch.polar(
+                    torch.ones(*phases.shape[:-1], padn, device=phases.device),
+                    torch.zeros(*phases.shape[:-1], padn, device=phases.device)
+                )], dim=-1)
+        return phases

trellis2/modules/image_feature_extractor.py

@@ -0,0 +1,123 @@
+from typing import *
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+from transformers import DINOv3ViTModel
+import numpy as np
+from PIL import Image
+
+
+class DinoV2FeatureExtractor:
+    """
+    Feature extractor for DINOv2 models.
+    """
+    def __init__(self, model_name: str):
+        self.model_name = model_name
+        self.model = torch.hub.load('facebookresearch/dinov2', model_name, pretrained=True)
+        self.model.eval()
+        self.transform = transforms.Compose([
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+    def to(self, device):
+        self.model.to(device)
+
+    def cuda(self):
+        self.model.cuda()
+
+    def cpu(self):
+        self.model.cpu()
+    
+    @torch.no_grad()
+    def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
+        """
+        Extract features from the image.
+        
+        Args:
+            image: A batch of images as a tensor of shape (B, C, H, W) or a list of PIL images.
+        
+        Returns:
+            A tensor of shape (B, N, D) where N is the number of patches and D is the feature dimension.
+        """
+        if isinstance(image, torch.Tensor):
+            assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
+        elif isinstance(image, list):
+            assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
+            image = [i.resize((518, 518), Image.LANCZOS) for i in image]
+            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+            image = torch.stack(image).cuda()
+        else:
+            raise ValueError(f"Unsupported type of image: {type(image)}")
+        
+        image = self.transform(image).cuda()
+        features = self.model(image, is_training=True)['x_prenorm']
+        patchtokens = F.layer_norm(features, features.shape[-1:])
+        return patchtokens
+    
+
+class DinoV3FeatureExtractor:
+    """
+    Feature extractor for DINOv3 models.
+    """
+    def __init__(self, model_name: str, image_size=512):
+        self.model_name = model_name
+        # Try loading with local_files_only first (for cached gated models)
+        try:
+            self.model = DINOv3ViTModel.from_pretrained(model_name, local_files_only=True)
+        except Exception:
+            # Fall back to remote loading
+            self.model = DINOv3ViTModel.from_pretrained(model_name)
+        self.model.eval()
+        self.image_size = image_size
+        self.transform = transforms.Compose([
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+
+    def to(self, device):
+        self.model.to(device)
+
+    def cuda(self):
+        self.model.cuda()
+
+    def cpu(self):
+        self.model.cpu()
+
+    def extract_features(self, image: torch.Tensor) -> torch.Tensor:
+        image = image.to(self.model.embeddings.patch_embeddings.weight.dtype)
+        hidden_states = self.model.embeddings(image, bool_masked_pos=None)
+        position_embeddings = self.model.rope_embeddings(image)
+
+        for i, layer_module in enumerate(self.model.layer):
+            hidden_states = layer_module(
+                hidden_states,
+                position_embeddings=position_embeddings,
+            )
+
+        return F.layer_norm(hidden_states, hidden_states.shape[-1:])
+        
+    @torch.no_grad()
+    def __call__(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
+        """
+        Extract features from the image.
+        
+        Args:
+            image: A batch of images as a tensor of shape (B, C, H, W) or a list of PIL images.
+        
+        Returns:
+            A tensor of shape (B, N, D) where N is the number of patches and D is the feature dimension.
+        """
+        if isinstance(image, torch.Tensor):
+            assert image.ndim == 4, "Image tensor should be batched (B, C, H, W)"
+        elif isinstance(image, list):
+            assert all(isinstance(i, Image.Image) for i in image), "Image list should be list of PIL images"
+            image = [i.resize((self.image_size, self.image_size), Image.LANCZOS) for i in image]
+            image = [np.array(i.convert('RGB')).astype(np.float32) / 255 for i in image]
+            image = [torch.from_numpy(i).permute(2, 0, 1).float() for i in image]
+            image = torch.stack(image).cuda()
+        else:
+            raise ValueError(f"Unsupported type of image: {type(image)}")
+        
+        image = self.transform(image).cuda()
+        features = self.extract_features(image)
+        return features

trellis2/modules/norm.py

@@ -0,0 +1,32 @@
+import torch
+import torch.nn as nn
+from .utils import manual_cast
+
+
+class LayerNorm32(nn.LayerNorm):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_dtype = x.dtype
+        x = manual_cast(x, torch.float32)
+        o = super().forward(x)
+        return manual_cast(o, x_dtype)
+    
+
+class GroupNorm32(nn.GroupNorm):
+    """
+    A GroupNorm layer that converts to float32 before the forward pass.
+    """
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        x_dtype = x.dtype
+        x = manual_cast(x, torch.float32)
+        o = super().forward(x)
+        return manual_cast(o, x_dtype)
+    
+    
+class ChannelLayerNorm32(LayerNorm32):
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        DIM = x.dim()
+        x = x.permute(0, *range(2, DIM), 1).contiguous()
+        x = super().forward(x)
+        x = x.permute(0, DIM-1, *range(1, DIM-1)).contiguous()
+        return x
+    

trellis2/modules/sparse/__init__.py

@@ -0,0 +1,69 @@
+from . import config
+import importlib
+
+__attributes = {
+    'VarLenTensor': 'basic',
+    'varlen_cat': 'basic',
+    'varlen_unbind': 'basic',
+    'SparseTensor': 'basic',
+    'sparse_cat': 'basic',
+    'sparse_unbind': 'basic',
+    'SparseGroupNorm': 'norm',
+    'SparseLayerNorm': 'norm',
+    'SparseGroupNorm32': 'norm',
+    'SparseLayerNorm32': 'norm',
+    'SparseReLU': 'nonlinearity',
+    'SparseSiLU': 'nonlinearity',
+    'SparseGELU': 'nonlinearity',
+    'SparseActivation': 'nonlinearity',
+    'SparseLinear': 'linear',
+    'sparse_scaled_dot_product_attention': 'attention',
+    'SerializeMode': 'attention',
+    'sparse_serialized_scaled_dot_product_self_attention': 'attention',
+    'sparse_windowed_scaled_dot_product_self_attention': 'attention',
+    'sparse_windowed_scaled_dot_product_cross_attention': 'attention',
+    'SparseRotaryPositionEmbedder':  'attention',
+    'SparseMultiHeadAttention': 'attention',
+    'SparseConv3d': 'conv',
+    'SparseInverseConv3d': 'conv',
+    'SparseDownsample': 'spatial',
+    'SparseUpsample': 'spatial',
+    'SparseSubdivide': 'spatial',
+    'SparseSpatial2Channel': 'spatial',
+    'SparseChannel2Spatial': 'spatial',
+    'sparse_nearest_interpolate': 'spatial',
+    'sparse_trilinear_interpolate': 'spatial',
+    'encode_seq': 'serialize',
+    'decode_seq': 'serialize',
+}
+
+__submodules = ['transformer', 'conv']
+
+__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]
+
+
+# For Pylance
+if __name__ == '__main__':
+    from .basic import *
+    from .norm import *
+    from .nonlinearity import *
+    from .linear import *
+    from .attention import *
+    from .conv import *
+    from .spatial import *
+    from .serialize import *
+    import transformer
+    import conv

trellis2/modules/sparse/attention/__init__.py

@@ -0,0 +1,3 @@
+from .full_attn import *
+from .windowed_attn import *
+from .modules import *