Trellis update

trellis/datasets/__init__.py

@@ -0,0 +1,58 @@
+import importlib
+
+__attributes = {
+    'SparseStructure': 'sparse_structure',
+    
+    'SparseFeat2Render': 'sparse_feat2render',
+    'SLat2Render':'structured_latent2render',
+    'Slat2RenderGeo':'structured_latent2render',
+    
+    'SparseStructureLatent': 'sparse_structure_latent',
+    'TextConditionedSparseStructureLatent': 'sparse_structure_latent',
+    'ImageConditionedSparseStructureLatent': 'sparse_structure_latent',
+    
+    'SLat': 'structured_latent',
+    'TextConditionedSLat': 'structured_latent',
+    'ImageConditionedSLat': 'structured_latent',
+}
+
+__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 .sparse_structure import SparseStructure
+    
+    from .sparse_feat2render import SparseFeat2Render
+    from .structured_latent2render import (
+        SLat2Render,
+        Slat2RenderGeo,
+    )
+    
+    from .sparse_structure_latent import (
+        SparseStructureLatent,
+        TextConditionedSparseStructureLatent,
+        ImageConditionedSparseStructureLatent,
+    )
+    
+    from .structured_latent import (
+        SLat,
+        TextConditionedSLat,
+        ImageConditionedSLat,
+    )
+    

trellis/datasets/components.py

@@ -0,0 +1,137 @@
+from typing import *
+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__()
+        self.roots = roots.split(',')
+        self.instances = []
+        self.metadata = pd.DataFrame()
+        
+        self._stats = {}
+        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: str, 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(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 TextConditionedMixin:
+    def __init__(self, roots, **kwargs):
+        super().__init__(roots, **kwargs)
+        self.captions = {}
+        for instance in self.instances:
+            sha256 = instance[1]
+            self.captions[sha256] = json.loads(self.metadata.loc[sha256]['captions'])
+    
+    def filter_metadata(self, metadata):
+        metadata, stats = super().filter_metadata(metadata)
+        metadata = metadata[metadata['captions'].notna()]
+        stats['With captions'] = len(metadata)
+        return metadata, stats
+    
+    def get_instance(self, root, instance):
+        pack = super().get_instance(root, instance)
+        text = np.random.choice(self.captions[instance])
+        pack['cond'] = text
+        return pack
+    
+    
+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[f'cond_rendered']]
+        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, 'renders_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_size_ratio = 1.2
+        aug_hsize = hsize * aug_size_ratio
+        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
+    

trellis/datasets/sparse_feat2render.py

@@ -0,0 +1,134 @@
+import os
+from PIL import Image
+import json
+import numpy as np
+import pandas as pd
+import torch
+import utils3d.torch
+from ..modules.sparse.basic import SparseTensor
+from .components import StandardDatasetBase
+
+
+class SparseFeat2Render(StandardDatasetBase):
+    """
+    SparseFeat2Render dataset.
+    
+    Args:
+        roots (str): paths to the dataset
+        image_size (int): size of the image
+        model (str): model name
+        resolution (int): resolution of the data
+        min_aesthetic_score (float): minimum aesthetic score
+        max_num_voxels (int): maximum number of voxels
+    """
+    def __init__(
+        self,
+        roots: str,
+        image_size: int,
+        model: str = 'dinov2_vitl14_reg',
+        resolution: int = 64,
+        min_aesthetic_score: float = 5.0,
+        max_num_voxels: int = 32768,
+    ):
+        self.image_size = image_size
+        self.model = model
+        self.resolution = resolution
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_num_voxels = max_num_voxels
+        self.value_range = (0, 1)
+        
+        super().__init__(roots)
+        
+    def filter_metadata(self, metadata):
+        stats = {}
+        metadata = metadata[metadata[f'feature_{self.model}']]
+        stats['With features'] = len(metadata)
+        metadata = metadata[metadata['aesthetic_score'] >= self.min_aesthetic_score]
+        stats[f'Aesthetic score >= {self.min_aesthetic_score}'] = len(metadata)
+        metadata = metadata[metadata['num_voxels'] <= self.max_num_voxels]
+        stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata)
+        return metadata, stats
+
+    def _get_image(self, root, instance):
+        with open(os.path.join(root, 'renders', instance, 'transforms.json')) as f:
+            metadata = json.load(f)
+        n_views = len(metadata['frames'])
+        view = np.random.randint(n_views)
+        metadata = metadata['frames'][view]
+        fov = metadata['camera_angle_x']
+        intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov))
+        c2w = torch.tensor(metadata['transform_matrix'])
+        c2w[:3, 1:3] *= -1
+        extrinsics = torch.inverse(c2w)
+
+        image_path = os.path.join(root, 'renders', instance, metadata['file_path'])
+        image = Image.open(image_path)
+        alpha = image.getchannel(3)
+        image = image.convert('RGB')
+        image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+        alpha = alpha.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+        image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
+        alpha = torch.tensor(np.array(alpha)).float() / 255.0
+        
+        return {
+            'image': image,
+            'alpha': alpha,
+            'extrinsics': extrinsics,
+            'intrinsics': intrinsics,
+        }
+    
+    def _get_feat(self, root, instance):
+        DATA_RESOLUTION = 64
+        feats_path = os.path.join(root, 'features', self.model, f'{instance}.npz')
+        feats = np.load(feats_path, allow_pickle=True)
+        coords = torch.tensor(feats['indices']).int()
+        feats = torch.tensor(feats['patchtokens']).float()
+        
+        if self.resolution != DATA_RESOLUTION:
+            factor = DATA_RESOLUTION // self.resolution
+            coords = coords // factor
+            coords, idx = coords.unique(return_inverse=True, dim=0)
+            feats = torch.scatter_reduce(
+                torch.zeros(coords.shape[0], feats.shape[1], device=feats.device),
+                dim=0,
+                index=idx.unsqueeze(-1).expand(-1, feats.shape[1]),
+                src=feats,
+                reduce='mean'
+            )
+        
+        return {
+            'coords': coords,
+            'feats': feats,
+        }
+
+    @torch.no_grad()
+    def visualize_sample(self, sample: dict):
+        return sample['image']
+
+    @staticmethod
+    def collate_fn(batch):
+        pack = {}
+        coords = []
+        for i, b in enumerate(batch):
+            coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
+        coords = torch.cat(coords)
+        feats = torch.cat([b['feats'] for b in batch])
+        pack['feats'] = SparseTensor(
+            coords=coords,
+            feats=feats,
+        )
+        
+        pack['image'] = torch.stack([b['image'] for b in batch])
+        pack['alpha'] = torch.stack([b['alpha'] for b in batch])
+        pack['extrinsics'] = torch.stack([b['extrinsics'] for b in batch])
+        pack['intrinsics'] = torch.stack([b['intrinsics'] for b in batch])
+
+        return pack
+
+    def get_instance(self, root, instance):
+        image = self._get_image(root, instance)
+        feat = self._get_feat(root, instance)
+        return {
+            **image,
+            **feat,
+        }

trellis/datasets/sparse_structure.py

@@ -0,0 +1,107 @@
+import os
+import json
+from typing import Union
+import numpy as np
+import pandas as pd
+import torch
+from torch.utils.data import Dataset
+import utils3d
+from .components import StandardDatasetBase
+from ..representations.octree import DfsOctree as Octree
+from ..renderers import OctreeRenderer
+
+
+class SparseStructure(StandardDatasetBase):
+    """
+    Sparse structure 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 = 64,
+        min_aesthetic_score: float = 5.0,
+    ):
+        self.resolution = resolution
+        self.min_aesthetic_score = min_aesthetic_score
+        self.value_range = (0, 1)
+
+        super().__init__(roots)
+        
+    def filter_metadata(self, metadata):
+        stats = {}
+        metadata = metadata[metadata[f'voxelized']]
+        stats['Voxelized'] = 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):
+        position = utils3d.io.read_ply(os.path.join(root, 'voxels', f'{instance}.ply'))[0]
+        coords = ((torch.tensor(position) + 0.5) * self.resolution).int().contiguous()
+        ss = torch.zeros(1, self.resolution, self.resolution, self.resolution, dtype=torch.long)
+        ss[:, coords[:, 0], coords[:, 1], coords[:, 2]] = 1
+        return {'ss': ss}
+
+    @torch.no_grad()
+    def visualize_sample(self, ss: Union[torch.Tensor, dict]):
+        ss = ss if isinstance(ss, torch.Tensor) else ss['ss']
+        
+        renderer = OctreeRenderer()
+        renderer.rendering_options.resolution = 512
+        renderer.rendering_options.near = 0.8
+        renderer.rendering_options.far = 1.6
+        renderer.rendering_options.bg_color = (0, 0, 0)
+        renderer.rendering_options.ssaa = 4
+        renderer.pipe.primitive = 'voxel'
+        
+        # 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 = []
+        
+        # Build each representation
+        ss = ss.cuda()
+        for i in range(ss.shape[0]):
+            representation = Octree(
+                depth=10,
+                aabb=[-0.5, -0.5, -0.5, 1, 1, 1],
+                device='cuda',
+                primitive='voxel',
+                sh_degree=0,
+                primitive_config={'solid': True},
+            )
+            coords = torch.nonzero(ss[i, 0], as_tuple=False)
+            representation.position = coords.float() / self.resolution
+            representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(self.resolution)), dtype=torch.uint8, device='cuda')
+
+            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, colors_overwrite=representation.position)
+                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)
+       

trellis/datasets/sparse_structure_latent.py

@@ -0,0 +1,189 @@
+import os
+import json
+from typing import *
+import numpy as np
+import torch
+import utils3d
+from ..representations.octree import DfsOctree as Octree
+from ..renderers import OctreeRenderer
+from .components import StandardDatasetBase, TextConditionedMixin, ImageConditionedMixin
+from .. import models
+from ..utils.dist_utils import read_file_dist
+
+
+class SparseStructureLatentVisMixin:
+    def __init__(
+        self,
+        *args,
+        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,
+        **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(read_file_dist(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 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):
+            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 = OctreeRenderer()
+        renderer.rendering_options.resolution = 512
+        renderer.rendering_options.near = 0.8
+        renderer.rendering_options.far = 1.6
+        renderer.rendering_options.bg_color = (0, 0, 0)
+        renderer.rendering_options.ssaa = 4
+        renderer.pipe.primitive = 'voxel'
+        
+        # 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 = []
+        
+        # Build each representation
+        x_0 = x_0.cuda()
+        for i in range(x_0.shape[0]):
+            representation = Octree(
+                depth=10,
+                aabb=[-0.5, -0.5, -0.5, 1, 1, 1],
+                device='cuda',
+                primitive='voxel',
+                sh_degree=0,
+                primitive_config={'solid': True},
+            )
+            coords = torch.nonzero(x_0[i, 0] > 0, as_tuple=False)
+            resolution = x_0.shape[-1]
+            representation.position = coords.float() / resolution
+            representation.depth = torch.full((representation.position.shape[0], 1), int(np.log2(resolution)), dtype=torch.uint8, device='cuda')
+
+            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, colors_overwrite=representation.position)
+                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
+        latent_model (str): name of the latent model
+        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,
+        *,
+        latent_model: 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.latent_model = latent_model
+        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[f'ss_latent_{self.latent_model}']]
+        stats['With sparse structure latents'] = 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_latents', self.latent_model, f'{instance}.npz'))
+        z = torch.tensor(latent['mean']).float()
+        if self.normalization is not None:
+            z = (z - self.mean) / self.std
+
+        pack = {
+            'x_0': z,
+        }
+        return pack
+    
+
+class TextConditionedSparseStructureLatent(TextConditionedMixin, SparseStructureLatent):
+    """
+    Text-conditioned sparse structure dataset
+    """
+    pass
+
+
+class ImageConditionedSparseStructureLatent(ImageConditionedMixin, SparseStructureLatent):
+    """
+    Image-conditioned sparse structure dataset
+    """
+    pass
+    

trellis/datasets/structured_latent.py

@@ -0,0 +1,218 @@
+import json
+import os
+from typing import *
+import numpy as np
+import torch
+import utils3d.torch
+from .components import StandardDatasetBase, TextConditionedMixin, ImageConditionedMixin
+from ..modules.sparse.basic import SparseTensor
+from .. import models
+from ..utils.render_utils import get_renderer
+from ..utils.dist_utils import read_file_dist
+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(read_file_dist(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 dataset
+    
+    Args:
+        roots (str): path to the dataset
+        latent_model (str): name of the latent model
+        min_aesthetic_score (float): minimum aesthetic score
+        max_num_voxels (int): maximum number of voxels
+        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,
+        *,
+        latent_model: str,
+        min_aesthetic_score: float = 5.0,
+        max_num_voxels: int = 32768,
+        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.latent_model = latent_model
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_num_voxels = max_num_voxels
+        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, 'num_voxels'] 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'latent_{self.latent_model}']]
+        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['num_voxels'] <= self.max_num_voxels]
+        stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata)
+        return metadata, stats
+
+    def get_instance(self, root, instance):
+        data = np.load(os.path.join(root, 'latents', self.latent_model, 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 TextConditionedSLat(TextConditionedMixin, SLat):
+    """
+    Text conditioned structured latent dataset
+    """
+    pass
+
+
+class ImageConditionedSLat(ImageConditionedMixin, SLat):
+    """
+    Image conditioned structured latent dataset
+    """
+    pass

trellis/datasets/structured_latent2render.py

@@ -0,0 +1,160 @@
+import os
+from PIL import Image
+import json
+import numpy as np
+import torch
+import utils3d.torch
+from ..modules.sparse.basic import SparseTensor
+from .components import StandardDatasetBase
+
+
+class SLat2Render(StandardDatasetBase):
+    """
+    Dataset for Structured Latent and rendered images.
+    
+    Args:
+        roots (str): paths to the dataset
+        image_size (int): size of the image
+        latent_model (str): latent model name
+        min_aesthetic_score (float): minimum aesthetic score
+        max_num_voxels (int): maximum number of voxels
+    """
+    def __init__(
+        self,
+        roots: str,
+        image_size: int,
+        latent_model: str,
+        min_aesthetic_score: float = 5.0,
+        max_num_voxels: int = 32768,
+    ):
+        self.image_size = image_size
+        self.latent_model = latent_model
+        self.min_aesthetic_score = min_aesthetic_score
+        self.max_num_voxels = max_num_voxels
+        self.value_range = (0, 1)
+        
+        super().__init__(roots)
+        
+    def filter_metadata(self, metadata):
+        stats = {}
+        metadata = metadata[metadata[f'latent_{self.latent_model}']]
+        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['num_voxels'] <= self.max_num_voxels]
+        stats[f'Num voxels <= {self.max_num_voxels}'] = len(metadata)
+        return metadata, stats
+
+    def _get_image(self, root, instance):
+        with open(os.path.join(root, 'renders', instance, 'transforms.json')) as f:
+            metadata = json.load(f)
+        n_views = len(metadata['frames'])
+        view = np.random.randint(n_views)
+        metadata = metadata['frames'][view]
+        fov = metadata['camera_angle_x']
+        intrinsics = utils3d.torch.intrinsics_from_fov_xy(torch.tensor(fov), torch.tensor(fov))
+        c2w = torch.tensor(metadata['transform_matrix'])
+        c2w[:3, 1:3] *= -1
+        extrinsics = torch.inverse(c2w)
+
+        image_path = os.path.join(root, 'renders', instance, metadata['file_path'])
+        image = Image.open(image_path)
+        alpha = image.getchannel(3)
+        image = image.convert('RGB')
+        image = image.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+        alpha = alpha.resize((self.image_size, self.image_size), Image.Resampling.LANCZOS)
+        image = torch.tensor(np.array(image)).permute(2, 0, 1).float() / 255.0
+        alpha = torch.tensor(np.array(alpha)).float() / 255.0
+        
+        return {
+            'image': image,
+            'alpha': alpha,
+            'extrinsics': extrinsics,
+            'intrinsics': intrinsics,
+        }
+    
+    def _get_latent(self, root, instance):
+        data = np.load(os.path.join(root, 'latents', self.latent_model, f'{instance}.npz'))
+        coords = torch.tensor(data['coords']).int()
+        feats = torch.tensor(data['feats']).float()
+        return {
+            'coords': coords,
+            'feats': feats,
+        }
+
+    @torch.no_grad()
+    def visualize_sample(self, sample: dict):
+        return sample['image']
+
+    @staticmethod
+    def collate_fn(batch):
+        pack = {}
+        coords = []
+        for i, b in enumerate(batch):
+            coords.append(torch.cat([torch.full((b['coords'].shape[0], 1), i, dtype=torch.int32), b['coords']], dim=-1))
+        coords = torch.cat(coords)
+        feats = torch.cat([b['feats'] for b in batch])
+        pack['latents'] = SparseTensor(
+            coords=coords,
+            feats=feats,
+        )
+        
+        # collate other data
+        keys = [k for k in batch[0].keys() if k not in ['coords', 'feats']]
+        for k in keys:
+            if isinstance(batch[0][k], torch.Tensor):
+                pack[k] = torch.stack([b[k] for b in batch])
+            elif isinstance(batch[0][k], list):
+                pack[k] = sum([b[k] for b in batch], [])
+            else:
+                pack[k] = [b[k] for b in batch]
+
+        return pack
+
+    def get_instance(self, root, instance):
+        image = self._get_image(root, instance)
+        latent = self._get_latent(root, instance)
+        return {
+            **image,
+            **latent,
+        }
+        
+
+class Slat2RenderGeo(SLat2Render):
+    def __init__(
+        self,
+        roots: str,
+        image_size: int,
+        latent_model: str,
+        min_aesthetic_score: float = 5.0,
+        max_num_voxels: int = 32768,
+    ):
+        super().__init__(
+            roots,
+            image_size,
+            latent_model,
+            min_aesthetic_score,
+            max_num_voxels,
+        )
+        
+    def _get_geo(self, root, instance):
+        verts, face = utils3d.io.read_ply(os.path.join(root, 'renders', instance, 'mesh.ply'))
+        mesh = {
+            "vertices" : torch.from_numpy(verts),
+            "faces" : torch.from_numpy(face),
+        }
+        return  {
+            "mesh" : mesh,
+        }
+        
+    def get_instance(self, root, instance):
+        image = self._get_image(root, instance)
+        latent = self._get_latent(root, instance)
+        geo = self._get_geo(root, instance)
+        return {
+            **image,
+            **latent,
+            **geo,
+        }
+        
+        

trellis/models/__init__.py

@@ -3,12 +3,20 @@ import importlib
 __attributes = {
     'SparseStructureEncoder': 'sparse_structure_vae',
     'SparseStructureDecoder': 'sparse_structure_vae',
+    
     'SparseStructureFlowModel': 'sparse_structure_flow',
+    
     'SLatEncoder': 'structured_latent_vae',
     'SLatGaussianDecoder': 'structured_latent_vae',
     'SLatRadianceFieldDecoder': 'structured_latent_vae',
     'SLatMeshDecoder': 'structured_latent_vae',
+    'ElasticSLatEncoder': 'structured_latent_vae',
+    'ElasticSLatGaussianDecoder': 'structured_latent_vae',
+    'ElasticSLatRadianceFieldDecoder': 'structured_latent_vae',
+    'ElasticSLatMeshDecoder': 'structured_latent_vae',
+    
     'SLatFlowModel': 'structured_latent_flow',
+    'ElasticSLatFlowModel': 'structured_latent_flow',
 }
 
 __submodules = []
@@ -64,7 +72,25 @@ def from_pretrained(path: str, **kwargs):
 
 # For Pylance
 if __name__ == '__main__':
-    from .sparse_structure_vae import SparseStructureEncoder, SparseStructureDecoder
+    from .sparse_structure_vae import (
+        SparseStructureEncoder, 
+        SparseStructureDecoder,
+    )
+    
     from .sparse_structure_flow import SparseStructureFlowModel
-    from .structured_latent_vae import SLatEncoder, SLatGaussianDecoder, SLatRadianceFieldDecoder, SLatMeshDecoder
-    from .structured_latent_flow import SLatFlowModel
+    
+    from .structured_latent_vae import (
+        SLatEncoder,
+        SLatGaussianDecoder,
+        SLatRadianceFieldDecoder,
+        SLatMeshDecoder,
+        ElasticSLatEncoder,
+        ElasticSLatGaussianDecoder,
+        ElasticSLatRadianceFieldDecoder,
+        ElasticSLatMeshDecoder,
+    )
+    
+    from .structured_latent_flow import (
+        SLatFlowModel,
+        ElasticSLatFlowModel,
+    )

trellis/models/sparse_elastic_mixin.py

@@ -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

trellis/models/structured_latent_flow.py

@@ -9,6 +9,7 @@ from ..modules.norm import LayerNorm32
 from ..modules import sparse as sp
 from ..modules.sparse.transformer import ModulatedSparseTransformerCrossBlock
 from .sparse_structure_flow import TimestepEmbedder
+from .sparse_elastic_mixin import SparseTransformerElasticMixin
 
 
 class SparseResBlock3d(nn.Module):
@@ -109,8 +110,9 @@ class SLatFlowModel(nn.Module):
         self.qk_rms_norm_cross = qk_rms_norm_cross
         self.dtype = torch.float16 if use_fp16 else torch.float32
 
-        assert int(np.log2(patch_size)) == np.log2(patch_size), "Patch size must be a power of 2"
-        assert np.log2(patch_size) == len(io_block_channels), "Number of IO ResBlocks must match the number of stages"
+        if self.io_block_channels is not None:
+            assert int(np.log2(patch_size)) == np.log2(patch_size), "Patch size must be a power of 2"
+            assert np.log2(patch_size) == len(io_block_channels), "Number of IO ResBlocks must match the number of stages"
 
         self.t_embedder = TimestepEmbedder(model_channels)
         if share_mod:
@@ -122,25 +124,27 @@ class SLatFlowModel(nn.Module):
         if pe_mode == "ape":
             self.pos_embedder = AbsolutePositionEmbedder(model_channels)
 
-        self.input_layer = sp.SparseLinear(in_channels, io_block_channels[0])
+        self.input_layer = sp.SparseLinear(in_channels, model_channels if io_block_channels is None else io_block_channels[0])
+        
         self.input_blocks = nn.ModuleList([])
-        for chs, next_chs in zip(io_block_channels, io_block_channels[1:] + [model_channels]):
-            self.input_blocks.extend([
-                SparseResBlock3d(
-                    chs,
-                    model_channels,
-                    out_channels=chs,
-                )
-                for _ in range(num_io_res_blocks-1)
-            ])
-            self.input_blocks.append(
-                SparseResBlock3d(
-                    chs,
-                    model_channels,
-                    out_channels=next_chs,
-                    downsample=True,
+        if io_block_channels is not None:
+            for chs, next_chs in zip(io_block_channels, io_block_channels[1:] + [model_channels]):
+                self.input_blocks.extend([
+                    SparseResBlock3d(
+                        chs,
+                        model_channels,
+                        out_channels=chs,
+                    )
+                    for _ in range(num_io_res_blocks-1)
+                ])
+                self.input_blocks.append(
+                    SparseResBlock3d(
+                        chs,
+                        model_channels,
+                        out_channels=next_chs,
+                        downsample=True,
+                    )
                 )
-            )
             
         self.blocks = nn.ModuleList([
             ModulatedSparseTransformerCrossBlock(
@@ -159,24 +163,26 @@ class SLatFlowModel(nn.Module):
         ])
 
         self.out_blocks = nn.ModuleList([])
-        for chs, prev_chs in zip(reversed(io_block_channels), [model_channels] + list(reversed(io_block_channels[1:]))):
-            self.out_blocks.append(
-                SparseResBlock3d(
-                    prev_chs * 2 if self.use_skip_connection else prev_chs,
-                    model_channels,
-                    out_channels=chs,
-                    upsample=True,
-                )
-            )
-            self.out_blocks.extend([
-                SparseResBlock3d(
-                    chs * 2 if self.use_skip_connection else chs,
-                    model_channels,
-                    out_channels=chs,
+        if io_block_channels is not None:
+            for chs, prev_chs in zip(reversed(io_block_channels), [model_channels] + list(reversed(io_block_channels[1:]))):
+                self.out_blocks.append(
+                    SparseResBlock3d(
+                        prev_chs * 2 if self.use_skip_connection else prev_chs,
+                        model_channels,
+                        out_channels=chs,
+                        upsample=True,
+                    )
                 )
-                for _ in range(num_io_res_blocks-1)
-            ])
-        self.out_layer = sp.SparseLinear(io_block_channels[0], out_channels)
+                self.out_blocks.extend([
+                    SparseResBlock3d(
+                        chs * 2 if self.use_skip_connection else chs,
+                        model_channels,
+                        out_channels=chs,
+                    )
+                    for _ in range(num_io_res_blocks-1)
+                ])
+            
+        self.out_layer = sp.SparseLinear(model_channels if io_block_channels is None else io_block_channels[0], out_channels)
 
         self.initialize_weights()
         if use_fp16:
@@ -260,3 +266,11 @@ class SLatFlowModel(nn.Module):
         h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
         h = self.out_layer(h.type(x.dtype))
         return h
+    
+
+class ElasticSLatFlowModel(SparseTransformerElasticMixin, SLatFlowModel):
+    """
+    SLat Flow Model with elastic memory management.
+    Used for training with low VRAM.
+    """
+    pass

trellis/models/structured_latent_vae/__init__.py

@@ -1,4 +1,4 @@
-from .encoder import SLatEncoder
-from .decoder_gs import SLatGaussianDecoder
-from .decoder_rf import SLatRadianceFieldDecoder
-from .decoder_mesh import SLatMeshDecoder
+from .encoder import SLatEncoder, ElasticSLatEncoder
+from .decoder_gs import SLatGaussianDecoder, ElasticSLatGaussianDecoder
+from .decoder_rf import SLatRadianceFieldDecoder, ElasticSLatRadianceFieldDecoder
+from .decoder_mesh import SLatMeshDecoder, ElasticSLatMeshDecoder

trellis/models/structured_latent_vae/decoder_gs.py

@@ -6,6 +6,7 @@ from ...modules import sparse as sp
 from ...utils.random_utils import hammersley_sequence
 from .base import SparseTransformerBase
 from ...representations import Gaussian
+from ..sparse_elastic_mixin import SparseTransformerElasticMixin
 
 
 class SLatGaussianDecoder(SparseTransformerBase):
@@ -120,3 +121,11 @@ class SLatGaussianDecoder(SparseTransformerBase):
         h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
         h = self.out_layer(h)
         return self.to_representation(h)
+    
+
+class ElasticSLatGaussianDecoder(SparseTransformerElasticMixin, SLatGaussianDecoder):
+    """
+    Slat VAE Gaussian decoder with elastic memory management.
+    Used for training with low VRAM.
+    """
+    pass

trellis/models/structured_latent_vae/decoder_mesh.py

@@ -8,6 +8,7 @@ from ...modules import sparse as sp
 from .base import SparseTransformerBase
 from ...representations import MeshExtractResult
 from ...representations.mesh import SparseFeatures2Mesh
+from ..sparse_elastic_mixin import SparseTransformerElasticMixin
 
 
 class SparseSubdivideBlock3d(nn.Module):
@@ -165,3 +166,11 @@ class SLatMeshDecoder(SparseTransformerBase):
         h = h.type(x.dtype)
         h = self.out_layer(h)
         return self.to_representation(h)
+    
+
+class ElasticSLatMeshDecoder(SparseTransformerElasticMixin, SLatMeshDecoder):
+    """
+    Slat VAE Mesh decoder with elastic memory management.
+    Used for training with low VRAM.
+    """
+    pass

trellis/models/structured_latent_vae/decoder_rf.py

@@ -6,6 +6,7 @@ import numpy as np
 from ...modules import sparse as sp
 from .base import SparseTransformerBase
 from ...representations import Strivec
+from ..sparse_elastic_mixin import SparseTransformerElasticMixin
 
 
 class SLatRadianceFieldDecoder(SparseTransformerBase):
@@ -102,3 +103,11 @@ class SLatRadianceFieldDecoder(SparseTransformerBase):
         h = h.replace(F.layer_norm(h.feats, h.feats.shape[-1:]))
         h = self.out_layer(h)
         return self.to_representation(h)
+
+
+class ElasticSLatRadianceFieldDecoder(SparseTransformerElasticMixin, SLatRadianceFieldDecoder):
+    """
+    Slat VAE Radiance Field Decoder with elastic memory management.
+    Used for training with low VRAM.
+    """
+    pass

trellis/models/structured_latent_vae/encoder.py

@@ -4,6 +4,7 @@ import torch.nn as nn
 import torch.nn.functional as F
 from ...modules import sparse as sp
 from .base import SparseTransformerBase
+from ..sparse_elastic_mixin import SparseTransformerElasticMixin
 
 
 class SLatEncoder(SparseTransformerBase):
@@ -70,3 +71,10 @@ class SLatEncoder(SparseTransformerBase):
             return z, mean, logvar
         else:
             return z
+        
+
+class ElasticSLatEncoder(SparseTransformerElasticMixin, SLatEncoder):
+    """
+    SLat VAE encoder with elastic memory management.
+    Used for training with low VRAM.
+    """

trellis/pipelines/__init__.py

@@ -1,5 +1,6 @@
 from . import samplers
 from .trellis_image_to_3d import TrellisImageTo3DPipeline
+from .trellis_text_to_3d import TrellisTextTo3DPipeline
 
 
 def from_pretrained(path: str):

trellis/pipelines/base.py

@@ -36,10 +36,12 @@ class Pipeline:
         with open(config_file, 'r') as f:
             args = json.load(f)['args']
 
-        _models = {
-            k: models.from_pretrained(f"{path}/{v}")
-            for k, v in args['models'].items()
-        }
+        _models = {}
+        for k, v in args['models'].items():
+            try:
+                _models[k] = models.from_pretrained(f"{path}/{v}")
+            except:
+                _models[k] = models.from_pretrained(v)
 
         new_pipeline = Pipeline(_models)
         new_pipeline._pretrained_args = args

trellis/pipelines/samplers/flow_euler.py

@@ -37,6 +37,8 @@ class FlowEulerSampler(Sampler):
 
     def _inference_model(self, model, x_t, t, cond=None, **kwargs):
         t = torch.tensor([1000 * t] * x_t.shape[0], device=x_t.device, dtype=torch.float32)
+        if cond is not None and cond.shape[0] == 1 and x_t.shape[0] > 1:
+            cond = cond.repeat(x_t.shape[0], *([1] * (len(cond.shape) - 1)))
         return model(x_t, t, cond, **kwargs)
 
     def _get_model_prediction(self, model, x_t, t, cond=None, **kwargs):

trellis/pipelines/trellis_image_to_3d.py

@@ -4,15 +4,12 @@ import torch
 import torch.nn as nn
 import torch.nn.functional as F
 import numpy as np
-from tqdm import tqdm
-from easydict import EasyDict as edict
 from torchvision import transforms
 from PIL import Image
 import rembg
 from .base import Pipeline
 from . import samplers
 from ..modules import sparse as sp
-from ..representations import Gaussian, Strivec, MeshExtractResult
 
 
 class TrellisImageTo3DPipeline(Pipeline):
@@ -271,8 +268,10 @@ class TrellisImageTo3DPipeline(Pipeline):
         Args:
             image (Image.Image): The image prompt.
             num_samples (int): The number of samples to generate.
+            seed (int): The random seed.
             sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler.
             slat_sampler_params (dict): Additional parameters for the structured latent sampler.
+            formats (List[str]): The formats to decode the structured latent to.
             preprocess_image (bool): Whether to preprocess the image.
         """
         if preprocess_image:

trellis/pipelines/trellis_text_to_3d.py

@@ -0,0 +1,278 @@
+from typing import *
+import torch
+import torch.nn as nn
+import numpy as np
+from transformers import CLIPTextModel, AutoTokenizer
+import open3d as o3d
+from .base import Pipeline
+from . import samplers
+from ..modules import sparse as sp
+
+
+class TrellisTextTo3DPipeline(Pipeline):
+    """
+    Pipeline for inferring Trellis text-to-3D models.
+
+    Args:
+        models (dict[str, nn.Module]): The models to use in the pipeline.
+        sparse_structure_sampler (samplers.Sampler): The sampler for the sparse structure.
+        slat_sampler (samplers.Sampler): The sampler for the structured latent.
+        slat_normalization (dict): The normalization parameters for the structured latent.
+        text_cond_model (str): The name of the text conditioning model.
+    """
+    def __init__(
+        self,
+        models: dict[str, nn.Module] = None,
+        sparse_structure_sampler: samplers.Sampler = None,
+        slat_sampler: samplers.Sampler = None,
+        slat_normalization: dict = None,
+        text_cond_model: str = None,
+    ):
+        if models is None:
+            return
+        super().__init__(models)
+        self.sparse_structure_sampler = sparse_structure_sampler
+        self.slat_sampler = slat_sampler
+        self.sparse_structure_sampler_params = {}
+        self.slat_sampler_params = {}
+        self.slat_normalization = slat_normalization
+        self._init_text_cond_model(text_cond_model)
+
+    @staticmethod
+    def from_pretrained(path: str) -> "TrellisTextTo3DPipeline":
+        """
+        Load a pretrained model.
+
+        Args:
+            path (str): The path to the model. Can be either local path or a Hugging Face repository.
+        """
+        pipeline = super(TrellisTextTo3DPipeline, TrellisTextTo3DPipeline).from_pretrained(path)
+        new_pipeline = TrellisTextTo3DPipeline()
+        new_pipeline.__dict__ = pipeline.__dict__
+        args = pipeline._pretrained_args
+
+        new_pipeline.sparse_structure_sampler = getattr(samplers, args['sparse_structure_sampler']['name'])(**args['sparse_structure_sampler']['args'])
+        new_pipeline.sparse_structure_sampler_params = args['sparse_structure_sampler']['params']
+
+        new_pipeline.slat_sampler = getattr(samplers, args['slat_sampler']['name'])(**args['slat_sampler']['args'])
+        new_pipeline.slat_sampler_params = args['slat_sampler']['params']
+
+        new_pipeline.slat_normalization = args['slat_normalization']
+
+        new_pipeline._init_text_cond_model(args['text_cond_model'])
+
+        return new_pipeline
+    
+    def _init_text_cond_model(self, name: str):
+        """
+        Initialize the text conditioning model.
+        """
+        # load model
+        model = CLIPTextModel.from_pretrained(name)
+        tokenizer = AutoTokenizer.from_pretrained(name)
+        model.eval()
+        model = model.cuda()
+        self.text_cond_model = {
+            'model': model,
+            'tokenizer': tokenizer,
+        }
+        self.text_cond_model['null_cond'] = self.encode_text([''])
+
+    @torch.no_grad()
+    def encode_text(self, text: List[str]) -> torch.Tensor:
+        """
+        Encode the text.
+        """
+        assert isinstance(text, list) and all(isinstance(t, str) for t in text), "text must be a list of strings"
+        encoding = self.text_cond_model['tokenizer'](text, max_length=77, padding='max_length', truncation=True, return_tensors='pt')
+        tokens = encoding['input_ids'].cuda()
+        embeddings = self.text_cond_model['model'](input_ids=tokens).last_hidden_state
+        
+        return embeddings
+        
+    def get_cond(self, prompt: List[str]) -> dict:
+        """
+        Get the conditioning information for the model.
+
+        Args:
+            prompt (List[str]): The text prompt.
+
+        Returns:
+            dict: The conditioning information
+        """
+        cond = self.encode_text(prompt)
+        neg_cond = self.text_cond_model['null_cond']
+        return {
+            'cond': cond,
+            'neg_cond': neg_cond,
+        }
+
+    def sample_sparse_structure(
+        self,
+        cond: dict,
+        num_samples: int = 1,
+        sampler_params: dict = {},
+    ) -> torch.Tensor:
+        """
+        Sample sparse structures with the given conditioning.
+        
+        Args:
+            cond (dict): The conditioning information.
+            num_samples (int): The number of samples to generate.
+            sampler_params (dict): Additional parameters for the sampler.
+        """
+        # Sample occupancy latent
+        flow_model = self.models['sparse_structure_flow_model']
+        reso = flow_model.resolution
+        noise = torch.randn(num_samples, flow_model.in_channels, reso, reso, reso).to(self.device)
+        sampler_params = {**self.sparse_structure_sampler_params, **sampler_params}
+        z_s = self.sparse_structure_sampler.sample(
+            flow_model,
+            noise,
+            **cond,
+            **sampler_params,
+            verbose=True
+        ).samples
+        
+        # Decode occupancy latent
+        decoder = self.models['sparse_structure_decoder']
+        coords = torch.argwhere(decoder(z_s)>0)[:, [0, 2, 3, 4]].int()
+
+        return coords
+
+    def decode_slat(
+        self,
+        slat: sp.SparseTensor,
+        formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
+    ) -> dict:
+        """
+        Decode the structured latent.
+
+        Args:
+            slat (sp.SparseTensor): The structured latent.
+            formats (List[str]): The formats to decode the structured latent to.
+
+        Returns:
+            dict: The decoded structured latent.
+        """
+        ret = {}
+        if 'mesh' in formats:
+            ret['mesh'] = self.models['slat_decoder_mesh'](slat)
+        if 'gaussian' in formats:
+            ret['gaussian'] = self.models['slat_decoder_gs'](slat)
+        if 'radiance_field' in formats:
+            ret['radiance_field'] = self.models['slat_decoder_rf'](slat)
+        return ret
+    
+    def sample_slat(
+        self,
+        cond: dict,
+        coords: torch.Tensor,
+        sampler_params: dict = {},
+    ) -> sp.SparseTensor:
+        """
+        Sample structured latent with the given conditioning.
+        
+        Args:
+            cond (dict): The conditioning information.
+            coords (torch.Tensor): The coordinates of the sparse structure.
+            sampler_params (dict): Additional parameters for the sampler.
+        """
+        # Sample structured latent
+        flow_model = self.models['slat_flow_model']
+        noise = sp.SparseTensor(
+            feats=torch.randn(coords.shape[0], flow_model.in_channels).to(self.device),
+            coords=coords,
+        )
+        sampler_params = {**self.slat_sampler_params, **sampler_params}
+        slat = self.slat_sampler.sample(
+            flow_model,
+            noise,
+            **cond,
+            **sampler_params,
+            verbose=True
+        ).samples
+
+        std = torch.tensor(self.slat_normalization['std'])[None].to(slat.device)
+        mean = torch.tensor(self.slat_normalization['mean'])[None].to(slat.device)
+        slat = slat * std + mean
+        
+        return slat
+
+    @torch.no_grad()
+    def run(
+        self,
+        prompt: str,
+        num_samples: int = 1,
+        seed: int = 42,
+        sparse_structure_sampler_params: dict = {},
+        slat_sampler_params: dict = {},
+        formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
+    ) -> dict:
+        """
+        Run the pipeline.
+
+        Args:
+            prompt (str): The text prompt.
+            num_samples (int): The number of samples to generate.
+            seed (int): The random seed.
+            sparse_structure_sampler_params (dict): Additional parameters for the sparse structure sampler.
+            slat_sampler_params (dict): Additional parameters for the structured latent sampler.
+            formats (List[str]): The formats to decode the structured latent to.
+        """
+        cond = self.get_cond([prompt])
+        torch.manual_seed(seed)
+        coords = self.sample_sparse_structure(cond, num_samples, sparse_structure_sampler_params)
+        slat = self.sample_slat(cond, coords, slat_sampler_params)
+        return self.decode_slat(slat, formats)
+    
+    def voxelize(self, mesh: o3d.geometry.TriangleMesh) -> torch.Tensor:
+        """
+        Voxelize a mesh.
+
+        Args:
+            mesh (o3d.geometry.TriangleMesh): The mesh to voxelize.
+            sha256 (str): The SHA256 hash of the mesh.
+            output_dir (str): The output directory.
+        """
+        vertices = np.asarray(mesh.vertices)
+        aabb = np.stack([vertices.min(0), vertices.max(0)])
+        center = (aabb[0] + aabb[1]) / 2
+        scale = (aabb[1] - aabb[0]).max()
+        vertices = (vertices - center) / scale
+        vertices = np.clip(vertices, -0.5 + 1e-6, 0.5 - 1e-6)
+        mesh.vertices = o3d.utility.Vector3dVector(vertices)
+        voxel_grid = o3d.geometry.VoxelGrid.create_from_triangle_mesh_within_bounds(mesh, voxel_size=1/64, min_bound=(-0.5, -0.5, -0.5), max_bound=(0.5, 0.5, 0.5))
+        vertices = np.array([voxel.grid_index for voxel in voxel_grid.get_voxels()])
+        return torch.tensor(vertices).int().cuda()
+
+    @torch.no_grad()
+    def run_variant(
+        self,
+        mesh: o3d.geometry.TriangleMesh,
+        prompt: str,
+        num_samples: int = 1,
+        seed: int = 42,
+        slat_sampler_params: dict = {},
+        formats: List[str] = ['mesh', 'gaussian', 'radiance_field'],
+    ) -> dict:
+        """
+        Run the pipeline for making variants of an asset.
+
+        Args:
+            mesh (o3d.geometry.TriangleMesh): The base mesh.
+            prompt (str): The text prompt.
+            num_samples (int): The number of samples to generate.
+            seed (int): The random seed
+            slat_sampler_params (dict): Additional parameters for the structured latent sampler.
+            formats (List[str]): The formats to decode the structured latent to.
+        """
+        cond = self.get_cond([prompt])
+        coords = self.voxelize(mesh)
+        coords = torch.cat([
+            torch.arange(num_samples).repeat_interleave(coords.shape[0], 0)[:, None].int().cuda(),
+            coords.repeat(num_samples, 1)
+        ], 1)
+        torch.manual_seed(seed)
+        slat = self.sample_slat(cond, coords, slat_sampler_params)
+        return self.decode_slat(slat, formats)

trellis/representations/mesh/cube2mesh.py

@@ -1,15 +1,8 @@
-# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
-#
-# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
-# and proprietary rights in and to this software, related documentation
-# and any modifications thereto.  Any use, reproduction, disclosure or
-# distribution of this software and related documentation without an express
-# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
 import torch
 from ...modules.sparse import SparseTensor
 from easydict import EasyDict as edict
 from .utils_cube import *
-from .flexicube import FlexiCubes
+from .flexicubes.flexicubes import FlexiCubes
 
 
 class MeshExtractResult:

trellis/representations/mesh/flexicubes/LICENSE.txt

@@ -0,0 +1,90 @@
+Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+
+NVIDIA Source Code License for FlexiCubes
+
+
+=======================================================================
+
+1. Definitions
+
+“Licensor” means any person or entity that distributes its Work.
+
+“Work” means (a) the original work of authorship made available under
+this license, which may include software, documentation, or other files,
+and (b) any additions to or derivative works  thereof  that are made
+available under this license.
+
+The terms “reproduce,” “reproduction,” “derivative works,” and
+“distribution” have the meaning as provided under U.S. copyright law;
+provided, however, that for the purposes of this license, derivative works
+shall not include works that remain separable from, or merely link
+(or bind by name) to the interfaces of, the Work.
+
+Works are “made available” under this license by including in or with
+the Work either (a) a copyright notice referencing the applicability of
+this license to the Work, or (b) a copy of this license.
+
+2. License Grant
+
+    2.1 Copyright Grant. Subject to the terms and conditions of this license,
+     each Licensor grants to you a perpetual, worldwide, non-exclusive,
+     royalty-free, copyright license to use, reproduce, prepare derivative
+     works of, publicly display, publicly perform, sublicense and distribute
+     its Work and any resulting derivative works in any form.
+
+3. Limitations
+
+    3.1 Redistribution. You may reproduce or distribute the Work only if
+    (a) you do so under this license, (b) you include a complete copy of
+    this license with your distribution, and (c) you retain without
+    modification any copyright, patent, trademark, or attribution notices
+    that are present in the Work.
+
+    3.2 Derivative Works. You may specify that additional or different terms
+     apply to the use, reproduction, and distribution of your derivative
+     works of the Work (“Your Terms”) only if (a) Your Terms provide that the
+     use limitation in Section 3.3 applies to your derivative works, and (b)
+     you identify the specific derivative works that are subject to Your Terms.
+     Notwithstanding Your Terms, this license (including the redistribution
+     requirements in Section 3.1) will continue to apply to the Work itself.
+
+    3.3 Use Limitation. The Work and any derivative works thereof only may be
+     used or intended for use non-commercially. Notwithstanding the foregoing,
+     NVIDIA Corporation and its affiliates may use the Work and any derivative
+     works commercially. As used herein, “non-commercially” means for research
+     or evaluation purposes only.
+
+    3.4 Patent Claims. If you bring or threaten to bring a patent claim against
+     any Licensor (including any claim, cross-claim or counterclaim in a lawsuit)
+     to enforce any patents that you allege are infringed by any Work, then your
+     rights under this license from such Licensor (including the grant in
+     Section 2.1) will terminate immediately.
+
+    3.5 Trademarks. This license does not grant any rights to use any Licensor’s
+     or its affiliates’ names, logos, or trademarks, except as necessary to
+     reproduce the notices described in this license.
+
+    3.6 Termination. If you violate any term of this license, then your rights
+     under this license (including the grant in Section 2.1) will terminate
+     immediately.
+
+4. Disclaimer of Warranty.
+
+THE WORK IS PROVIDED “AS IS” WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND,
+EITHER EXPRESS OR IMPLIED, INCLUDING WARRANTIES OR CONDITIONS OF
+MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE, TITLE OR NON-INFRINGEMENT.
+YOU BEAR THE RISK OF UNDERTAKING ANY ACTIVITIES UNDER THIS LICENSE.
+
+5. Limitation of Liability.
+
+EXCEPT AS PROHIBITED BY APPLICABLE LAW, IN NO EVENT AND UNDER NO LEGAL THEORY,
+WHETHER IN TORT (INCLUDING NEGLIGENCE), CONTRACT, OR OTHERWISE SHALL ANY
+LICENSOR BE LIABLE TO YOU FOR DAMAGES, INCLUDING ANY DIRECT, INDIRECT, SPECIAL,
+INCIDENTAL, OR CONSEQUENTIAL DAMAGES ARISING OUT OF OR RELATED TO THIS LICENSE,
+THE USE OR INABILITY TO USE THE WORK (INCLUDING BUT NOT LIMITED TO LOSS OF
+GOODWILL, BUSINESS INTERRUPTION, LOST PROFITS OR DATA, COMPUTER FAILURE OR
+MALFUNCTION, OR ANY OTHER DAMAGES OR LOSSES), EVEN IF THE LICENSOR HAS BEEN
+ADVISED OF THE POSSIBILITY OF SUCH DAMAGES.
+
+=======================================================================

trellis/representations/mesh/flexicubes/README.md

@@ -0,0 +1,110 @@
+## Flexible Isosurface Extraction for Gradient-Based Mesh Optimization (FlexiCubes)<br><sub>Official PyTorch implementation </sub>
+
+![Teaser image](<images/teaser_top.png>)
+
+FlexiCubes is a high-quality isosurface representation specifically designed for gradient-based mesh optimization with respect to geometric, visual, or even physical objectives. For more details, please refer to our [paper](https://arxiv.org/abs/2308.05371) and [project page](https://research.nvidia.com/labs/toronto-ai/flexicubes/).
+
+## Highlights
+* [Getting started](https://github.com/nv-tlabs/FlexiCubes#getting-started)
+* [Basic workflow](https://github.com/nv-tlabs/FlexiCubes#example-usage)
+* [nvdiffrec: image-based reconstruction example](https://github.com/NVlabs/nvdiffrec#news)
+* [GET3D: generative AI example](https://github.com/nv-tlabs/GET3D#employing-flexicubes)
+* [Bibtex](https://github.com/nv-tlabs/FlexiCubes#citation)
+
+## Getting Started
+
+The core functions of FlexiCubes are now in [Kaolin](https://github.com/NVIDIAGameWorks/kaolin/) starting from v0.15.0. See installation instructions [here](https://kaolin.readthedocs.io/en/latest/notes/installation.html) and API documentations [here](https://kaolin.readthedocs.io/en/latest/modules/kaolin.non_commercial.html#kaolin.non_commercial.FlexiCubes)
+
+The original code of the paper is still visible in `flexicube.py`.
+
+## Example Usage
+
+### Gradient-Based Mesh Optimization
+We provide examples demonstrating how to use FlexiCubes for reconstructing unknown meshes through gradient-based optimization. Specifically, starting from randomly initialized SDF, we optimize the shape towards the reference mesh by minimizing their geometric difference, measured by multiview mask and depth losses. This workflow is a simplified version of `nvdiffrec` with code largely borrowed from the [nvdiffrec GitHub](https://github.com/NVlabs/nvdiffrec). We use the same pipeline to conduct the analysis in Section 3 and the main experiments described in Section 5 of our paper. We provide a detailed tutorial in `examples/optimization.ipynb`, along with an optimization script in `examples/optimize.py` which accepts command-line arguments.
+
+
+To run the examples, it is suggested to install the Conda environment as detailed below:
+```sh
+conda create -n flexicubes python=3.9
+conda activate flexicubes
+conda install pytorch==1.12.0 torchvision==0.13.0 torchaudio==0.12.0 cudatoolkit=11.3 -c pytorch
+pip install imageio trimesh tqdm matplotlib torch_scatter ninja
+pip install git+https://github.com/NVlabs/nvdiffrast/
+pip install kaolin==0.15.0 -f https://nvidia-kaolin.s3.us-east-2.amazonaws.com/torch-1.12.0_cu113.html
+```
+
+Then download the dataset collected by [Myles et al.](https://vcg.isti.cnr.it/Publications/2014/MPZ14/) as follows. We include one shape in 'examples/data/inputmodels/block.obj' if you want to test without downloading the full dataset.
+
+```sh
+cd examples
+python download_data.py
+```
+
+After downloading the data, run shape optimization with the following example command:
+```sh
+python optimize.py --ref_mesh data/inputmodels/block.obj --out_dir out/block
+```
+You can find visualization and output meshes in the `out/block`. Below, we show the initial and final shapes during optimization, with the reference shape on the right.
+
+<img src="images/block_init.png" alt="block_init" width="80%" height="80%">
+
+<img src="images/block_final.png" alt="block_final" width="80%" height="80%">
+
+
+To further demonstrate the flexibility of our FlexiCubes representation, which can accommodates both reconstruction objectives and regularizers defined on the extracted mesh, you can add a developability regularizer (proposed by [Stein et al.](https://www.cs.cmu.edu/~kmcrane/Projects/DiscreteDevelopable/)) to the previous reconstruction pipeline to encourage fabricability from panels:
+```sh
+python optimize.py --ref_mesh data/inputmodels/david.obj --out_dir out/david_dev --develop_reg True  --iter=1250
+```
+
+### Extract mesh from known signed distance field
+While not its designated use case, our function can extract a mesh from a known Signed Distance Field (SDF) without optimization. Please refer to the tutorial found in `examples/extraction.ipynb` for details.
+
+## Tips for using FlexiCubes
+### Regularization losses: 
+We commonly use three regularizers in our mesh optimization pipelines, referenced in lines `L104-L106` in `examples/optimize.py`. The weights of these regularizers should be scaled according to the your application objectives. Initially, it is suggested to employ low weights because strong regularization can hinder convergence. You can incrementally increase the weights if you notice artifacts appearing in the optimized meshes. Specifically:
+
+* The loss function at `L104` helps to remove floaters in areas of the shape that are not supervised by the application objective, such as internal faces when using image supervision only.
+* The L_dev loss at `L105` can be increased if you observe artifacts in flat areas, as illustrated in the image below.
+* Generally, the L1 regularizer on flexible weights at `L106` does not have a significant impact during the optimization of a single shape. However, we found it to be effective in stabilizing training in generative pipelines such as GET3D.
+<img src="images/ablate_L_dev.jpg" alt="Ablating L_dev" width="80%" height="80%">
+
+### Resolution of voxel grid vs. tetrahedral grid:
+If you are switching from our previous work, DMTet, it's important to note the difference in grid resolution when compared to FlexiCubes. In both implementations, the resolution is defined by the edge length: a grid resolution of `n` means the grid edge length is 1/n for both the voxel and tetrahedral grids. However, a tetrahedral grid with a resolution of `n` contains only `(n/2+1)³` grid vertices, in contrast to the `(n+1)³` vertices in a voxel grid. Consequently, if you are switching from DMTet to FlexiCubes while maintaining the same resolution, you will notice not only a denser output mesh but also a substantial increase in computational cost. To align the triangle count in the output meshes more closely, we recommend adopting a 4:5 resolution ratio between the voxel grid and the tetrahedral grid. For instance, in our paper, `64³` FlexiCubes generate approximately the same number of triangles as `80³` DMTet.
+
+## Applications
+FlexiCubes is now integrated into NVIDIA applications as a drop-in replacement for DMTet. You can visit their GitHub pages to see how FlexiCubes is used in advanced photogrammetry and 3D generative pipelines.
+
+[Extracting Triangular 3D Models, Materials, and Lighting From Images (nvdiffrec)](https://github.com/NVlabs/nvdiffrec#news)
+
+[GET3D: A Generative Model of High Quality 3D Textured Shapes Learned from Images](https://github.com/nv-tlabs/GET3D#employing-flexicubes)
+
+
+
+## License
+Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+
+This work is made available under the [Nvidia Source Code License](LICENSE.txt).
+
+For business inquiries, please visit our website and submit the form: [NVIDIA Research Licensing](https://www.nvidia.com/en-us/research/inquiries/).
+
+## Citation
+```bibtex
+@article{shen2023flexicubes,
+author = {Shen, Tianchang and Munkberg, Jacob and Hasselgren, Jon and Yin, Kangxue and Wang, Zian 
+        and Chen, Wenzheng and Gojcic, Zan and Fidler, Sanja and Sharp, Nicholas and Gao, Jun},
+title = {Flexible Isosurface Extraction for Gradient-Based Mesh Optimization},
+year = {2023},
+issue_date = {August 2023},
+publisher = {Association for Computing Machinery},
+address = {New York, NY, USA},
+volume = {42},
+number = {4},
+issn = {0730-0301},
+url = {https://doi.org/10.1145/3592430},
+doi = {10.1145/3592430},
+journal = {ACM Trans. Graph.},
+month = {jul},
+articleno = {37},
+numpages = {16}
+}
+```

trellis/representations/mesh/flexicubes/examples/download_data.py

@@ -0,0 +1,41 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
+import requests
+from zipfile import ZipFile
+from tqdm import tqdm
+import os
+
+def download_file(url, output_path):
+    response = requests.get(url, stream=True)
+    response.raise_for_status()
+    total_size_in_bytes = int(response.headers.get('content-length', 0))
+    block_size = 1024 #1 Kibibyte
+    progress_bar = tqdm(total=total_size_in_bytes, unit='iB', unit_scale=True)
+    
+    with open(output_path, 'wb') as file:
+        for data in response.iter_content(block_size):
+            progress_bar.update(len(data))
+            file.write(data)
+    progress_bar.close()
+    if total_size_in_bytes != 0 and progress_bar.n != total_size_in_bytes:
+        raise Exception("ERROR, something went wrong")
+
+
+url = "https://vcg.isti.cnr.it/Publications/2014/MPZ14/inputmodels.zip"
+zip_file_path = './data/inputmodels.zip'
+
+os.makedirs('./data', exist_ok=True)
+
+download_file(url, zip_file_path)
+
+with ZipFile(zip_file_path, 'r') as zip_ref:
+    zip_ref.extractall('./data')
+
+os.remove(zip_file_path)
+
+print("Download and extraction complete.")

trellis/representations/mesh/flexicubes/examples/loss.py

@@ -0,0 +1,95 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
+import torch
+import torch_scatter
+
+###############################################################################
+# Pytorch implementation of the developability regularizer introduced in paper 
+# "Developability of Triangle Meshes" by Stein et al.
+###############################################################################
+def mesh_developable_reg(mesh):
+
+    verts = mesh.vertices
+    tris = mesh.faces
+
+    device = verts.device
+    V = verts.shape[0]
+    F = tris.shape[0]
+
+    POS_EPS = 1e-6
+    REL_EPS = 1e-6
+
+    def normalize(vecs):
+        return vecs / (torch.linalg.norm(vecs, dim=-1, keepdim=True) + POS_EPS)
+
+    tri_pos = verts[tris]
+
+    vert_normal_covariance_sum = torch.zeros((V, 9), device=device)
+    vert_area = torch.zeros(V, device=device)
+    vert_degree = torch.zeros(V, dtype=torch.int32, device=device)
+
+    for iC in range(3):  # loop over three corners of each triangle
+
+        # gather tri verts
+        pRoot = tri_pos[:, iC, :]
+        pA = tri_pos[:, (iC + 1) % 3, :]
+        pB = tri_pos[:, (iC + 2) % 3, :]
+
+        # compute the corner angle & normal
+        vA = pA - pRoot
+        vAn = normalize(vA)
+        vB = pB - pRoot
+        vBn = normalize(vB)
+        area_normal = torch.linalg.cross(vA, vB, dim=-1)
+        face_area = 0.5 * torch.linalg.norm(area_normal, dim=-1)
+        normal = normalize(area_normal)
+        corner_angle = torch.acos(torch.clamp(torch.sum(vAn * vBn, dim=-1), min=-1., max=1.))
+
+        # add up the contribution to the covariance matrix
+        outer = normal[:, :, None] @ normal[:, None, :]
+        contrib = corner_angle[:, None] * outer.reshape(-1, 9)
+
+        # scatter the result to the appropriate matrices
+        vert_normal_covariance_sum = torch_scatter.scatter_add(src=contrib,
+                                                               index=tris[:, iC],
+                                                               dim=-2,
+                                                               out=vert_normal_covariance_sum)
+
+        vert_area = torch_scatter.scatter_add(src=face_area / 3.,
+                                              index=tris[:, iC],
+                                              dim=-1,
+                                              out=vert_area)
+
+        vert_degree = torch_scatter.scatter_add(src=torch.ones(F, dtype=torch.int32, device=device),
+                                                index=tris[:, iC],
+                                                dim=-1,
+                                                out=vert_degree)
+
+    # The energy is the smallest eigenvalue of the outer-product matrix
+    vert_normal_covariance_sum = vert_normal_covariance_sum.reshape(
+        -1, 3, 3)  # reshape to a batch of matrices
+    vert_normal_covariance_sum = vert_normal_covariance_sum + torch.eye(
+        3, device=device)[None, :, :] * REL_EPS
+
+    min_eigvals = torch.min(torch.linalg.eigvals(vert_normal_covariance_sum).abs(), dim=-1).values
+
+    # Mask out degree-3 vertices
+    vert_area = torch.where(vert_degree == 3, torch.tensor(0, dtype=vert_area.dtype,device=vert_area.device), vert_area)
+
+    # Adjust the vertex area weighting so it is unit-less, and 1 on average
+    vert_area = vert_area * (V / torch.sum(vert_area, dim=-1, keepdim=True))
+
+    return vert_area * min_eigvals 
+
+def sdf_reg_loss(sdf, all_edges):
+    sdf_f1x6x2 = sdf[all_edges.reshape(-1)].reshape(-1,2)
+    mask = torch.sign(sdf_f1x6x2[...,0]) != torch.sign(sdf_f1x6x2[...,1])
+    sdf_f1x6x2 = sdf_f1x6x2[mask]
+    sdf_diff = torch.nn.functional.binary_cross_entropy_with_logits(sdf_f1x6x2[...,0], (sdf_f1x6x2[...,1] > 0).float()) + \
+            torch.nn.functional.binary_cross_entropy_with_logits(sdf_f1x6x2[...,1], (sdf_f1x6x2[...,0] > 0).float())
+    return sdf_diff

trellis/representations/mesh/flexicubes/examples/optimize.py

@@ -0,0 +1,150 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
+import argparse
+import numpy as np
+import torch
+import nvdiffrast.torch as dr
+import trimesh
+import os
+from util import *
+import render
+import loss
+import imageio
+
+import sys
+sys.path.append('..')
+from flexicubes import FlexiCubes
+
+###############################################################################
+# Functions adapted from https://github.com/NVlabs/nvdiffrec
+###############################################################################
+
+def lr_schedule(iter):
+    return max(0.0, 10**(-(iter)*0.0002)) # Exponential falloff from [1.0, 0.1] over 5k epochs.    
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser(description='flexicubes optimization')
+    parser.add_argument('-o', '--out_dir', type=str, default=None)
+    parser.add_argument('-rm', '--ref_mesh', type=str)    
+    
+    parser.add_argument('-i', '--iter', type=int, default=1000)
+    parser.add_argument('-b', '--batch', type=int, default=8)
+    parser.add_argument('-r', '--train_res', nargs=2, type=int, default=[2048, 2048])
+    parser.add_argument('-lr', '--learning_rate', type=float, default=0.01)
+    parser.add_argument('--voxel_grid_res', type=int, default=64)
+    
+    parser.add_argument('--sdf_loss', type=bool, default=True)
+    parser.add_argument('--develop_reg', type=bool, default=False)
+    parser.add_argument('--sdf_regularizer', type=float, default=0.2)
+    
+    parser.add_argument('-dr', '--display_res', nargs=2, type=int, default=[512, 512])
+    parser.add_argument('-si', '--save_interval', type=int, default=20)
+    FLAGS = parser.parse_args()
+    device = 'cuda'
+    
+    os.makedirs(FLAGS.out_dir, exist_ok=True)
+    glctx = dr.RasterizeGLContext()
+    
+    # Load GT mesh
+    gt_mesh = load_mesh(FLAGS.ref_mesh, device)
+    gt_mesh.auto_normals() # compute face normals for visualization
+    
+    # ==============================================================================================
+    #  Create and initialize FlexiCubes
+    # ==============================================================================================
+    fc = FlexiCubes(device)
+    x_nx3, cube_fx8 = fc.construct_voxel_grid(FLAGS.voxel_grid_res)
+    x_nx3 *= 2 # scale up the grid so that it's larger than the target object
+    
+    sdf = torch.rand_like(x_nx3[:,0]) - 0.1 # randomly init SDF
+    sdf    = torch.nn.Parameter(sdf.clone().detach(), requires_grad=True)
+    # set per-cube learnable weights to zeros
+    weight = torch.zeros((cube_fx8.shape[0], 21), dtype=torch.float, device='cuda') 
+    weight    = torch.nn.Parameter(weight.clone().detach(), requires_grad=True)
+    deform = torch.nn.Parameter(torch.zeros_like(x_nx3), requires_grad=True)
+    
+    #  Retrieve all the edges of the voxel grid; these edges will be utilized to 
+    #  compute the regularization loss in subsequent steps of the process.    
+    all_edges = cube_fx8[:, fc.cube_edges].reshape(-1, 2)
+    grid_edges = torch.unique(all_edges, dim=0)
+    
+    # ==============================================================================================
+    #  Setup optimizer
+    # ==============================================================================================
+    optimizer = torch.optim.Adam([sdf, weight,deform], lr=FLAGS.learning_rate)
+    scheduler = torch.optim.lr_scheduler.LambdaLR(optimizer, lr_lambda=lambda x: lr_schedule(x)) 
+    
+    # ==============================================================================================
+    #  Train loop
+    # ==============================================================================================   
+    for it in range(FLAGS.iter): 
+        optimizer.zero_grad()
+        # sample random camera poses
+        mv, mvp = render.get_random_camera_batch(FLAGS.batch, iter_res=FLAGS.train_res, device=device, use_kaolin=False)
+        # render gt mesh
+        target = render.render_mesh_paper(gt_mesh, mv, mvp, FLAGS.train_res)
+        # extract and render FlexiCubes mesh
+        grid_verts = x_nx3 + (2-1e-8) / (FLAGS.voxel_grid_res * 2) * torch.tanh(deform)
+        vertices, faces, L_dev = fc(grid_verts, sdf, cube_fx8, FLAGS.voxel_grid_res, beta_fx12=weight[:,:12], alpha_fx8=weight[:,12:20],
+            gamma_f=weight[:,20], training=True)
+        flexicubes_mesh = Mesh(vertices, faces)
+        buffers = render.render_mesh_paper(flexicubes_mesh, mv, mvp, FLAGS.train_res)
+        
+        # evaluate reconstruction loss
+        mask_loss = (buffers['mask'] - target['mask']).abs().mean()
+        depth_loss = (((((buffers['depth'] - (target['depth']))* target['mask'])**2).sum(-1)+1e-8)).sqrt().mean() * 10
+    
+        t_iter = it / FLAGS.iter
+        sdf_weight = FLAGS.sdf_regularizer - (FLAGS.sdf_regularizer - FLAGS.sdf_regularizer/20)*min(1.0, 4.0 * t_iter)
+        reg_loss = loss.sdf_reg_loss(sdf, grid_edges).mean() * sdf_weight # Loss to eliminate internal floaters that are not visible
+        reg_loss += L_dev.mean() * 0.5
+        reg_loss += (weight[:,:20]).abs().mean() * 0.1
+        total_loss = mask_loss + depth_loss + reg_loss
+        
+        if FLAGS.sdf_loss: # optionally add SDF loss to eliminate internal structures
+            with torch.no_grad():
+                pts = sample_random_points(1000, gt_mesh)
+                gt_sdf = compute_sdf(pts, gt_mesh.vertices, gt_mesh.faces)
+            pred_sdf = compute_sdf(pts, flexicubes_mesh.vertices, flexicubes_mesh.faces)
+            total_loss += torch.nn.functional.mse_loss(pred_sdf, gt_sdf) * 2e3
+        
+        # optionally add developability regularizer, as described in paper section 5.2
+        if FLAGS.develop_reg:
+            reg_weight = max(0, t_iter - 0.8) * 5
+            if reg_weight > 0: # only applied after shape converges
+                reg_loss = loss.mesh_developable_reg(flexicubes_mesh).mean() * 10
+                reg_loss += (deform).abs().mean()
+                reg_loss += (weight[:,:20]).abs().mean()
+                total_loss = mask_loss + depth_loss + reg_loss 
+        
+        total_loss.backward()
+        optimizer.step()
+        scheduler.step()        
+        
+        if (it % FLAGS.save_interval == 0 or it == (FLAGS.iter-1)): # save normal image for visualization
+            with torch.no_grad():
+                # extract mesh with training=False
+                vertices, faces, L_dev = fc(grid_verts, sdf, cube_fx8, FLAGS.voxel_grid_res, beta_fx12=weight[:,:12], alpha_fx8=weight[:,12:20],
+                gamma_f=weight[:,20], training=False)
+                flexicubes_mesh = Mesh(vertices, faces)
+                
+                flexicubes_mesh.auto_normals() # compute face normals for visualization
+                mv, mvp = render.get_rotate_camera(it//FLAGS.save_interval, iter_res=FLAGS.display_res, device=device,use_kaolin=False)
+                val_buffers = render.render_mesh_paper(flexicubes_mesh, mv.unsqueeze(0), mvp.unsqueeze(0), FLAGS.display_res, return_types=["normal"], white_bg=True)
+                val_image = ((val_buffers["normal"][0].detach().cpu().numpy()+1)/2*255).astype(np.uint8)
+                
+                gt_buffers = render.render_mesh_paper(gt_mesh, mv.unsqueeze(0), mvp.unsqueeze(0), FLAGS.display_res, return_types=["normal"], white_bg=True)
+                gt_image = ((gt_buffers["normal"][0].detach().cpu().numpy()+1)/2*255).astype(np.uint8)
+                imageio.imwrite(os.path.join(FLAGS.out_dir, '{:04d}.png'.format(it)), np.concatenate([val_image, gt_image], 1))
+                print(f"Optimization Step [{it}/{FLAGS.iter}], Loss: {total_loss.item():.4f}")
+            
+    # ==============================================================================================
+    #  Save ouput
+    # ==============================================================================================     
+    mesh_np = trimesh.Trimesh(vertices = vertices.detach().cpu().numpy(), faces=faces.detach().cpu().numpy(), process=False)
+    mesh_np.export(os.path.join(FLAGS.out_dir, 'output_mesh.obj'))

trellis/representations/mesh/flexicubes/examples/render.py

@@ -0,0 +1,267 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
+import numpy as np
+import copy
+import math
+from ipywidgets import interactive, HBox, VBox, FloatLogSlider, IntSlider
+
+import torch
+import nvdiffrast.torch as dr
+import kaolin as kal
+import util
+
+###############################################################################
+# Functions adapted from https://github.com/NVlabs/nvdiffrec
+###############################################################################
+
+def get_random_camera_batch(batch_size, fovy = np.deg2rad(45), iter_res=[512,512], cam_near_far=[0.1, 1000.0], cam_radius=3.0, device="cuda", use_kaolin=True):
+    if use_kaolin:
+        camera_pos = torch.stack(kal.ops.coords.spherical2cartesian(
+            *kal.ops.random.sample_spherical_coords((batch_size,), azimuth_low=0., azimuth_high=math.pi * 2,
+                                                    elevation_low=-math.pi / 2., elevation_high=math.pi / 2., device='cuda'),
+            cam_radius
+        ), dim=-1)
+        return kal.render.camera.Camera.from_args(
+            eye=camera_pos + torch.rand((batch_size, 1), device='cuda') * 0.5 - 0.25,
+            at=torch.zeros(batch_size, 3),
+            up=torch.tensor([[0., 1., 0.]]),
+            fov=fovy,
+            near=cam_near_far[0], far=cam_near_far[1],
+            height=iter_res[0], width=iter_res[1],
+            device='cuda'
+        )
+    else:
+        def get_random_camera():
+            proj_mtx = util.perspective(fovy, iter_res[1] / iter_res[0], cam_near_far[0], cam_near_far[1])
+            mv     = util.translate(0, 0, -cam_radius) @ util.random_rotation_translation(0.25)
+            mvp    = proj_mtx @ mv
+            return mv, mvp
+        mv_batch = []
+        mvp_batch = []
+        for i in range(batch_size):
+            mv, mvp = get_random_camera()
+            mv_batch.append(mv)
+            mvp_batch.append(mvp)
+        return torch.stack(mv_batch).to(device), torch.stack(mvp_batch).to(device)
+
+def get_rotate_camera(itr, fovy = np.deg2rad(45), iter_res=[512,512], cam_near_far=[0.1, 1000.0], cam_radius=3.0, device="cuda", use_kaolin=True):
+    if use_kaolin:
+        ang = (itr / 10) * np.pi * 2
+        camera_pos = torch.stack(kal.ops.coords.spherical2cartesian(torch.tensor(ang), torch.tensor(0.4), -torch.tensor(cam_radius)))
+        return kal.render.camera.Camera.from_args(
+            eye=camera_pos,
+            at=torch.zeros(3),
+            up=torch.tensor([0., 1., 0.]),
+            fov=fovy,
+            near=cam_near_far[0], far=cam_near_far[1],
+            height=iter_res[0], width=iter_res[1],
+            device='cuda'
+        )
+    else:
+        proj_mtx = util.perspective(fovy, iter_res[1] / iter_res[0], cam_near_far[0], cam_near_far[1])
+
+        # Smooth rotation for display.
+        ang    = (itr / 10) * np.pi * 2
+        mv     = util.translate(0, 0, -cam_radius) @ (util.rotate_x(-0.4) @ util.rotate_y(ang))
+        mvp    = proj_mtx @ mv
+        return mv.to(device), mvp.to(device)
+
+glctx = dr.RasterizeGLContext()
+def render_mesh(mesh, camera, iter_res, return_types = ["mask", "depth"], white_bg=False, wireframe_thickness=0.4):
+    vertices_camera = camera.extrinsics.transform(mesh.vertices)
+    face_vertices_camera = kal.ops.mesh.index_vertices_by_faces(
+        vertices_camera, mesh.faces
+    )
+
+    # Projection: nvdiffrast take clip coordinates as input to apply barycentric perspective correction.
+    # Using `camera.intrinsics.transform(vertices_camera) would return the normalized device coordinates.
+    proj = camera.projection_matrix().unsqueeze(1)
+    proj[:, :, 1, 1] = -proj[:, :, 1, 1]
+    homogeneous_vecs = kal.render.camera.up_to_homogeneous(
+        vertices_camera
+    )
+    vertices_clip = (proj @ homogeneous_vecs.unsqueeze(-1)).squeeze(-1)
+    faces_int = mesh.faces.int()
+
+    rast, _ = dr.rasterize(
+        glctx, vertices_clip, faces_int, iter_res)
+
+    out_dict = {}
+    for type in return_types:
+        if type == "mask" :
+            img = dr.antialias((rast[..., -1:] > 0).float(), rast, vertices_clip, faces_int)
+        elif type == "depth":
+            img = dr.interpolate(homogeneous_vecs, rast, faces_int)[0]
+        elif type == "wireframe":
+            img = torch.logical_or(
+                torch.logical_or(rast[..., 0] < wireframe_thickness, rast[..., 1] < wireframe_thickness),
+                (rast[..., 0] + rast[..., 1]) > (1. - wireframe_thickness)
+            ).unsqueeze(-1)
+        elif type == "normals" :
+            img = dr.interpolate(
+                mesh.face_normals.reshape(len(mesh), -1, 3), rast,
+                torch.arange(mesh.faces.shape[0] * 3, device='cuda', dtype=torch.int).reshape(-1, 3)
+            )[0]
+        if white_bg:
+            bg = torch.ones_like(img)
+            alpha = (rast[..., -1:] > 0).float() 
+            img = torch.lerp(bg, img, alpha)
+        out_dict[type] = img
+
+        
+    return out_dict
+
+def render_mesh_paper(mesh, mv, mvp, iter_res, return_types = ["mask", "depth"], white_bg=False):
+    '''
+    The rendering function used to produce the results in the paper.
+    '''
+    v_pos_clip = util.xfm_points(mesh.vertices.unsqueeze(0), mvp)  # Rotate it to camera coordinates
+    rast, db = dr.rasterize(
+        dr.RasterizeGLContext(), v_pos_clip, mesh.faces.int(), iter_res)
+
+    out_dict = {}
+    for type in return_types:
+        if type == "mask" :
+            img = dr.antialias((rast[..., -1:] > 0).float(), rast, v_pos_clip, mesh.faces.int()) 
+        elif type == "depth":
+            v_pos_cam = util.xfm_points(mesh.vertices.unsqueeze(0), mv)
+            img, _ = util.interpolate(v_pos_cam, rast, mesh.faces.int())
+        elif type == "normal" :
+            normal_indices = (torch.arange(0, mesh.nrm.shape[0], dtype=torch.int64, device='cuda')[:, None]).repeat(1, 3)
+            img, _ = util.interpolate(mesh.nrm.unsqueeze(0).contiguous(), rast, normal_indices.int())
+        elif type == "vertex_normal":
+            img, _ = util.interpolate(mesh.v_nrm.unsqueeze(0).contiguous(), rast, mesh.faces.int())
+            img = dr.antialias((img + 1) * 0.5, rast, v_pos_clip, mesh.faces.int()) 
+        if white_bg:
+            bg = torch.ones_like(img)
+            alpha = (rast[..., -1:] > 0).float() 
+            img = torch.lerp(bg, img, alpha)
+        out_dict[type] = img
+    return out_dict
+
+class SplitVisualizer():
+    def __init__(self, lh_mesh, rh_mesh, height, width):
+        self.lh_mesh = lh_mesh
+        self.rh_mesh = rh_mesh
+        self.height = height
+        self.width = width
+        self.wireframe_thickness = 0.4
+        
+
+    def render(self, camera):
+        lh_outputs = render_mesh(
+            self.lh_mesh, camera, (self.height, self.width),
+            return_types=["normals", "wireframe"], wireframe_thickness=self.wireframe_thickness
+        )
+        rh_outputs = render_mesh(
+            self.rh_mesh, camera, (self.height, self.width),
+            return_types=["normals", "wireframe"], wireframe_thickness=self.wireframe_thickness
+        )
+        outputs = {
+            k: torch.cat(
+                [lh_outputs[k][0].permute(1, 0, 2), rh_outputs[k][0].permute(1, 0, 2)],
+                dim=0
+            ).permute(1, 0, 2) for k in ["normals", "wireframe"]
+        }
+        return {
+            'img': (outputs['wireframe'] * ((outputs['normals'] + 1.) / 2.) * 255).to(torch.uint8),
+            'normals': outputs['normals']
+        }
+
+    def show(self, init_camera):
+        visualizer = kal.visualize.IpyTurntableVisualizer(
+            self.height, self.width * 2, copy.deepcopy(init_camera), self.render,
+            max_fps=24, world_up_axis=1)
+
+        def slider_callback(new_wireframe_thickness):
+            """ipywidgets sliders callback"""
+            with visualizer.out: # This is in case of bug
+                self.wireframe_thickness = new_wireframe_thickness
+                # this is how we request a new update
+                visualizer.render_update()
+                
+        wireframe_thickness_slider = FloatLogSlider(
+            value=self.wireframe_thickness,
+            base=10,
+            min=-3,
+            max=-0.4,
+            step=0.1,
+            description='wireframe_thickness',
+            continuous_update=True,
+            readout=True,
+            readout_format='.3f',
+        )
+        
+        interactive_slider = interactive(
+            slider_callback,
+            new_wireframe_thickness=wireframe_thickness_slider,
+        )
+        
+        full_output = VBox([visualizer.canvas, interactive_slider])
+        display(full_output, visualizer.out)
+
+class TimelineVisualizer():
+    def __init__(self, meshes, height, width):
+        self.meshes = meshes
+        self.height = height
+        self.width = width
+        self.wireframe_thickness = 0.4
+        self.idx = len(meshes) - 1
+
+    def render(self, camera):
+        outputs = render_mesh(
+            self.meshes[self.idx], camera, (self.height, self.width),
+            return_types=["normals", "wireframe"], wireframe_thickness=self.wireframe_thickness
+        )
+
+        return {
+            'img': (outputs['wireframe'] * ((outputs['normals'] + 1.) / 2.) * 255).to(torch.uint8)[0],
+            'normals': outputs['normals'][0]
+        }
+
+    def show(self, init_camera):
+        visualizer = kal.visualize.IpyTurntableVisualizer(
+            self.height, self.width, copy.deepcopy(init_camera), self.render,
+            max_fps=24, world_up_axis=1)
+
+        def slider_callback(new_wireframe_thickness, new_idx):
+            """ipywidgets sliders callback"""
+            with visualizer.out: # This is in case of bug
+                self.wireframe_thickness = new_wireframe_thickness
+                self.idx = new_idx
+                # this is how we request a new update
+                visualizer.render_update()
+
+        wireframe_thickness_slider = FloatLogSlider(
+            value=self.wireframe_thickness,
+            base=10,
+            min=-3,
+            max=-0.4,
+            step=0.1,
+            description='wireframe_thickness',
+            continuous_update=True,
+            readout=True,
+            readout_format='.3f',
+        )
+
+        idx_slider = IntSlider(
+            value=self.idx,
+            min=0,
+            max=len(self.meshes) - 1,
+            description='idx',
+            continuous_update=True,
+            readout=True
+        )
+
+        interactive_slider = interactive(
+            slider_callback,
+            new_wireframe_thickness=wireframe_thickness_slider,
+            new_idx=idx_slider
+        )
+        full_output = HBox([visualizer.canvas, interactive_slider])
+        display(full_output, visualizer.out)

trellis/representations/mesh/flexicubes/examples/util.py

@@ -0,0 +1,122 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
+import numpy as np
+import torch
+import trimesh
+import kaolin
+import nvdiffrast.torch as dr
+
+###############################################################################
+# Functions adapted from https://github.com/NVlabs/nvdiffrec
+###############################################################################
+
+def dot(x: torch.Tensor, y: torch.Tensor) -> torch.Tensor:
+    return torch.sum(x*y, -1, keepdim=True)
+
+def length(x: torch.Tensor, eps: float =1e-8) -> torch.Tensor:
+    return torch.sqrt(torch.clamp(dot(x,x), min=eps)) # Clamp to avoid nan gradients because grad(sqrt(0)) = NaN
+
+def safe_normalize(x: torch.Tensor, eps: float =1e-8) -> torch.Tensor:
+    return x / length(x, eps)
+
+def perspective(fovy=0.7854, aspect=1.0, n=0.1, f=1000.0, device=None):
+    y = np.tan(fovy / 2)
+    return torch.tensor([[1/(y*aspect),    0,            0,              0], 
+                         [           0, 1/-y,            0,              0], 
+                         [           0,    0, -(f+n)/(f-n), -(2*f*n)/(f-n)], 
+                         [           0,    0,           -1,              0]], dtype=torch.float32, device=device)
+
+def translate(x, y, z, device=None):
+    return torch.tensor([[1, 0, 0, x], 
+                         [0, 1, 0, y], 
+                         [0, 0, 1, z], 
+                         [0, 0, 0, 1]], dtype=torch.float32, device=device)
+
+@torch.no_grad()
+def random_rotation_translation(t, device=None):
+    m = np.random.normal(size=[3, 3])
+    m[1] = np.cross(m[0], m[2])
+    m[2] = np.cross(m[0], m[1])
+    m = m / np.linalg.norm(m, axis=1, keepdims=True)
+    m = np.pad(m, [[0, 1], [0, 1]], mode='constant')
+    m[3, 3] = 1.0
+    m[:3, 3] = np.random.uniform(-t, t, size=[3])
+    return torch.tensor(m, dtype=torch.float32, device=device)
+
+def rotate_x(a, device=None):
+    s, c = np.sin(a), np.cos(a)
+    return torch.tensor([[1,  0, 0, 0], 
+                         [0,  c, s, 0], 
+                         [0, -s, c, 0], 
+                         [0,  0, 0, 1]], dtype=torch.float32, device=device)
+
+def rotate_y(a, device=None):
+    s, c = np.sin(a), np.cos(a)
+    return torch.tensor([[ c, 0, s, 0], 
+                         [ 0, 1, 0, 0], 
+                         [-s, 0, c, 0], 
+                         [ 0, 0, 0, 1]], dtype=torch.float32, device=device)
+    
+class Mesh:
+    def __init__(self, vertices, faces):
+        self.vertices = vertices
+        self.faces = faces
+        
+    def auto_normals(self):
+        v0 = self.vertices[self.faces[:, 0], :]
+        v1 = self.vertices[self.faces[:, 1], :]
+        v2 = self.vertices[self.faces[:, 2], :]
+        nrm = safe_normalize(torch.cross(v1 - v0, v2 - v0))
+        self.nrm = nrm
+
+def load_mesh(path, device):
+    mesh_np = trimesh.load(path)
+    vertices = torch.tensor(mesh_np.vertices, device=device, dtype=torch.float)
+    faces = torch.tensor(mesh_np.faces, device=device, dtype=torch.long)
+    
+    # Normalize
+    vmin, vmax = vertices.min(dim=0)[0], vertices.max(dim=0)[0]
+    scale = 1.8 / torch.max(vmax - vmin).item()
+    vertices = vertices - (vmax + vmin) / 2 # Center mesh on origin
+    vertices = vertices * scale # Rescale to [-0.9, 0.9]
+    return Mesh(vertices, faces)
+
+def compute_sdf(points, vertices, faces):
+    face_vertices = kaolin.ops.mesh.index_vertices_by_faces(vertices.clone().unsqueeze(0), faces)
+    distance = kaolin.metrics.trianglemesh.point_to_mesh_distance(points.unsqueeze(0), face_vertices)[0]
+    with torch.no_grad():
+        sign = (kaolin.ops.mesh.check_sign(vertices.unsqueeze(0), faces, points.unsqueeze(0))<1).float() * 2 - 1
+    sdf = (sign*distance).squeeze(0)
+    return sdf
+
+def sample_random_points(n, mesh):
+    pts_random = (torch.rand((n//2,3),device='cuda') - 0.5) * 2
+    pts_surface = kaolin.ops.mesh.sample_points(mesh.vertices.unsqueeze(0), mesh.faces, 500)[0].squeeze(0)
+    pts_surface += torch.randn_like(pts_surface) * 0.05
+    pts = torch.cat([pts_random, pts_surface])
+    return pts
+
+def xfm_points(points, matrix):
+    '''Transform points.
+    Args:
+        points: Tensor containing 3D points with shape [minibatch_size, num_vertices, 3] or [1, num_vertices, 3]
+        matrix: A 4x4 transform matrix with shape [minibatch_size, 4, 4]
+        use_python: Use PyTorch's torch.matmul (for validation)
+    Returns:
+        Transformed points in homogeneous 4D with shape [minibatch_size, num_vertices, 4].
+    '''
+    out = torch.matmul(
+        torch.nn.functional.pad(points, pad=(0, 1), mode='constant', value=1.0), torch.transpose(matrix, 1, 2))
+    if torch.is_anomaly_enabled():
+        assert torch.all(torch.isfinite(out)), "Output of xfm_points contains inf or NaN"
+    return out
+
+def interpolate(attr, rast, attr_idx, rast_db=None):
+    return dr.interpolate(
+        attr, rast, attr_idx, rast_db=rast_db,
+        diff_attrs=None if rast_db is None else 'all')

trellis/representations/mesh/flexicubes/flexicubes.py

@@ -0,0 +1,384 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
+
+import torch
+from .tables import *
+from kaolin.utils.testing import check_tensor
+
+__all__ = [
+    'FlexiCubes'
+]
+
+
+class FlexiCubes:
+    def __init__(self, device="cuda"):
+
+        self.device = device
+        self.dmc_table = torch.tensor(dmc_table, dtype=torch.long, device=device, requires_grad=False)
+        self.num_vd_table = torch.tensor(num_vd_table,
+                                         dtype=torch.long, device=device, requires_grad=False)
+        self.check_table = torch.tensor(
+            check_table,
+            dtype=torch.long, device=device, requires_grad=False)
+
+        self.tet_table = torch.tensor(tet_table, dtype=torch.long, device=device, requires_grad=False)
+        self.quad_split_1 = torch.tensor([0, 1, 2, 0, 2, 3], dtype=torch.long, device=device, requires_grad=False)
+        self.quad_split_2 = torch.tensor([0, 1, 3, 3, 1, 2], dtype=torch.long, device=device, requires_grad=False)
+        self.quad_split_train = torch.tensor(
+            [0, 1, 1, 2, 2, 3, 3, 0], dtype=torch.long, device=device, requires_grad=False)
+
+        self.cube_corners = torch.tensor([[0, 0, 0], [1, 0, 0], [0, 1, 0], [1, 1, 0], [0, 0, 1], [
+                                         1, 0, 1], [0, 1, 1], [1, 1, 1]], dtype=torch.float, device=device)
+        self.cube_corners_idx = torch.pow(2, torch.arange(8, requires_grad=False))
+        self.cube_edges = torch.tensor([0, 1, 1, 5, 4, 5, 0, 4, 2, 3, 3, 7, 6, 7, 2, 6,
+                                       2, 0, 3, 1, 7, 5, 6, 4], dtype=torch.long, device=device, requires_grad=False)
+
+        self.edge_dir_table = torch.tensor([0, 2, 0, 2, 0, 2, 0, 2, 1, 1, 1, 1],
+                                           dtype=torch.long, device=device)
+        self.dir_faces_table = torch.tensor([
+            [[5, 4], [3, 2], [4, 5], [2, 3]],
+            [[5, 4], [1, 0], [4, 5], [0, 1]],
+            [[3, 2], [1, 0], [2, 3], [0, 1]]
+        ], dtype=torch.long, device=device)
+        self.adj_pairs = torch.tensor([0, 1, 1, 3, 3, 2, 2, 0], dtype=torch.long, device=device)
+
+    def __call__(self, voxelgrid_vertices, scalar_field, cube_idx, resolution, qef_reg_scale=1e-3,
+                 weight_scale=0.99, beta=None, alpha=None, gamma_f=None, voxelgrid_colors=None, training=False):
+        assert torch.is_tensor(voxelgrid_vertices) and \
+            check_tensor(voxelgrid_vertices, (None, 3), throw=False), \
+            "'voxelgrid_vertices' should be a tensor of shape (num_vertices, 3)"
+        num_vertices = voxelgrid_vertices.shape[0]
+        assert torch.is_tensor(scalar_field) and \
+            check_tensor(scalar_field, (num_vertices,), throw=False), \
+            "'scalar_field' should be a tensor of shape (num_vertices,)"
+        assert torch.is_tensor(cube_idx) and \
+            check_tensor(cube_idx, (None, 8), throw=False), \
+            "'cube_idx' should be a tensor of shape (num_cubes, 8)"
+        num_cubes = cube_idx.shape[0]
+        assert beta is None or (
+            torch.is_tensor(beta) and
+            check_tensor(beta, (num_cubes, 12), throw=False)
+        ), "'beta' should be a tensor of shape (num_cubes, 12)"
+        assert alpha is None or (
+            torch.is_tensor(alpha) and
+            check_tensor(alpha, (num_cubes, 8), throw=False)
+        ), "'alpha' should be a tensor of shape (num_cubes, 8)"
+        assert gamma_f is None or (
+            torch.is_tensor(gamma_f) and
+            check_tensor(gamma_f, (num_cubes,), throw=False)
+        ), "'gamma_f' should be a tensor of shape (num_cubes,)"
+
+        surf_cubes, occ_fx8 = self._identify_surf_cubes(scalar_field, cube_idx)
+        if surf_cubes.sum() == 0:
+            return (
+                torch.zeros((0, 3), device=self.device),
+                torch.zeros((0, 3), dtype=torch.long, device=self.device),
+                torch.zeros((0), device=self.device),
+                torch.zeros((0, voxelgrid_colors.shape[-1]), device=self.device) if voxelgrid_colors is not None else None
+            )
+        beta, alpha, gamma_f = self._normalize_weights(
+            beta, alpha, gamma_f, surf_cubes, weight_scale)
+        
+        if voxelgrid_colors is not None:
+            voxelgrid_colors = torch.sigmoid(voxelgrid_colors)
+
+        case_ids = self._get_case_id(occ_fx8, surf_cubes, resolution)
+
+        surf_edges, idx_map, edge_counts, surf_edges_mask = self._identify_surf_edges(
+            scalar_field, cube_idx, surf_cubes
+        )
+
+        vd, L_dev, vd_gamma, vd_idx_map, vd_color = self._compute_vd(
+            voxelgrid_vertices, cube_idx[surf_cubes], surf_edges, scalar_field,
+            case_ids, beta, alpha, gamma_f, idx_map, qef_reg_scale, voxelgrid_colors)
+        vertices, faces, s_edges, edge_indices, vertices_color = self._triangulate(
+            scalar_field, surf_edges, vd, vd_gamma, edge_counts, idx_map,
+            vd_idx_map, surf_edges_mask, training, vd_color)
+        return vertices, faces, L_dev, vertices_color
+
+    def _compute_reg_loss(self, vd, ue, edge_group_to_vd, vd_num_edges):
+        """
+        Regularizer L_dev as in Equation 8
+        """
+        dist = torch.norm(ue - torch.index_select(input=vd, index=edge_group_to_vd, dim=0), dim=-1)
+        mean_l2 = torch.zeros_like(vd[:, 0])
+        mean_l2 = (mean_l2).index_add_(0, edge_group_to_vd, dist) / vd_num_edges.squeeze(1).float()
+        mad = (dist - torch.index_select(input=mean_l2, index=edge_group_to_vd, dim=0)).abs()
+        return mad
+
+    def _normalize_weights(self, beta, alpha, gamma_f, surf_cubes, weight_scale):
+        """
+        Normalizes the given weights to be non-negative. If input weights are None, it creates and returns a set of weights of ones.
+        """
+        n_cubes = surf_cubes.shape[0]
+
+        if beta is not None:
+            beta = (torch.tanh(beta) * weight_scale + 1)
+        else:
+            beta = torch.ones((n_cubes, 12), dtype=torch.float, device=self.device)
+
+        if alpha is not None:
+            alpha = (torch.tanh(alpha) * weight_scale + 1)
+        else:
+            alpha = torch.ones((n_cubes, 8), dtype=torch.float, device=self.device)
+
+        if gamma_f is not None:
+            gamma_f = torch.sigmoid(gamma_f) * weight_scale + (1 - weight_scale) / 2
+        else:
+            gamma_f = torch.ones((n_cubes), dtype=torch.float, device=self.device)
+
+        return beta[surf_cubes], alpha[surf_cubes], gamma_f[surf_cubes]
+
+    @torch.no_grad()
+    def _get_case_id(self, occ_fx8, surf_cubes, res):
+        """
+        Obtains the ID of topology cases based on cell corner occupancy. This function resolves the 
+        ambiguity in the Dual Marching Cubes (DMC) configurations as described in Section 1.3 of the 
+        supplementary material. It should be noted that this function assumes a regular grid.
+        """
+        case_ids = (occ_fx8[surf_cubes] * self.cube_corners_idx.to(self.device).unsqueeze(0)).sum(-1)
+
+        problem_config = self.check_table.to(self.device)[case_ids]
+        to_check = problem_config[..., 0] == 1
+        problem_config = problem_config[to_check]
+        if not isinstance(res, (list, tuple)):
+            res = [res, res, res]
+
+        # The 'problematic_configs' only contain configurations for surface cubes. Next, we construct a 3D array,
+        # 'problem_config_full', to store configurations for all cubes (with default config for non-surface cubes).
+        # This allows efficient checking on adjacent cubes.
+        problem_config_full = torch.zeros(list(res) + [5], device=self.device, dtype=torch.long)
+        vol_idx = torch.nonzero(problem_config_full[..., 0] == 0)  # N, 3
+        vol_idx_problem = vol_idx[surf_cubes][to_check]
+        problem_config_full[vol_idx_problem[..., 0], vol_idx_problem[..., 1], vol_idx_problem[..., 2]] = problem_config
+        vol_idx_problem_adj = vol_idx_problem + problem_config[..., 1:4]
+
+        within_range = (
+            vol_idx_problem_adj[..., 0] >= 0) & (
+            vol_idx_problem_adj[..., 0] < res[0]) & (
+            vol_idx_problem_adj[..., 1] >= 0) & (
+            vol_idx_problem_adj[..., 1] < res[1]) & (
+            vol_idx_problem_adj[..., 2] >= 0) & (
+            vol_idx_problem_adj[..., 2] < res[2])
+
+        vol_idx_problem = vol_idx_problem[within_range]
+        vol_idx_problem_adj = vol_idx_problem_adj[within_range]
+        problem_config = problem_config[within_range]
+        problem_config_adj = problem_config_full[vol_idx_problem_adj[..., 0],
+                                                 vol_idx_problem_adj[..., 1], vol_idx_problem_adj[..., 2]]
+        # If two cubes with cases C16 and C19 share an ambiguous face, both cases are inverted.
+        to_invert = (problem_config_adj[..., 0] == 1)
+        idx = torch.arange(case_ids.shape[0], device=self.device)[to_check][within_range][to_invert]
+        case_ids.index_put_((idx,), problem_config[to_invert][..., -1])
+        return case_ids
+
+    @torch.no_grad()
+    def _identify_surf_edges(self, scalar_field, cube_idx, surf_cubes):
+        """
+        Identifies grid edges that intersect with the underlying surface by checking for opposite signs. As each edge 
+        can be shared by multiple cubes, this function also assigns a unique index to each surface-intersecting edge 
+        and marks the cube edges with this index.
+        """
+        occ_n = scalar_field < 0
+        all_edges = cube_idx[surf_cubes][:, self.cube_edges].reshape(-1, 2)
+        unique_edges, _idx_map, counts = torch.unique(all_edges, dim=0, return_inverse=True, return_counts=True)
+
+        unique_edges = unique_edges.long()
+        mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 1
+
+        surf_edges_mask = mask_edges[_idx_map]
+        counts = counts[_idx_map]
+
+        mapping = torch.ones((unique_edges.shape[0]), dtype=torch.long, device=cube_idx.device) * -1
+        mapping[mask_edges] = torch.arange(mask_edges.sum(), device=cube_idx.device)
+        # Shaped as [number of cubes x 12 edges per cube]. This is later used to map a cube edge to the unique index
+        # for a surface-intersecting edge. Non-surface-intersecting edges are marked with -1.
+        idx_map = mapping[_idx_map]
+        surf_edges = unique_edges[mask_edges]
+        return surf_edges, idx_map, counts, surf_edges_mask
+
+    @torch.no_grad()
+    def _identify_surf_cubes(self, scalar_field, cube_idx):
+        """
+        Identifies grid cubes that intersect with the underlying surface by checking if the signs at 
+        all corners are not identical.
+        """
+        occ_n = scalar_field < 0
+        occ_fx8 = occ_n[cube_idx.reshape(-1)].reshape(-1, 8)
+        _occ_sum = torch.sum(occ_fx8, -1)
+        surf_cubes = (_occ_sum > 0) & (_occ_sum < 8)
+        return surf_cubes, occ_fx8
+
+    def _linear_interp(self, edges_weight, edges_x):
+        """
+        Computes the location of zero-crossings on 'edges_x' using linear interpolation with 'edges_weight'.
+        """
+        edge_dim = edges_weight.dim() - 2
+        assert edges_weight.shape[edge_dim] == 2
+        edges_weight = torch.cat([torch.index_select(input=edges_weight, index=torch.tensor(1, device=self.device), dim=edge_dim), -
+                                 torch.index_select(input=edges_weight, index=torch.tensor(0, device=self.device), dim=edge_dim)]
+                                 , edge_dim)
+        denominator = edges_weight.sum(edge_dim)
+        ue = (edges_x * edges_weight).sum(edge_dim) / denominator
+        return ue
+
+    def _solve_vd_QEF(self, p_bxnx3, norm_bxnx3, c_bx3, qef_reg_scale):
+        p_bxnx3 = p_bxnx3.reshape(-1, 7, 3)
+        norm_bxnx3 = norm_bxnx3.reshape(-1, 7, 3)
+        c_bx3 = c_bx3.reshape(-1, 3)
+        A = norm_bxnx3
+        B = ((p_bxnx3) * norm_bxnx3).sum(-1, keepdims=True)
+
+        A_reg = (torch.eye(3, device=p_bxnx3.device) * qef_reg_scale).unsqueeze(0).repeat(p_bxnx3.shape[0], 1, 1)
+        B_reg = (qef_reg_scale * c_bx3).unsqueeze(-1)
+        A = torch.cat([A, A_reg], 1)
+        B = torch.cat([B, B_reg], 1)
+        dual_verts = torch.linalg.lstsq(A, B).solution.squeeze(-1)
+        return dual_verts
+
+    def _compute_vd(self, voxelgrid_vertices, surf_cubes_fx8, surf_edges, scalar_field,
+                    case_ids, beta, alpha, gamma_f, idx_map, qef_reg_scale, voxelgrid_colors):
+        """
+        Computes the location of dual vertices as described in Section 4.2
+        """
+        alpha_nx12x2 = torch.index_select(input=alpha, index=self.cube_edges, dim=1).reshape(-1, 12, 2)
+        surf_edges_x = torch.index_select(input=voxelgrid_vertices, index=surf_edges.reshape(-1), dim=0).reshape(-1, 2, 3)
+        surf_edges_s = torch.index_select(input=scalar_field, index=surf_edges.reshape(-1), dim=0).reshape(-1, 2, 1)
+        zero_crossing = self._linear_interp(surf_edges_s, surf_edges_x)
+        
+        if voxelgrid_colors is not None:
+            C = voxelgrid_colors.shape[-1]
+            surf_edges_c = torch.index_select(input=voxelgrid_colors, index=surf_edges.reshape(-1), dim=0).reshape(-1, 2, C)
+
+        idx_map = idx_map.reshape(-1, 12)
+        num_vd = torch.index_select(input=self.num_vd_table, index=case_ids, dim=0)
+        edge_group, edge_group_to_vd, edge_group_to_cube, vd_num_edges, vd_gamma = [], [], [], [], []
+        
+        # if color is not None:
+        #     vd_color = []
+
+        total_num_vd = 0
+        vd_idx_map = torch.zeros((case_ids.shape[0], 12), dtype=torch.long, device=self.device, requires_grad=False)
+
+        for num in torch.unique(num_vd):
+            cur_cubes = (num_vd == num)  # consider cubes with the same numbers of vd emitted (for batching)
+            curr_num_vd = cur_cubes.sum() * num
+            curr_edge_group = self.dmc_table[case_ids[cur_cubes], :num].reshape(-1, num * 7)
+            curr_edge_group_to_vd = torch.arange(
+                curr_num_vd, device=self.device).unsqueeze(-1).repeat(1, 7) + total_num_vd
+            total_num_vd += curr_num_vd
+            curr_edge_group_to_cube = torch.arange(idx_map.shape[0], device=self.device)[
+                cur_cubes].unsqueeze(-1).repeat(1, num * 7).reshape_as(curr_edge_group)
+
+            curr_mask = (curr_edge_group != -1)
+            edge_group.append(torch.masked_select(curr_edge_group, curr_mask))
+            edge_group_to_vd.append(torch.masked_select(curr_edge_group_to_vd.reshape_as(curr_edge_group), curr_mask))
+            edge_group_to_cube.append(torch.masked_select(curr_edge_group_to_cube, curr_mask))
+            vd_num_edges.append(curr_mask.reshape(-1, 7).sum(-1, keepdims=True))
+            vd_gamma.append(torch.masked_select(gamma_f, cur_cubes).unsqueeze(-1).repeat(1, num).reshape(-1))
+            # if color is not None:
+            #     vd_color.append(color[cur_cubes].unsqueeze(1).repeat(1, num, 1).reshape(-1, 3))
+            
+        edge_group = torch.cat(edge_group)
+        edge_group_to_vd = torch.cat(edge_group_to_vd)
+        edge_group_to_cube = torch.cat(edge_group_to_cube)
+        vd_num_edges = torch.cat(vd_num_edges)
+        vd_gamma = torch.cat(vd_gamma)
+        # if color is not None:
+        #     vd_color = torch.cat(vd_color)
+        # else:
+        #     vd_color = None
+
+        vd = torch.zeros((total_num_vd, 3), device=self.device)
+        beta_sum = torch.zeros((total_num_vd, 1), device=self.device)
+
+        idx_group = torch.gather(input=idx_map.reshape(-1), dim=0, index=edge_group_to_cube * 12 + edge_group)
+
+        x_group = torch.index_select(input=surf_edges_x, index=idx_group.reshape(-1), dim=0).reshape(-1, 2, 3)
+        s_group = torch.index_select(input=surf_edges_s, index=idx_group.reshape(-1), dim=0).reshape(-1, 2, 1)
+        
+
+        zero_crossing_group = torch.index_select(
+            input=zero_crossing, index=idx_group.reshape(-1), dim=0).reshape(-1, 3)
+
+        alpha_group = torch.index_select(input=alpha_nx12x2.reshape(-1, 2), dim=0,
+                                            index=edge_group_to_cube * 12 + edge_group).reshape(-1, 2, 1)
+        ue_group = self._linear_interp(s_group * alpha_group, x_group)
+
+        beta_group = torch.gather(input=beta.reshape(-1), dim=0,
+                                    index=edge_group_to_cube * 12 + edge_group).reshape(-1, 1)
+        beta_sum = beta_sum.index_add_(0, index=edge_group_to_vd, source=beta_group)
+        vd = vd.index_add_(0, index=edge_group_to_vd, source=ue_group * beta_group) / beta_sum
+        
+        '''
+        interpolate colors use the same method as dual vertices
+        '''
+        if voxelgrid_colors is not None:
+            vd_color = torch.zeros((total_num_vd, C), device=self.device)
+            c_group = torch.index_select(input=surf_edges_c, index=idx_group.reshape(-1), dim=0).reshape(-1, 2, C)
+            uc_group = self._linear_interp(s_group * alpha_group, c_group)
+            vd_color = vd_color.index_add_(0, index=edge_group_to_vd, source=uc_group * beta_group) / beta_sum
+        else:
+            vd_color = None
+        
+        L_dev = self._compute_reg_loss(vd, zero_crossing_group, edge_group_to_vd, vd_num_edges)
+
+        v_idx = torch.arange(vd.shape[0], device=self.device)  # + total_num_vd
+
+        vd_idx_map = (vd_idx_map.reshape(-1)).scatter(dim=0, index=edge_group_to_cube *
+                                                      12 + edge_group, src=v_idx[edge_group_to_vd])
+
+        return vd, L_dev, vd_gamma, vd_idx_map, vd_color
+
+    def _triangulate(self, scalar_field, surf_edges, vd, vd_gamma, edge_counts, idx_map, vd_idx_map, surf_edges_mask, training, vd_color):
+        """
+        Connects four neighboring dual vertices to form a quadrilateral. The quadrilaterals are then split into 
+        triangles based on the gamma parameter, as described in Section 4.3.
+        """
+        with torch.no_grad():
+            group_mask = (edge_counts == 4) & surf_edges_mask  # surface edges shared by 4 cubes.
+            group = idx_map.reshape(-1)[group_mask]
+            vd_idx = vd_idx_map[group_mask]
+            edge_indices, indices = torch.sort(group, stable=True)
+            quad_vd_idx = vd_idx[indices].reshape(-1, 4)
+
+            # Ensure all face directions point towards the positive SDF to maintain consistent winding.
+            s_edges = scalar_field[surf_edges[edge_indices.reshape(-1, 4)[:, 0]].reshape(-1)].reshape(-1, 2)
+            flip_mask = s_edges[:, 0] > 0
+            quad_vd_idx = torch.cat((quad_vd_idx[flip_mask][:, [0, 1, 3, 2]],
+                                     quad_vd_idx[~flip_mask][:, [2, 3, 1, 0]]))
+
+        quad_gamma = torch.index_select(input=vd_gamma, index=quad_vd_idx.reshape(-1), dim=0).reshape(-1, 4)
+        gamma_02 = quad_gamma[:, 0] * quad_gamma[:, 2]
+        gamma_13 = quad_gamma[:, 1] * quad_gamma[:, 3]
+        if not training:
+            mask = (gamma_02 > gamma_13)
+            faces = torch.zeros((quad_gamma.shape[0], 6), dtype=torch.long, device=quad_vd_idx.device)
+            faces[mask] = quad_vd_idx[mask][:, self.quad_split_1]
+            faces[~mask] = quad_vd_idx[~mask][:, self.quad_split_2]
+            faces = faces.reshape(-1, 3)
+        else:
+            vd_quad = torch.index_select(input=vd, index=quad_vd_idx.reshape(-1), dim=0).reshape(-1, 4, 3)
+            vd_02 = (vd_quad[:, 0] + vd_quad[:, 2]) / 2
+            vd_13 = (vd_quad[:, 1] + vd_quad[:, 3]) / 2
+            weight_sum = (gamma_02 + gamma_13) + 1e-8
+            vd_center = (vd_02 * gamma_02.unsqueeze(-1) + vd_13 * gamma_13.unsqueeze(-1)) / weight_sum.unsqueeze(-1)
+            
+            if vd_color is not None:
+                color_quad = torch.index_select(input=vd_color, index=quad_vd_idx.reshape(-1), dim=0).reshape(-1, 4, vd_color.shape[-1])
+                color_02 = (color_quad[:, 0] + color_quad[:, 2]) / 2
+                color_13 = (color_quad[:, 1] + color_quad[:, 3]) / 2
+                color_center = (color_02 * gamma_02.unsqueeze(-1) + color_13 * gamma_13.unsqueeze(-1)) / weight_sum.unsqueeze(-1)
+                vd_color = torch.cat([vd_color, color_center])
+            
+            
+            vd_center_idx = torch.arange(vd_center.shape[0], device=self.device) + vd.shape[0]
+            vd = torch.cat([vd, vd_center])
+            faces = quad_vd_idx[:, self.quad_split_train].reshape(-1, 4, 2)
+            faces = torch.cat([faces, vd_center_idx.reshape(-1, 1, 1).repeat(1, 4, 1)], -1).reshape(-1, 3)
+        return vd, faces, s_edges, edge_indices, vd_color

trellis/representations/mesh/flexicubes/tables.py

@@ -0,0 +1,791 @@
+# Copyright (c) 2023, NVIDIA CORPORATION & AFFILIATES.  All rights reserved.
+#
+# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
+dmc_table = [
+[[-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 5, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 5, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 7, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 5, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 5, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 5, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 4, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 4, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 8, 11, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 5, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 4, 5, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 7, 8, 9, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 5, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 5, 7, 8, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 5, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 7, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 9, 10, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 4, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 4, 5, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 4, 5, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 7, 8, 9, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 5, 7, 9, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 5, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 5, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[8, 9, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 8, -1, -1, -1, -1], [1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 7, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 9, 10, 11, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 9, -1, -1, -1, -1], [1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 8, 10, 11, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 5, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 7, 8, 9, -1, -1, -1], [1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 5, 7, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 5, 7, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 8, 9, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 5, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 5, 8, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 6, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 5, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 5, 6, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 5, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 6, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 6, 7, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 4, 6, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 4, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [2, 3, 4, 6, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 4, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 9, -1, -1, -1, -1], [2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 6, 7, 8, -1, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 5, -1, -1, -1], [2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 4, 5, 6, 7, 8], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 5, 6, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 5, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 2, 3, 5, 6, 8], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 9, 10, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 8, 9, 10, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 6, 8, 11, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 6, 11, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 9, 10, -1, -1, -1], [4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 4, 6, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [4, 5, 9, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1]],
+[[0, 2, 4, 5, 10, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 4, 5, 8, 10, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 6, 8, 9, 11, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 5, 6, 9, 11, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 5, 6, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 5, 6, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 6, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 6, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 6, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[6, 7, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 6, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 6, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 6, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 9, -1, -1, -1, -1], [1, 3, 6, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 6, 7, 8, 10, -1], [4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 5, 6, 7, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 6, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 5, 6, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 5, 6, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 8, 9, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 7, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 7, 9, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 6, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 6, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[6, 7, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 6, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 6, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 6, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 8, 11, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 8, 9, 11, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 4, 7, 11, -1, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [4, 7, 8, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [5, 6, 10, -1, -1, -1, -1]],
+[[1, 2, 4, 7, 9, 11, -1], [5, 6, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 6, 9, 10, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 8, 11, -1, -1, -1], [4, 6, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 6, 10, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 4, 6, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[6, 7, 8, 9, 10, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 6, 7, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 6, 7, 8, 10, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 6, 7, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 5, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 5, 6, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 8, -1, -1, -1, -1], [1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 7, -1, -1, -1], [1, 2, 5, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 5, 6, 9, -1, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 4, 5, 6, 7, 9], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 4, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [1, 2, 4, 6, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 4, 6, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 4, 6, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 6, 7, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 2, 3, 6, 7, 9], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 6, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 6, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 5, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 5, 6, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 5, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 6, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 8, -1, -1, -1, -1], [1, 3, 5, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 5, 6, 7, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 5, 6, 9, 11, -1], [4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 6, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 6, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 6, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 6, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 6, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 6, 7, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 6, 7, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[6, 7, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 8, 9, -1, -1, -1], [5, 7, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 5, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [4, 5, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 5, 9, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [4, 7, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 7, 10, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 7, 8, 10, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[8, 9, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 9, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 8, 10, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 10, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 5, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 5, 7, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [2, 3, 5, 7, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 5, 7, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 4, 5, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 4, 5, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [2, 3, 4, 5, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 4, 5, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 4, 7, 9, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 4, 7, 8, 9, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 2, 3, 4, 7, 10], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 8, -1, -1, -1, -1], [1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 8, 9, 10, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 9, 10, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 2, 3, 8, 10, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 10, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 5, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [1, 2, 5, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 5, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 5, 7, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 4, 5, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 2, 3, 4, 5, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 4, 5, 8, 9, 11], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 9, -1, -1, -1, -1], [2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 4, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [1, 2, 4, 7, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 4, 7, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 4, 7, 8, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 2, 8, 9, 11, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 2, 3, 9, 11, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 2, 8, 11, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[2, 3, 11, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 5, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 5, 7, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 5, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[5, 7, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 5, 8, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 5, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 5, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 5, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 4, 7, 9, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 4, 7, 8, 9, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 4, 7, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[4, 7, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[1, 3, 8, 9, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 1, 9, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[0, 3, 8, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]],
+[[-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1], [-1, -1, -1, -1, -1, -1, -1]]
+]
+num_vd_table = [0, 1, 1, 1, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 2, 1, 3, 1, 2, 2,
+2, 1, 2, 1, 2, 1, 1, 2, 1, 1, 2, 2, 2, 1, 2, 3, 1, 1, 2, 2, 1, 1, 1, 1, 1, 1, 2,
+1, 2, 1, 2, 2, 1, 1, 2, 1, 1, 1, 1, 2, 2, 2, 1, 1, 2, 1, 2, 3, 2, 2, 1, 1, 1, 1,
+1, 1, 2, 1, 1, 1, 2, 1, 2, 2, 2, 1, 1, 1, 1, 1, 2, 3, 2, 2, 2, 2, 2, 1, 3, 4, 2,
+2, 2, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 2, 1, 1, 1, 1, 2, 1, 1, 2, 2, 2, 2, 2,
+3, 2, 1, 2, 1, 1, 1, 1, 1, 1, 2, 2, 3, 2, 3, 2, 4, 2, 2, 2, 2, 1, 2, 1, 2, 1, 1,
+2, 1, 1, 2, 2, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 2, 1, 1, 1, 1, 1,
+1, 2, 1, 1, 1, 2, 2, 2, 1, 1, 2, 1, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1, 2, 1, 1, 1, 2,
+1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 2, 1, 1, 1, 1, 1, 2, 1, 2, 1, 1, 1, 1, 1, 1, 1, 1,
+1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0]
+check_table = [
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 1, 0, 0, 194],
+[1, -1, 0, 0, 193],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, 1, 0, 164],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, -1, 0, 161],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, 0, 1, 152],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, 0, 1, 145],
+[1, 0, 0, 1, 144],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, 0, -1, 137],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, 1, 0, 133],
+[1, 0, 1, 0, 132],
+[1, 1, 0, 0, 131],
+[1, 1, 0, 0, 130],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, 0, 1, 100],
+[0, 0, 0, 0, 0],
+[1, 0, 0, 1, 98],
+[0, 0, 0, 0, 0],
+[1, 0, 0, 1, 96],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, 1, 0, 88],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, -1, 0, 82],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, 1, 0, 74],
+[0, 0, 0, 0, 0],
+[1, 0, 1, 0, 72],
+[0, 0, 0, 0, 0],
+[1, 0, 0, -1, 70],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, -1, 0, 0, 67],
+[0, 0, 0, 0, 0],
+[1, -1, 0, 0, 65],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 1, 0, 0, 56],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, -1, 0, 0, 52],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 1, 0, 0, 44],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 1, 0, 0, 40],
+[0, 0, 0, 0, 0],
+[1, 0, 0, -1, 38],
+[1, 0, -1, 0, 37],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, -1, 0, 33],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, -1, 0, 0, 28],
+[0, 0, 0, 0, 0],
+[1, 0, -1, 0, 26],
+[1, 0, 0, -1, 25],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, -1, 0, 0, 20],
+[0, 0, 0, 0, 0],
+[1, 0, -1, 0, 18],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, 0, -1, 9],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[1, 0, 0, -1, 6],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0]
+]
+tet_table = [
+[-1, -1, -1, -1, -1, -1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[1, 1, 1, 1, 1, 1],
+[4, 4, 4, 4, 4, 4],
+[0, 0, 0, 0, 0, 0],
+[4, 0, 0, 4, 4, -1],
+[1, 1, 1, 1, 1, 1],
+[4, 4, 4, 4, 4, 4],
+[0, 4, 0, 4, 4, -1],
+[0, 0, 0, 0, 0, 0],
+[1, 1, 1, 1, 1, 1],
+[5, 5, 5, 5, 5, 5],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[1, 1, 1, 1, 1, 1],
+[2, 2, 2, 2, 2, 2],
+[0, 0, 0, 0, 0, 0],
+[2, 0, 2, -1, 0, 2],
+[1, 1, 1, 1, 1, 1],
+[2, -1, 2, 4, 4, 2],
+[0, 0, 0, 0, 0, 0],
+[2, 0, 2, 4, 4, 2],
+[1, 1, 1, 1, 1, 1],
+[2, 4, 2, 4, 4, 2],
+[0, 4, 0, 4, 4, 0],
+[2, 0, 2, 0, 0, 2],
+[1, 1, 1, 1, 1, 1],
+[2, 5, 2, 5, 5, 2],
+[0, 0, 0, 0, 0, 0],
+[2, 0, 2, 0, 0, 2],
+[1, 1, 1, 1, 1, 1],
+[1, 1, 1, 1, 1, 1],
+[0, 1, 1, -1, 0, 1],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[4, 1, 1, 4, 4, 1],
+[0, 1, 1, 0, 0, 1],
+[4, 0, 0, 4, 4, 0],
+[2, 2, 2, 2, 2, 2],
+[-1, 1, 1, 4, 4, 1],
+[0, 1, 1, 4, 4, 1],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[5, 1, 1, 5, 5, 1],
+[0, 1, 1, 0, 0, 1],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[8, 8, 8, 8, 8, 8],
+[1, 1, 1, 4, 4, 1],
+[0, 0, 0, 0, 0, 0],
+[4, 0, 0, 4, 4, 0],
+[4, 4, 4, 4, 4, 4],
+[1, 1, 1, 4, 4, 1],
+[0, 4, 0, 4, 4, 0],
+[0, 0, 0, 0, 0, 0],
+[4, 4, 4, 4, 4, 4],
+[1, 1, 1, 5, 5, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[5, 5, 5, 5, 5, 5],
+[6, 6, 6, 6, 6, 6],
+[6, -1, 0, 6, 0, 6],
+[6, 0, 0, 6, 0, 6],
+[6, 1, 1, 6, 1, 6],
+[4, 4, 4, 4, 4, 4],
+[0, 0, 0, 0, 0, 0],
+[4, 0, 0, 4, 4, 4],
+[1, 1, 1, 1, 1, 1],
+[6, 4, -1, 6, 4, 6],
+[6, 4, 0, 6, 4, 6],
+[6, 0, 0, 6, 0, 6],
+[6, 1, 1, 6, 1, 6],
+[5, 5, 5, 5, 5, 5],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[1, 1, 1, 1, 1, 1],
+[2, 2, 2, 2, 2, 2],
+[0, 0, 0, 0, 0, 0],
+[2, 0, 2, 2, 0, 2],
+[1, 1, 1, 1, 1, 1],
+[2, 2, 2, 2, 2, 2],
+[0, 0, 0, 0, 0, 0],
+[2, 0, 2, 2, 2, 2],
+[1, 1, 1, 1, 1, 1],
+[2, 4, 2, 2, 4, 2],
+[0, 4, 0, 4, 4, 0],
+[2, 0, 2, 2, 0, 2],
+[1, 1, 1, 1, 1, 1],
+[2, 2, 2, 2, 2, 2],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[1, 1, 1, 1, 1, 1],
+[6, 1, 1, 6, -1, 6],
+[6, 1, 1, 6, 0, 6],
+[6, 0, 0, 6, 0, 6],
+[6, 2, 2, 6, 2, 6],
+[4, 1, 1, 4, 4, 1],
+[0, 1, 1, 0, 0, 1],
+[4, 0, 0, 4, 4, 4],
+[2, 2, 2, 2, 2, 2],
+[6, 1, 1, 6, 4, 6],
+[6, 1, 1, 6, 4, 6],
+[6, 0, 0, 6, 0, 6],
+[6, 2, 2, 6, 2, 6],
+[5, 1, 1, 5, 5, 1],
+[0, 1, 1, 0, 0, 1],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[6, 6, 6, 6, 6, 6],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[4, 4, 4, 4, 4, 4],
+[1, 1, 1, 1, 4, 1],
+[0, 4, 0, 4, 4, 0],
+[0, 0, 0, 0, 0, 0],
+[4, 4, 4, 4, 4, 4],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 5, 0, 5, 0, 5],
+[5, 5, 5, 5, 5, 5],
+[5, 5, 5, 5, 5, 5],
+[0, 5, 0, 5, 0, 5],
+[-1, 5, 0, 5, 0, 5],
+[1, 5, 1, 5, 1, 5],
+[4, 5, -1, 5, 4, 5],
+[0, 5, 0, 5, 0, 5],
+[4, 5, 0, 5, 4, 5],
+[1, 5, 1, 5, 1, 5],
+[4, 4, 4, 4, 4, 4],
+[0, 4, 0, 4, 4, 4],
+[0, 0, 0, 0, 0, 0],
+[1, 1, 1, 1, 1, 1],
+[6, 6, 6, 6, 6, 6],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[1, 1, 1, 1, 1, 1],
+[2, 5, 2, 5, -1, 5],
+[0, 5, 0, 5, 0, 5],
+[2, 5, 2, 5, 0, 5],
+[1, 5, 1, 5, 1, 5],
+[2, 5, 2, 5, 4, 5],
+[0, 5, 0, 5, 0, 5],
+[2, 5, 2, 5, 4, 5],
+[1, 5, 1, 5, 1, 5],
+[2, 4, 2, 4, 4, 2],
+[0, 4, 0, 4, 4, 4],
+[2, 0, 2, 0, 0, 2],
+[1, 1, 1, 1, 1, 1],
+[2, 6, 2, 6, 6, 2],
+[0, 0, 0, 0, 0, 0],
+[2, 0, 2, 0, 0, 2],
+[1, 1, 1, 1, 1, 1],
+[1, 1, 1, 1, 1, 1],
+[0, 1, 1, 1, 0, 1],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[4, 1, 1, 1, 4, 1],
+[0, 1, 1, 1, 0, 1],
+[4, 0, 0, 4, 4, 0],
+[2, 2, 2, 2, 2, 2],
+[1, 1, 1, 1, 1, 1],
+[0, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[5, 5, 5, 5, 5, 5],
+[1, 1, 1, 1, 4, 1],
+[0, 0, 0, 0, 0, 0],
+[4, 0, 0, 4, 4, 0],
+[4, 4, 4, 4, 4, 4],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[4, 4, 4, 4, 4, 4],
+[1, 1, 1, 1, 1, 1],
+[6, 0, 0, 6, 0, 6],
+[0, 0, 0, 0, 0, 0],
+[6, 6, 6, 6, 6, 6],
+[5, 5, 5, 5, 5, 5],
+[5, 5, 0, 5, 0, 5],
+[5, 5, 0, 5, 0, 5],
+[5, 5, 1, 5, 1, 5],
+[4, 4, 4, 4, 4, 4],
+[0, 0, 0, 0, 0, 0],
+[4, 4, 0, 4, 4, 4],
+[1, 1, 1, 1, 1, 1],
+[4, 4, 4, 4, 4, 4],
+[4, 4, 0, 4, 4, 4],
+[0, 0, 0, 0, 0, 0],
+[1, 1, 1, 1, 1, 1],
+[8, 8, 8, 8, 8, 8],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[1, 1, 1, 1, 1, 1],
+[2, 2, 2, 2, 2, 2],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 0, 2],
+[1, 1, 1, 1, 1, 1],
+[2, 2, 2, 2, 2, 2],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[1, 1, 1, 1, 1, 1],
+[2, 2, 2, 2, 2, 2],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[4, 1, 1, 4, 4, 1],
+[2, 2, 2, 2, 2, 2],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[1, 1, 1, 1, 1, 1],
+[1, 1, 1, 1, 1, 1],
+[1, 1, 1, 1, 0, 1],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[2, 4, 2, 4, 4, 2],
+[1, 1, 1, 1, 1, 1],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[2, 2, 2, 2, 2, 2],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[5, 5, 5, 5, 5, 5],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[4, 4, 4, 4, 4, 4],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[4, 4, 4, 4, 4, 4],
+[1, 1, 1, 1, 1, 1],
+[0, 0, 0, 0, 0, 0],
+[0, 0, 0, 0, 0, 0],
+[12, 12, 12, 12, 12, 12]
+]

trellis/representations/octree/octree_dfs.py

@@ -3,21 +3,6 @@ import torch.nn as nn
 import torch.nn.functional as F
 
 
-DEFAULT_TRIVEC_CONFIG = {
-    'dim': 8,
-    'rank': 8,
-}
-
-DEFAULT_VOXEL_CONFIG = {
-    'solid': False,
-}
-
-DEFAULT_DECOPOLY_CONFIG = {
-    'degree': 8,
-    'rank': 16,
-}
-
-
 class DfsOctree:
     """
     Sparse Voxel Octree (SVO) implementation for PyTorch.
@@ -145,8 +130,8 @@ class DfsOctree:
 
     @property
     def get_density(self):
-        if self.primitive == 'voxel' and self.voxel_config['solid']:
-            return torch.full((self.position.shape[0], 1), 1000, dtype=torch.float32, device=self.device)
+        if self.primitive == 'voxel' and self.primitive_config.get('solid', False):
+            return torch.full((self.position.shape[0], 1), torch.finfo(torch.float32).max, dtype=torch.float32, device=self.device)
         return self.density_activation(self.density)
     
     @property
@@ -172,7 +157,7 @@ class DfsOctree:
         return torch.cat([self.features_dc, self.features_ac], dim=-2)
 
     def state_dict(self):
-        ret = {'structure': self.structure, 'position': self.position, 'depth': self.depth, 'sh_degree': self.sh_degree, 'active_sh_degree': self.active_sh_degree, 'trivec_config': self.trivec_config, 'voxel_config': self.voxel_config, 'primitive': self.primitive}
+        ret = {'structure': self.structure, 'position': self.position, 'depth': self.depth, 'sh_degree': self.sh_degree, 'active_sh_degree': self.active_sh_degree, 'primitive_config': self.primitive_config, 'primitive': self.primitive}
         if hasattr(self, 'density_shift'):
             ret['density_shift'] = self.density_shift
         for data in set(self.data + self.param_names):

trellis/trainers/__init__.py

@@ -0,0 +1,63 @@
+import importlib
+
+__attributes = {
+    'BasicTrainer': 'basic',
+    
+    'SparseStructureVaeTrainer': 'vae.sparse_structure_vae',
+    
+    'SLatVaeGaussianTrainer': 'vae.structured_latent_vae_gaussian',
+    'SLatVaeRadianceFieldDecoderTrainer': 'vae.structured_latent_vae_rf_dec',
+    'SLatVaeMeshDecoderTrainer': 'vae.structured_latent_vae_mesh_dec',
+    
+    'FlowMatchingTrainer': 'flow_matching.flow_matching',
+    'FlowMatchingCFGTrainer': 'flow_matching.flow_matching',
+    'TextConditionedFlowMatchingCFGTrainer': 'flow_matching.flow_matching',
+    'ImageConditionedFlowMatchingCFGTrainer': 'flow_matching.flow_matching',
+    
+    'SparseFlowMatchingTrainer': 'flow_matching.sparse_flow_matching',
+    'SparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
+    'TextConditionedSparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
+    'ImageConditionedSparseFlowMatchingCFGTrainer': 'flow_matching.sparse_flow_matching',
+}
+
+__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 .basic import BasicTrainer
+
+    from .vae.sparse_structure_vae import SparseStructureVaeTrainer
+
+    from .vae.structured_latent_vae_gaussian import SLatVaeGaussianTrainer
+    from .vae.structured_latent_vae_rf_dec import SLatVaeRadianceFieldDecoderTrainer
+    from .vae.structured_latent_vae_mesh_dec import SLatVaeMeshDecoderTrainer
+    
+    from .flow_matching.flow_matching import (
+        FlowMatchingTrainer,
+        FlowMatchingCFGTrainer,
+        TextConditionedFlowMatchingCFGTrainer,
+        ImageConditionedFlowMatchingCFGTrainer,
+    )
+    
+    from .flow_matching.sparse_flow_matching import (
+        SparseFlowMatchingTrainer,
+        SparseFlowMatchingCFGTrainer,
+        TextConditionedSparseFlowMatchingCFGTrainer,
+        ImageConditionedSparseFlowMatchingCFGTrainer,
+    )

trellis/trainers/base.py

@@ -0,0 +1,451 @@
+from abc import abstractmethod
+import os
+import time
+import json
+
+import torch
+import torch.distributed as dist
+from torch.utils.data import DataLoader
+import numpy as np
+
+from torchvision import utils
+from torch.utils.tensorboard import SummaryWriter
+
+from .utils import *
+from ..utils.general_utils import *
+from ..utils.data_utils import recursive_to_device, cycle, ResumableSampler
+
+
+class Trainer:
+    """
+    Base class for training.
+    """
+    def __init__(self,
+        models,
+        dataset,
+        *,
+        output_dir,
+        load_dir,
+        step,
+        max_steps,
+        batch_size=None,
+        batch_size_per_gpu=None,
+        batch_split=None,
+        optimizer={},
+        lr_scheduler=None,
+        elastic=None,
+        grad_clip=None,
+        ema_rate=0.9999,
+        fp16_mode='inflat_all',
+        fp16_scale_growth=1e-3,
+        finetune_ckpt=None,
+        log_param_stats=False,
+        prefetch_data=True,
+        i_print=1000,
+        i_log=500,
+        i_sample=10000,
+        i_save=10000,
+        i_ddpcheck=10000,
+        **kwargs
+    ):
+        assert batch_size is not None or batch_size_per_gpu is not None, 'Either batch_size or batch_size_per_gpu must be specified.'
+
+        self.models = models
+        self.dataset = dataset
+        self.batch_split = batch_split if batch_split is not None else 1
+        self.max_steps = max_steps
+        self.optimizer_config = optimizer
+        self.lr_scheduler_config = lr_scheduler
+        self.elastic_controller_config = elastic
+        self.grad_clip = grad_clip
+        self.ema_rate = [ema_rate] if isinstance(ema_rate, float) else ema_rate
+        self.fp16_mode = fp16_mode
+        self.fp16_scale_growth = fp16_scale_growth
+        self.log_param_stats = log_param_stats
+        self.prefetch_data = prefetch_data
+        if self.prefetch_data:
+            self._data_prefetched = None
+
+        self.output_dir = output_dir
+        self.i_print = i_print
+        self.i_log = i_log
+        self.i_sample = i_sample
+        self.i_save = i_save
+        self.i_ddpcheck = i_ddpcheck        
+
+        if dist.is_initialized():
+            # Multi-GPU params
+            self.world_size = dist.get_world_size()
+            self.rank = dist.get_rank()
+            self.local_rank = dist.get_rank() % torch.cuda.device_count()
+            self.is_master = self.rank == 0
+        else:
+            # Single-GPU params
+            self.world_size = 1
+            self.rank = 0
+            self.local_rank = 0
+            self.is_master = True
+
+        self.batch_size = batch_size if batch_size_per_gpu is None else batch_size_per_gpu * self.world_size
+        self.batch_size_per_gpu = batch_size_per_gpu if batch_size_per_gpu is not None else batch_size // self.world_size
+        assert self.batch_size % self.world_size == 0, 'Batch size must be divisible by the number of GPUs.'
+        assert self.batch_size_per_gpu % self.batch_split == 0, 'Batch size per GPU must be divisible by batch split.'
+
+        self.init_models_and_more(**kwargs)
+        self.prepare_dataloader(**kwargs)
+        
+        # Load checkpoint
+        self.step = 0
+        if load_dir is not None and step is not None:
+            self.load(load_dir, step)
+        elif finetune_ckpt is not None:
+            self.finetune_from(finetune_ckpt)
+        
+        if self.is_master:
+            os.makedirs(os.path.join(self.output_dir, 'ckpts'), exist_ok=True)
+            os.makedirs(os.path.join(self.output_dir, 'samples'), exist_ok=True)
+            self.writer = SummaryWriter(os.path.join(self.output_dir, 'tb_logs'))
+
+        if self.world_size > 1:
+            self.check_ddp()
+            
+        if self.is_master:
+            print('\n\nTrainer initialized.')
+            print(self)
+            
+    @property
+    def device(self):
+        for _, model in self.models.items():
+            if hasattr(model, 'device'):
+                return model.device
+        return next(list(self.models.values())[0].parameters()).device
+            
+    @abstractmethod
+    def init_models_and_more(self, **kwargs):
+        """
+        Initialize models and more.
+        """
+        pass
+    
+    def prepare_dataloader(self, **kwargs):
+        """
+        Prepare dataloader.
+        """
+        self.data_sampler = ResumableSampler(
+            self.dataset,
+            shuffle=True,
+        )
+        self.dataloader = DataLoader(
+            self.dataset,
+            batch_size=self.batch_size_per_gpu,
+            num_workers=int(np.ceil(os.cpu_count() / torch.cuda.device_count())),
+            pin_memory=True,
+            drop_last=True,
+            persistent_workers=True,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+            sampler=self.data_sampler,
+        )
+        self.data_iterator = cycle(self.dataloader)
+
+    @abstractmethod
+    def load(self, load_dir, step=0):
+        """
+        Load a checkpoint.
+        Should be called by all processes.
+        """
+        pass
+
+    @abstractmethod
+    def save(self):
+        """
+        Save a checkpoint.
+        Should be called only by the rank 0 process.
+        """
+        pass
+    
+    @abstractmethod
+    def finetune_from(self, finetune_ckpt):
+        """
+        Finetune from a checkpoint.
+        Should be called by all processes.
+        """
+        pass
+    
+    @abstractmethod
+    def run_snapshot(self, num_samples, batch_size=4, verbose=False, **kwargs):
+        """
+        Run a snapshot of the model.
+        """
+        pass
+
+    @torch.no_grad()
+    def visualize_sample(self, sample):
+        """
+        Convert a sample to an image.
+        """
+        if hasattr(self.dataset, 'visualize_sample'):
+            return self.dataset.visualize_sample(sample)
+        else:
+            return sample
+
+    @torch.no_grad()
+    def snapshot_dataset(self, num_samples=100):
+        """
+        Sample images from the dataset.
+        """
+        dataloader = torch.utils.data.DataLoader(
+            self.dataset,
+            batch_size=num_samples,
+            num_workers=0,
+            shuffle=True,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+        )
+        data = next(iter(dataloader))
+        data = recursive_to_device(data, self.device)
+        vis = self.visualize_sample(data)
+        if isinstance(vis, dict):
+            save_cfg = [(f'dataset_{k}', v) for k, v in vis.items()]
+        else:
+            save_cfg = [('dataset', vis)]
+        for name, image in save_cfg:
+            utils.save_image(
+                image,
+                os.path.join(self.output_dir, 'samples', f'{name}.jpg'),
+                nrow=int(np.sqrt(num_samples)),
+                normalize=True,
+                value_range=self.dataset.value_range,
+            )
+
+    @torch.no_grad()
+    def snapshot(self, suffix=None, num_samples=64, batch_size=4, verbose=False):
+        """
+        Sample images from the model.
+        NOTE: This function should be called by all processes.
+        """
+        if self.is_master:
+            print(f'\nSampling {num_samples} images...', end='')
+
+        if suffix is None:
+            suffix = f'step{self.step:07d}'
+
+        # Assign tasks
+        num_samples_per_process = int(np.ceil(num_samples / self.world_size))
+        samples = self.run_snapshot(num_samples_per_process, batch_size=batch_size, verbose=verbose)
+
+        # Preprocess images
+        for key in list(samples.keys()):
+            if samples[key]['type'] == 'sample':
+                vis = self.visualize_sample(samples[key]['value'])
+                if isinstance(vis, dict):
+                    for k, v in vis.items():
+                        samples[f'{key}_{k}'] = {'value': v, 'type': 'image'}
+                    del samples[key]
+                else:
+                    samples[key] = {'value': vis, 'type': 'image'}
+
+        # Gather results
+        if self.world_size > 1:
+            for key in samples.keys():
+                samples[key]['value'] = samples[key]['value'].contiguous()
+                if self.is_master:
+                    all_images = [torch.empty_like(samples[key]['value']) for _ in range(self.world_size)]
+                else:
+                    all_images = []
+                dist.gather(samples[key]['value'], all_images, dst=0)
+                if self.is_master:
+                    samples[key]['value'] = torch.cat(all_images, dim=0)[:num_samples]
+
+        # Save images
+        if self.is_master:
+            os.makedirs(os.path.join(self.output_dir, 'samples', suffix), exist_ok=True)
+            for key in samples.keys():
+                if samples[key]['type'] == 'image':
+                    utils.save_image(
+                        samples[key]['value'],
+                        os.path.join(self.output_dir, 'samples', suffix, f'{key}_{suffix}.jpg'),
+                        nrow=int(np.sqrt(num_samples)),
+                        normalize=True,
+                        value_range=self.dataset.value_range,
+                    )
+                elif samples[key]['type'] == 'number':
+                    min = samples[key]['value'].min()
+                    max = samples[key]['value'].max()
+                    images = (samples[key]['value'] - min) / (max - min)
+                    images = utils.make_grid(
+                        images,
+                        nrow=int(np.sqrt(num_samples)),
+                        normalize=False,
+                    )
+                    save_image_with_notes(
+                        images,
+                        os.path.join(self.output_dir, 'samples', suffix, f'{key}_{suffix}.jpg'),
+                        notes=f'{key} min: {min}, max: {max}',
+                    )
+
+        if self.is_master:
+            print(' Done.')
+
+    @abstractmethod
+    def update_ema(self):
+        """
+        Update exponential moving average.
+        Should only be called by the rank 0 process.
+        """
+        pass
+
+    @abstractmethod
+    def check_ddp(self):
+        """
+        Check if DDP is working properly.
+        Should be called by all process.
+        """
+        pass
+
+    @abstractmethod
+    def training_losses(**mb_data):
+        """
+        Compute training losses.
+        """
+        pass
+    
+    def load_data(self):
+        """
+        Load data.
+        """
+        if self.prefetch_data:
+            if self._data_prefetched is None:
+                self._data_prefetched = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True)
+            data = self._data_prefetched
+            self._data_prefetched = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True)
+        else:
+            data = recursive_to_device(next(self.data_iterator), self.device, non_blocking=True)
+        
+        # if the data is a dict, we need to split it into multiple dicts with batch_size_per_gpu
+        if isinstance(data, dict):
+            if self.batch_split == 1:
+                data_list = [data]
+            else:
+                batch_size = list(data.values())[0].shape[0]
+                data_list = [
+                    {k: v[i * batch_size // self.batch_split:(i + 1) * batch_size // self.batch_split] for k, v in data.items()}
+                    for i in range(self.batch_split)
+                ]
+        elif isinstance(data, list):
+            data_list = data
+        else:
+            raise ValueError('Data must be a dict or a list of dicts.')
+        
+        return data_list
+
+    @abstractmethod
+    def run_step(self, data_list):
+        """
+        Run a training step.
+        """
+        pass
+
+    def run(self):
+        """
+        Run training.
+        """
+        if self.is_master:
+            print('\nStarting training...')
+            self.snapshot_dataset()
+        if self.step == 0:
+            self.snapshot(suffix='init')
+        else: # resume
+            self.snapshot(suffix=f'resume_step{self.step:07d}')
+
+        log = []
+        time_last_print = 0.0
+        time_elapsed = 0.0
+        while self.step < self.max_steps:
+            time_start = time.time()
+
+            data_list = self.load_data()
+            step_log = self.run_step(data_list)
+
+            time_end = time.time()
+            time_elapsed += time_end - time_start
+
+            self.step += 1
+
+            # Print progress
+            if self.is_master and self.step % self.i_print == 0:
+                speed = self.i_print / (time_elapsed - time_last_print) * 3600
+                columns = [
+                    f'Step: {self.step}/{self.max_steps} ({self.step / self.max_steps * 100:.2f}%)',
+                    f'Elapsed: {time_elapsed / 3600:.2f} h',
+                    f'Speed: {speed:.2f} steps/h',
+                    f'ETA: {(self.max_steps - self.step) / speed:.2f} h',
+                ]
+                print(' | '.join([c.ljust(25) for c in columns]), flush=True)
+                time_last_print = time_elapsed
+
+            # Check ddp
+            if self.world_size > 1 and self.i_ddpcheck is not None and self.step % self.i_ddpcheck == 0:
+                self.check_ddp()
+
+            # Sample images
+            if self.step % self.i_sample == 0:
+                self.snapshot()
+
+            if self.is_master:
+                log.append((self.step, {}))
+
+                # Log time
+                log[-1][1]['time'] = {
+                    'step': time_end - time_start,
+                    'elapsed': time_elapsed,
+                }
+
+                # Log losses
+                if step_log is not None:
+                    log[-1][1].update(step_log)
+
+                # Log scale
+                if self.fp16_mode == 'amp':
+                    log[-1][1]['scale'] = self.scaler.get_scale()
+                elif self.fp16_mode == 'inflat_all':
+                    log[-1][1]['log_scale'] = self.log_scale
+
+                # Save log
+                if self.step % self.i_log == 0:
+                    ## save to log file
+                    log_str = '\n'.join([
+                        f'{step}: {json.dumps(log)}' for step, log in log
+                    ])
+                    with open(os.path.join(self.output_dir, 'log.txt'), 'a') as log_file:
+                        log_file.write(log_str + '\n')
+
+                    # show with mlflow
+                    log_show = [l for _, l in log if not dict_any(l, lambda x: np.isnan(x))]
+                    log_show = dict_reduce(log_show, lambda x: np.mean(x))
+                    log_show = dict_flatten(log_show, sep='/')
+                    for key, value in log_show.items():
+                        self.writer.add_scalar(key, value, self.step)
+                    log = []
+
+                # Save checkpoint
+                if self.step % self.i_save == 0:
+                    self.save()
+
+        if self.is_master:
+            self.snapshot(suffix='final')
+            self.writer.close()
+            print('Training finished.')
+            
+    def profile(self, wait=2, warmup=3, active=5):
+        """
+        Profile the training loop.
+        """
+        with torch.profiler.profile(
+            schedule=torch.profiler.schedule(wait=wait, warmup=warmup, active=active, repeat=1),
+            on_trace_ready=torch.profiler.tensorboard_trace_handler(os.path.join(self.output_dir, 'profile')),
+            profile_memory=True,
+            with_stack=True,
+        ) as prof:
+            for _ in range(wait + warmup + active):
+                self.run_step()
+                prof.step()
+            

trellis/trainers/basic.py

@@ -0,0 +1,438 @@
+import os
+import copy
+from functools import partial
+from contextlib import nullcontext
+
+import torch
+import torch.distributed as dist
+from torch.nn.parallel import DistributedDataParallel as DDP
+import numpy as np
+
+from .utils import *
+from .base import Trainer
+from ..utils.general_utils import *
+from ..utils.dist_utils import *
+from ..utils import grad_clip_utils, elastic_utils
+
+
+class BasicTrainer(Trainer):
+    """
+    Trainer for basic training loop.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+    """
+
+    def __str__(self):
+        lines = []
+        lines.append(self.__class__.__name__)
+        lines.append(f'  - Models:')
+        for name, model in self.models.items():
+            lines.append(f'    - {name}: {model.__class__.__name__}')
+        lines.append(f'  - Dataset: {indent(str(self.dataset), 2)}')
+        lines.append(f'  - Dataloader:')
+        lines.append(f'    - Sampler: {self.dataloader.sampler.__class__.__name__}')
+        lines.append(f'    - Num workers: {self.dataloader.num_workers}')
+        lines.append(f'  - Number of steps: {self.max_steps}')
+        lines.append(f'  - Number of GPUs: {self.world_size}')
+        lines.append(f'  - Batch size: {self.batch_size}')
+        lines.append(f'  - Batch size per GPU: {self.batch_size_per_gpu}')
+        lines.append(f'  - Batch split: {self.batch_split}')
+        lines.append(f'  - Optimizer: {self.optimizer.__class__.__name__}')
+        lines.append(f'  - Learning rate: {self.optimizer.param_groups[0]["lr"]}')
+        if self.lr_scheduler_config is not None:
+            lines.append(f'  - LR scheduler: {self.lr_scheduler.__class__.__name__}')
+        if self.elastic_controller_config is not None:
+            lines.append(f'  - Elastic memory: {indent(str(self.elastic_controller), 2)}')
+        if self.grad_clip is not None:
+            lines.append(f'  - Gradient clip: {indent(str(self.grad_clip), 2)}')
+        lines.append(f'  - EMA rate: {self.ema_rate}')
+        lines.append(f'  - FP16 mode: {self.fp16_mode}')
+        return '\n'.join(lines)
+            
+    def init_models_and_more(self, **kwargs):
+        """
+        Initialize models and more.
+        """
+        if self.world_size > 1:
+            # Prepare distributed data parallel
+            self.training_models = {
+                name: DDP(
+                    model,
+                    device_ids=[self.local_rank],
+                    output_device=self.local_rank,
+                    bucket_cap_mb=128,
+                    find_unused_parameters=False
+                )
+                for name, model in self.models.items()
+            }
+        else:
+            self.training_models = self.models
+
+        # Build master params
+        self.model_params = sum(
+            [[p for p in model.parameters() if p.requires_grad] for model in self.models.values()]
+        , [])
+        if self.fp16_mode == 'amp':
+            self.master_params = self.model_params
+            self.scaler = torch.GradScaler() if self.fp16_mode == 'amp' else None
+        elif self.fp16_mode == 'inflat_all':
+            self.master_params = make_master_params(self.model_params)
+            self.fp16_scale_growth = self.fp16_scale_growth
+            self.log_scale = 20.0
+        elif self.fp16_mode is None:
+            self.master_params = self.model_params
+        else:
+            raise NotImplementedError(f'FP16 mode {self.fp16_mode} is not implemented.')
+
+        # Build EMA params
+        if self.is_master:
+            self.ema_params = [copy.deepcopy(self.master_params) for _ in self.ema_rate]
+
+        # Initialize optimizer
+        if hasattr(torch.optim, self.optimizer_config['name']):
+            self.optimizer = getattr(torch.optim, self.optimizer_config['name'])(self.master_params, **self.optimizer_config['args'])
+        else:
+            self.optimizer = globals()[self.optimizer_config['name']](self.master_params, **self.optimizer_config['args'])
+        
+        # Initalize learning rate scheduler
+        if self.lr_scheduler_config is not None:
+            if hasattr(torch.optim.lr_scheduler, self.lr_scheduler_config['name']):
+                self.lr_scheduler = getattr(torch.optim.lr_scheduler, self.lr_scheduler_config['name'])(self.optimizer, **self.lr_scheduler_config['args'])
+            else:
+                self.lr_scheduler = globals()[self.lr_scheduler_config['name']](self.optimizer, **self.lr_scheduler_config['args'])
+
+        # Initialize elastic memory controller
+        if self.elastic_controller_config is not None:
+            assert any([isinstance(model, (elastic_utils.ElasticModule, elastic_utils.ElasticModuleMixin)) for model in self.models.values()]), \
+                'No elastic module found in models, please inherit from ElasticModule or ElasticModuleMixin'
+            self.elastic_controller = getattr(elastic_utils, self.elastic_controller_config['name'])(**self.elastic_controller_config['args'])
+            for model in self.models.values():
+                if isinstance(model, (elastic_utils.ElasticModule, elastic_utils.ElasticModuleMixin)):
+                    model.register_memory_controller(self.elastic_controller)
+
+        # Initialize gradient clipper
+        if self.grad_clip is not None:
+            if isinstance(self.grad_clip, (float, int)):
+                self.grad_clip = float(self.grad_clip)
+            else:
+                self.grad_clip = getattr(grad_clip_utils, self.grad_clip['name'])(**self.grad_clip['args'])
+
+    def _master_params_to_state_dicts(self, master_params):
+        """
+        Convert master params to dict of state_dicts.
+        """
+        if self.fp16_mode == 'inflat_all':
+            master_params = unflatten_master_params(self.model_params, master_params)
+        state_dicts = {name: model.state_dict() for name, model in self.models.items()}
+        master_params_names = sum(
+            [[(name, n) for n, p in model.named_parameters() if p.requires_grad] for name, model in self.models.items()]
+        , [])
+        for i, (model_name, param_name) in enumerate(master_params_names):
+            state_dicts[model_name][param_name] = master_params[i]
+        return state_dicts
+
+    def _state_dicts_to_master_params(self, master_params, state_dicts):
+        """
+        Convert a state_dict to master params.
+        """
+        master_params_names = sum(
+            [[(name, n) for n, p in model.named_parameters() if p.requires_grad] for name, model in self.models.items()]
+        , [])
+        params = [state_dicts[name][param_name] for name, param_name in master_params_names]
+        if self.fp16_mode == 'inflat_all':
+            model_params_to_master_params(params, master_params)
+        else:
+            for i, param in enumerate(params):
+                master_params[i].data.copy_(param.data)
+
+    def load(self, load_dir, step=0):
+        """
+        Load a checkpoint.
+        Should be called by all processes.
+        """
+        if self.is_master:
+            print(f'\nLoading checkpoint from step {step}...', end='')
+            
+        model_ckpts = {}
+        for name, model in self.models.items():
+            model_ckpt = torch.load(read_file_dist(os.path.join(load_dir, 'ckpts', f'{name}_step{step:07d}.pt')), map_location=self.device, weights_only=True)
+            model_ckpts[name] = model_ckpt
+            model.load_state_dict(model_ckpt)
+            if self.fp16_mode == 'inflat_all':
+                model.convert_to_fp16()
+        self._state_dicts_to_master_params(self.master_params, model_ckpts)
+        del model_ckpts
+
+        if self.is_master:
+            for i, ema_rate in enumerate(self.ema_rate):
+                ema_ckpts = {}
+                for name, model in self.models.items():
+                    ema_ckpt = torch.load(os.path.join(load_dir, 'ckpts', f'{name}_ema{ema_rate}_step{step:07d}.pt'), map_location=self.device, weights_only=True)
+                    ema_ckpts[name] = ema_ckpt
+                self._state_dicts_to_master_params(self.ema_params[i], ema_ckpts)
+                del ema_ckpts
+        
+        misc_ckpt = torch.load(read_file_dist(os.path.join(load_dir, 'ckpts', f'misc_step{step:07d}.pt')), map_location=torch.device('cpu'), weights_only=False)
+        self.optimizer.load_state_dict(misc_ckpt['optimizer'])
+        self.step = misc_ckpt['step']
+        self.data_sampler.load_state_dict(misc_ckpt['data_sampler'])
+        if self.fp16_mode == 'amp':
+            self.scaler.load_state_dict(misc_ckpt['scaler'])
+        elif self.fp16_mode == 'inflat_all':
+            self.log_scale = misc_ckpt['log_scale']
+        if self.lr_scheduler_config is not None:
+            self.lr_scheduler.load_state_dict(misc_ckpt['lr_scheduler'])
+        if self.elastic_controller_config is not None:
+            self.elastic_controller.load_state_dict(misc_ckpt['elastic_controller'])
+        if self.grad_clip is not None and not isinstance(self.grad_clip, float):
+            self.grad_clip.load_state_dict(misc_ckpt['grad_clip'])
+        del misc_ckpt
+
+        if self.world_size > 1:
+            dist.barrier()
+        if self.is_master:
+            print(' Done.')
+
+        if self.world_size > 1:
+            self.check_ddp()
+
+    def save(self):
+        """
+        Save a checkpoint.
+        Should be called only by the rank 0 process.
+        """
+        assert self.is_master, 'save() should be called only by the rank 0 process.'
+        print(f'\nSaving checkpoint at step {self.step}...', end='')
+        
+        model_ckpts = self._master_params_to_state_dicts(self.master_params)
+        for name, model_ckpt in model_ckpts.items():
+            torch.save(model_ckpt, os.path.join(self.output_dir, 'ckpts', f'{name}_step{self.step:07d}.pt'))
+        
+        for i, ema_rate in enumerate(self.ema_rate):
+            ema_ckpts = self._master_params_to_state_dicts(self.ema_params[i])
+            for name, ema_ckpt in ema_ckpts.items():
+                torch.save(ema_ckpt, os.path.join(self.output_dir, 'ckpts', f'{name}_ema{ema_rate}_step{self.step:07d}.pt'))
+
+        misc_ckpt = {
+            'optimizer': self.optimizer.state_dict(),
+            'step': self.step,
+            'data_sampler': self.data_sampler.state_dict(),
+        }
+        if self.fp16_mode == 'amp':
+            misc_ckpt['scaler'] = self.scaler.state_dict()
+        elif self.fp16_mode == 'inflat_all':
+            misc_ckpt['log_scale'] = self.log_scale
+        if self.lr_scheduler_config is not None:
+            misc_ckpt['lr_scheduler'] = self.lr_scheduler.state_dict()
+        if self.elastic_controller_config is not None:
+            misc_ckpt['elastic_controller'] = self.elastic_controller.state_dict()
+        if self.grad_clip is not None and not isinstance(self.grad_clip, float):
+            misc_ckpt['grad_clip'] = self.grad_clip.state_dict()
+        torch.save(misc_ckpt, os.path.join(self.output_dir, 'ckpts', f'misc_step{self.step:07d}.pt'))
+        print(' Done.')
+
+    def finetune_from(self, finetune_ckpt):
+        """
+        Finetune from a checkpoint.
+        Should be called by all processes.
+        """
+        if self.is_master:
+            print('\nFinetuning from:')
+            for name, path in finetune_ckpt.items():
+                print(f'  - {name}: {path}')
+        
+        model_ckpts = {}
+        for name, model in self.models.items():
+            model_state_dict = model.state_dict()
+            if name in finetune_ckpt:
+                model_ckpt = torch.load(read_file_dist(finetune_ckpt[name]), map_location=self.device, weights_only=True)
+                for k, v in model_ckpt.items():
+                    if model_ckpt[k].shape != model_state_dict[k].shape:
+                        if self.is_master:
+                            print(f'Warning: {k} shape mismatch, {model_ckpt[k].shape} vs {model_state_dict[k].shape}, skipped.')
+                        model_ckpt[k] = model_state_dict[k]
+                model_ckpts[name] = model_ckpt
+                model.load_state_dict(model_ckpt)
+                if self.fp16_mode == 'inflat_all':
+                    model.convert_to_fp16()
+            else:
+                if self.is_master:
+                    print(f'Warning: {name} not found in finetune_ckpt, skipped.')
+                model_ckpts[name] = model_state_dict
+        self._state_dicts_to_master_params(self.master_params, model_ckpts)
+        del model_ckpts
+
+        if self.world_size > 1:
+            dist.barrier()
+        if self.is_master:
+            print('Done.')
+
+        if self.world_size > 1:
+            self.check_ddp()
+
+    def update_ema(self):
+        """
+        Update exponential moving average.
+        Should only be called by the rank 0 process.
+        """
+        assert self.is_master, 'update_ema() should be called only by the rank 0 process.'
+        for i, ema_rate in enumerate(self.ema_rate):
+            for master_param, ema_param in zip(self.master_params, self.ema_params[i]):
+                ema_param.detach().mul_(ema_rate).add_(master_param, alpha=1.0 - ema_rate)
+
+    def check_ddp(self):
+        """
+        Check if DDP is working properly.
+        Should be called by all process.
+        """
+        if self.is_master:
+            print('\nPerforming DDP check...')
+
+        if self.is_master:
+            print('Checking if parameters are consistent across processes...')
+        dist.barrier()
+        try:
+            for p in self.master_params:
+                # split to avoid OOM
+                for i in range(0, p.numel(), 10000000):
+                    sub_size = min(10000000, p.numel() - i)
+                    sub_p = p.detach().view(-1)[i:i+sub_size]
+                    # gather from all processes
+                    sub_p_gather = [torch.empty_like(sub_p) for _ in range(self.world_size)]
+                    dist.all_gather(sub_p_gather, sub_p)
+                    # check if equal
+                    assert all([torch.equal(sub_p, sub_p_gather[i]) for i in range(self.world_size)]), 'parameters are not consistent across processes'
+        except AssertionError as e:
+            if self.is_master:
+                print(f'\n\033[91mError: {e}\033[0m')
+                print('DDP check failed.')
+            raise e
+
+        dist.barrier()
+        if self.is_master:
+            print('Done.')
+
+    def run_step(self, data_list):
+        """
+        Run a training step.
+        """
+        step_log = {'loss': {}, 'status': {}}
+        amp_context = partial(torch.autocast, device_type='cuda') if self.fp16_mode == 'amp' else nullcontext
+        elastic_controller_context = self.elastic_controller.record if self.elastic_controller_config is not None else nullcontext
+
+        # Train
+        losses = []
+        statuses = []
+        elastic_controller_logs = []
+        zero_grad(self.model_params)
+        for i, mb_data in enumerate(data_list):
+            ## sync at the end of each batch split
+            sync_contexts = [self.training_models[name].no_sync for name in self.training_models] if i != len(data_list) - 1 and self.world_size > 1 else [nullcontext]
+            with nested_contexts(*sync_contexts), elastic_controller_context():
+                with amp_context():
+                    loss, status = self.training_losses(**mb_data)
+                    l = loss['loss'] / len(data_list)
+                ## backward
+                if self.fp16_mode == 'amp':
+                    self.scaler.scale(l).backward()
+                elif self.fp16_mode == 'inflat_all':
+                    scaled_l = l * (2 ** self.log_scale)
+                    scaled_l.backward()
+                else:
+                    l.backward()
+            ## log
+            losses.append(dict_foreach(loss, lambda x: x.item() if isinstance(x, torch.Tensor) else x))
+            statuses.append(dict_foreach(status, lambda x: x.item() if isinstance(x, torch.Tensor) else x))
+            if self.elastic_controller_config is not None:
+                elastic_controller_logs.append(self.elastic_controller.log())
+        ## gradient clip
+        if self.grad_clip is not None:
+            if self.fp16_mode == 'amp':
+                self.scaler.unscale_(self.optimizer)
+            elif self.fp16_mode == 'inflat_all':
+                model_grads_to_master_grads(self.model_params, self.master_params)
+                self.master_params[0].grad.mul_(1.0 / (2 ** self.log_scale))
+            if isinstance(self.grad_clip, float):
+                grad_norm = torch.nn.utils.clip_grad_norm_(self.master_params, self.grad_clip)
+            else:
+                grad_norm = self.grad_clip(self.master_params)
+            if torch.isfinite(grad_norm):
+                statuses[-1]['grad_norm'] = grad_norm.item()
+        ## step
+        if self.fp16_mode == 'amp':
+            prev_scale = self.scaler.get_scale()
+            self.scaler.step(self.optimizer)
+            self.scaler.update()
+        elif self.fp16_mode == 'inflat_all':
+            prev_scale = 2 ** self.log_scale
+            if not any(not p.grad.isfinite().all() for p in self.model_params):
+                if self.grad_clip is None:
+                    model_grads_to_master_grads(self.model_params, self.master_params)
+                    self.master_params[0].grad.mul_(1.0 / (2 ** self.log_scale))
+                self.optimizer.step()
+                master_params_to_model_params(self.model_params, self.master_params)
+                self.log_scale += self.fp16_scale_growth
+            else:
+                self.log_scale -= 1
+        else:
+            prev_scale = 1.0
+            if not any(not p.grad.isfinite().all() for p in self.model_params):
+                self.optimizer.step()
+            else:
+                print('\n\033[93mWarning: NaN detected in gradients. Skipping update.\033[0m') 
+        ## adjust learning rate
+        if self.lr_scheduler_config is not None:
+            statuses[-1]['lr'] = self.lr_scheduler.get_last_lr()[0]
+            self.lr_scheduler.step()
+
+        # Logs
+        step_log['loss'] = dict_reduce(losses, lambda x: np.mean(x))
+        step_log['status'] = dict_reduce(statuses, lambda x: np.mean(x), special_func={'min': lambda x: np.min(x), 'max': lambda x: np.max(x)})
+        if self.elastic_controller_config is not None:
+            step_log['elastic'] = dict_reduce(elastic_controller_logs, lambda x: np.mean(x))
+        if self.grad_clip is not None:
+            step_log['grad_clip'] = self.grad_clip if isinstance(self.grad_clip, float) else self.grad_clip.log()
+            
+        # Check grad and norm of each param
+        if self.log_param_stats:
+            param_norms = {}
+            param_grads = {}
+            for name, param in self.backbone.named_parameters():
+                if param.requires_grad:
+                    param_norms[name] = param.norm().item()
+                    if param.grad is not None and torch.isfinite(param.grad).all():
+                        param_grads[name] = param.grad.norm().item() / prev_scale
+            step_log['param_norms'] = param_norms
+            step_log['param_grads'] = param_grads
+
+        # Update exponential moving average
+        if self.is_master:
+            self.update_ema()
+
+        return step_log

trellis/trainers/flow_matching/flow_matching.py

@@ -0,0 +1,353 @@
+from typing import *
+import copy
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+import numpy as np
+from easydict import EasyDict as edict
+
+from ..basic import BasicTrainer
+from ...pipelines import samplers 
+from ...utils.general_utils import dict_reduce
+from .mixins.classifier_free_guidance import ClassifierFreeGuidanceMixin
+from .mixins.text_conditioned import TextConditionedMixin
+from .mixins.image_conditioned import ImageConditionedMixin
+
+
+class FlowMatchingTrainer(BasicTrainer):
+    """
+    Trainer for diffusion model with flow matching objective.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+    """
+    def __init__(
+        self,
+        *args,
+        t_schedule: dict = {
+            'name': 'logitNormal',
+            'args': {
+                'mean': 0.0,
+                'std': 1.0,
+            }
+        },
+        sigma_min: float = 1e-5,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.t_schedule = t_schedule
+        self.sigma_min = sigma_min
+
+    def diffuse(self, x_0: torch.Tensor, t: torch.Tensor, noise: Optional[torch.Tensor] = None) -> torch.Tensor:
+        """
+        Diffuse the data for a given number of diffusion steps.
+        In other words, sample from q(x_t | x_0).
+
+        Args:
+            x_0: The [N x C x ...] tensor of noiseless inputs.
+            t: The [N] tensor of diffusion steps [0-1].
+            noise: If specified, use this noise instead of generating new noise.
+
+        Returns:
+            x_t, the noisy version of x_0 under timestep t.
+        """
+        if noise is None:
+            noise = torch.randn_like(x_0)
+        assert noise.shape == x_0.shape, "noise must have same shape as x_0"
+
+        t = t.view(-1, *[1 for _ in range(len(x_0.shape) - 1)])
+        x_t = (1 - t) * x_0 + (self.sigma_min + (1 - self.sigma_min) * t) * noise
+
+        return x_t
+
+    def reverse_diffuse(self, x_t: torch.Tensor, t: torch.Tensor, noise: torch.Tensor) -> torch.Tensor:
+        """
+        Get original image from noisy version under timestep t.
+        """
+        assert noise.shape == x_t.shape, "noise must have same shape as x_t"
+        t = t.view(-1, *[1 for _ in range(len(x_t.shape) - 1)])
+        x_0 = (x_t - (self.sigma_min + (1 - self.sigma_min) * t) * noise) / (1 - t)
+        return x_0
+
+    def get_v(self, x_0: torch.Tensor, noise: torch.Tensor, t: torch.Tensor) -> torch.Tensor:
+        """
+        Compute the velocity of the diffusion process at time t.
+        """
+        return (1 - self.sigma_min) * noise - x_0
+
+    def get_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data.
+        """
+        return cond
+    
+    def get_inference_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data for inference.
+        """
+        return {'cond': cond, **kwargs}
+
+    def get_sampler(self, **kwargs) -> samplers.FlowEulerSampler:
+        """
+        Get the sampler for the diffusion process.
+        """
+        return samplers.FlowEulerSampler(self.sigma_min)
+    
+    def vis_cond(self, **kwargs):
+        """
+        Visualize the conditioning data.
+        """
+        return {}
+
+    def sample_t(self, batch_size: int) -> torch.Tensor:
+        """
+        Sample timesteps.
+        """
+        if self.t_schedule['name'] == 'uniform':
+            t = torch.rand(batch_size)
+        elif self.t_schedule['name'] == 'logitNormal':
+            mean = self.t_schedule['args']['mean']
+            std = self.t_schedule['args']['std']
+            t = torch.sigmoid(torch.randn(batch_size) * std + mean)
+        else:
+            raise ValueError(f"Unknown t_schedule: {self.t_schedule['name']}")
+        return t
+
+    def training_losses(
+        self,
+        x_0: torch.Tensor,
+        cond=None,
+        **kwargs
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses for a single timestep.
+
+        Args:
+            x_0: The [N x C x ...] tensor of noiseless inputs.
+            cond: The [N x ...] tensor of additional conditions.
+            kwargs: Additional arguments to pass to the backbone.
+
+        Returns:
+            a dict with the key "loss" containing a tensor of shape [N].
+            may also contain other keys for different terms.
+        """
+        noise = torch.randn_like(x_0)
+        t = self.sample_t(x_0.shape[0]).to(x_0.device).float()
+        x_t = self.diffuse(x_0, t, noise=noise)
+        cond = self.get_cond(cond, **kwargs)
+        
+        pred = self.training_models['denoiser'](x_t, t * 1000, cond, **kwargs)
+        assert pred.shape == noise.shape == x_0.shape
+        target = self.get_v(x_0, noise, t)
+        terms = edict()
+        terms["mse"] = F.mse_loss(pred, target)
+        terms["loss"] = terms["mse"]
+
+        # log loss with time bins
+        mse_per_instance = np.array([
+            F.mse_loss(pred[i], target[i]).item()
+            for i in range(x_0.shape[0])
+        ])
+        time_bin = np.digitize(t.cpu().numpy(), np.linspace(0, 1, 11)) - 1
+        for i in range(10):
+            if (time_bin == i).sum() != 0:
+                terms[f"bin_{i}"] = {"mse": mse_per_instance[time_bin == i].mean()}
+
+        return terms, {}
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+        )
+
+        # inference
+        sampler = self.get_sampler()
+        sample_gt = []
+        sample = []
+        cond_vis = []
+        for i in range(0, num_samples, batch_size):
+            batch = min(batch_size, num_samples - i)
+            data = next(iter(dataloader))
+            data = {k: v[:batch].cuda() if isinstance(v, torch.Tensor) else v[:batch] for k, v in data.items()}
+            noise = torch.randn_like(data['x_0'])
+            sample_gt.append(data['x_0'])
+            cond_vis.append(self.vis_cond(**data))
+            del data['x_0']
+            args = self.get_inference_cond(**data)
+            res = sampler.sample(
+                self.models['denoiser'],
+                noise=noise,
+                **args,
+                steps=50, cfg_strength=3.0, verbose=verbose,
+            )
+            sample.append(res.samples)
+
+        sample_gt = torch.cat(sample_gt, dim=0)
+        sample = torch.cat(sample, dim=0)
+        sample_dict = {
+            'sample_gt': {'value': sample_gt, 'type': 'sample'},
+            'sample': {'value': sample, 'type': 'sample'},
+        }
+        sample_dict.update(dict_reduce(cond_vis, None, {
+            'value': lambda x: torch.cat(x, dim=0),
+            'type': lambda x: x[0],
+        }))
+        
+        return sample_dict
+
+    
+class FlowMatchingCFGTrainer(ClassifierFreeGuidanceMixin, FlowMatchingTrainer):
+    """
+    Trainer for diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+    """
+    pass
+
+
+class TextConditionedFlowMatchingCFGTrainer(TextConditionedMixin, FlowMatchingCFGTrainer):
+    """
+    Trainer for text-conditioned diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        text_cond_model(str): Text conditioning model.
+    """
+    pass
+
+
+class ImageConditionedFlowMatchingCFGTrainer(ImageConditionedMixin, FlowMatchingCFGTrainer):
+    """
+    Trainer for image-conditioned diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        image_cond_model (str): Image conditioning model.
+    """
+    pass

trellis/trainers/flow_matching/mixins/classifier_free_guidance.py

@@ -0,0 +1,59 @@
+import torch
+import numpy as np
+from ....utils.general_utils import dict_foreach
+from ....pipelines import samplers
+
+
+class ClassifierFreeGuidanceMixin:
+    def __init__(self, *args, p_uncond: float = 0.1, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.p_uncond = p_uncond
+
+    def get_cond(self, cond, neg_cond=None, **kwargs):
+        """
+        Get the conditioning data.
+        """
+        assert neg_cond is not None, "neg_cond must be provided for classifier-free guidance" 
+
+        if self.p_uncond > 0:
+            # randomly drop the class label
+            def get_batch_size(cond):
+                if isinstance(cond, torch.Tensor):
+                    return cond.shape[0]
+                elif isinstance(cond, list):
+                    return len(cond)
+                else:
+                    raise ValueError(f"Unsupported type of cond: {type(cond)}")
+                
+            ref_cond = cond if not isinstance(cond, dict) else cond[list(cond.keys())[0]]
+            B = get_batch_size(ref_cond)
+            
+            def select(cond, neg_cond, mask):
+                if isinstance(cond, torch.Tensor):
+                    mask = torch.tensor(mask, device=cond.device).reshape(-1, *[1] * (cond.ndim - 1))
+                    return torch.where(mask, neg_cond, cond)
+                elif isinstance(cond, list):
+                    return [nc if m else c for c, nc, m in zip(cond, neg_cond, mask)]
+                else:
+                    raise ValueError(f"Unsupported type of cond: {type(cond)}")
+            
+            mask = list(np.random.rand(B) < self.p_uncond)
+            if not isinstance(cond, dict):
+                cond = select(cond, neg_cond, mask)
+            else:
+                cond = dict_foreach([cond, neg_cond], lambda x: select(x[0], x[1], mask))
+    
+        return cond
+
+    def get_inference_cond(self, cond, neg_cond=None, **kwargs):
+        """
+        Get the conditioning data for inference.
+        """
+        assert neg_cond is not None, "neg_cond must be provided for classifier-free guidance"
+        return {'cond': cond, 'neg_cond': neg_cond, **kwargs}
+    
+    def get_sampler(self, **kwargs) -> samplers.FlowEulerCfgSampler:
+        """
+        Get the sampler for the diffusion process.
+        """
+        return samplers.FlowEulerCfgSampler(self.sigma_min)

trellis/trainers/flow_matching/mixins/image_conditioned.py

@@ -0,0 +1,93 @@
+from typing import *
+import torch
+import torch.nn.functional as F
+from torchvision import transforms
+import numpy as np
+from PIL import Image
+
+from ....utils import dist_utils
+
+
+class ImageConditionedMixin:
+    """
+    Mixin for image-conditioned models.
+    
+    Args:
+        image_cond_model: The image conditioning model.
+    """
+    def __init__(self, *args, image_cond_model: str = 'dinov2_vitl14_reg', **kwargs):
+        super().__init__(*args, **kwargs)
+        self.image_cond_model_name = image_cond_model
+        self.image_cond_model = None     # the model is init lazily
+        
+    @staticmethod
+    def prepare_for_training(image_cond_model: str, **kwargs):
+        """
+        Prepare for training.
+        """
+        if hasattr(super(ImageConditionedMixin, ImageConditionedMixin), 'prepare_for_training'):
+            super(ImageConditionedMixin, ImageConditionedMixin).prepare_for_training(**kwargs)
+        # download the model
+        torch.hub.load('facebookresearch/dinov2', image_cond_model, pretrained=True)
+        
+    def _init_image_cond_model(self):
+        """
+        Initialize the image conditioning model.
+        """
+        with dist_utils.local_master_first():
+            dinov2_model = torch.hub.load('facebookresearch/dinov2', self.image_cond_model_name, pretrained=True)
+        dinov2_model.eval().cuda()
+        transform = transforms.Compose([
+            transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
+        ])
+        self.image_cond_model = {
+            'model': dinov2_model,
+            'transform': transform,
+        }
+    
+    @torch.no_grad()
+    def encode_image(self, image: Union[torch.Tensor, List[Image.Image]]) -> torch.Tensor:
+        """
+        Encode the image.
+        """
+        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)}")
+        
+        if self.image_cond_model is None:
+            self._init_image_cond_model()
+        image = self.image_cond_model['transform'](image).cuda()
+        features = self.image_cond_model['model'](image, is_training=True)['x_prenorm']
+        patchtokens = F.layer_norm(features, features.shape[-1:])
+        return patchtokens
+        
+    def get_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data.
+        """
+        cond = self.encode_image(cond)
+        kwargs['neg_cond'] = torch.zeros_like(cond)
+        cond = super().get_cond(cond, **kwargs)
+        return cond
+    
+    def get_inference_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data for inference.
+        """
+        cond = self.encode_image(cond)
+        kwargs['neg_cond'] = torch.zeros_like(cond)
+        cond = super().get_inference_cond(cond, **kwargs)
+        return cond
+
+    def vis_cond(self, cond, **kwargs):
+        """
+        Visualize the conditioning data.
+        """
+        return {'image': {'value': cond, 'type': 'image'}}

trellis/trainers/flow_matching/mixins/text_conditioned.py

@@ -0,0 +1,68 @@
+from typing import *
+import os
+os.environ['TOKENIZERS_PARALLELISM'] = 'true'
+import torch
+from transformers import AutoTokenizer, CLIPTextModel
+
+from ....utils import dist_utils
+
+
+class TextConditionedMixin:
+    """
+    Mixin for text-conditioned models.
+    
+    Args:
+        text_cond_model: The text conditioning model.
+    """
+    def __init__(self, *args, text_cond_model: str = 'openai/clip-vit-large-patch14', **kwargs):
+        super().__init__(*args, **kwargs)
+        self.text_cond_model_name = text_cond_model
+        self.text_cond_model = None     # the model is init lazily
+        
+    def _init_text_cond_model(self):
+        """
+        Initialize the text conditioning model.
+        """
+        # load model
+        with dist_utils.local_master_first():
+            model = CLIPTextModel.from_pretrained(self.text_cond_model_name)
+            tokenizer = AutoTokenizer.from_pretrained(self.text_cond_model_name)
+        model.eval()
+        model = model.cuda()
+        self.text_cond_model = {
+            'model': model,
+            'tokenizer': tokenizer,
+        }
+        self.text_cond_model['null_cond'] = self.encode_text([''])
+        
+    @torch.no_grad()
+    def encode_text(self, text: List[str]) -> torch.Tensor:
+        """
+        Encode the text.
+        """
+        assert isinstance(text, list) and isinstance(text[0], str), "TextConditionedMixin only supports list of strings as cond"
+        if self.text_cond_model is None:
+            self._init_text_cond_model()
+        encoding = self.text_cond_model['tokenizer'](text, max_length=77, padding='max_length', truncation=True, return_tensors='pt')
+        tokens = encoding['input_ids'].cuda()
+        embeddings = self.text_cond_model['model'](input_ids=tokens).last_hidden_state
+        
+        return embeddings
+        
+    def get_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data.
+        """
+        cond = self.encode_text(cond)
+        kwargs['neg_cond'] = self.text_cond_model['null_cond'].repeat(cond.shape[0], 1, 1)
+        cond = super().get_cond(cond, **kwargs)
+        return cond
+    
+    def get_inference_cond(self, cond, **kwargs):
+        """
+        Get the conditioning data for inference.
+        """
+        cond = self.encode_text(cond)
+        kwargs['neg_cond'] = self.text_cond_model['null_cond'].repeat(cond.shape[0], 1, 1)
+        cond = super().get_inference_cond(cond, **kwargs)
+        return cond

trellis/trainers/flow_matching/sparse_flow_matching.py

@@ -0,0 +1,286 @@
+from typing import *
+import os
+import copy
+import functools
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+import numpy as np
+from easydict import EasyDict as edict
+
+from ...modules import sparse as sp
+from ...utils.general_utils import dict_reduce
+from ...utils.data_utils import cycle, BalancedResumableSampler
+from .flow_matching import FlowMatchingTrainer
+from .mixins.classifier_free_guidance import ClassifierFreeGuidanceMixin
+from .mixins.text_conditioned import TextConditionedMixin
+from .mixins.image_conditioned import ImageConditionedMixin
+
+
+class SparseFlowMatchingTrainer(FlowMatchingTrainer):
+    """
+    Trainer for sparse diffusion model with flow matching objective.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+    """
+    
+    def prepare_dataloader(self, **kwargs):
+        """
+        Prepare dataloader.
+        """
+        self.data_sampler = BalancedResumableSampler(
+            self.dataset,
+            shuffle=True,
+            batch_size=self.batch_size_per_gpu,
+        )
+        self.dataloader = DataLoader(
+            self.dataset,
+            batch_size=self.batch_size_per_gpu,
+            num_workers=int(np.ceil(os.cpu_count() / torch.cuda.device_count())),
+            pin_memory=True,
+            drop_last=True,
+            persistent_workers=True,
+            collate_fn=functools.partial(self.dataset.collate_fn, split_size=self.batch_split),
+            sampler=self.data_sampler,
+        )
+        self.data_iterator = cycle(self.dataloader)
+        
+    def training_losses(
+        self,
+        x_0: sp.SparseTensor,
+        cond=None,
+        **kwargs
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses for a single timestep.
+
+        Args:
+            x_0: The [N x ... x C] sparse tensor of the inputs.
+            cond: The [N x ...] tensor of additional conditions.
+            kwargs: Additional arguments to pass to the backbone.
+
+        Returns:
+            a dict with the key "loss" containing a tensor of shape [N].
+            may also contain other keys for different terms.
+        """
+        noise = x_0.replace(torch.randn_like(x_0.feats))
+        t = self.sample_t(x_0.shape[0]).to(x_0.device).float()
+        x_t = self.diffuse(x_0, t, noise=noise)
+        cond = self.get_cond(cond, **kwargs)
+        
+        pred = self.training_models['denoiser'](x_t, t * 1000, cond, **kwargs)
+        assert pred.shape == noise.shape == x_0.shape
+        target = self.get_v(x_0, noise, t)
+        terms = edict()
+        terms["mse"] = F.mse_loss(pred.feats, target.feats)
+        terms["loss"] = terms["mse"]
+
+        # log loss with time bins
+        mse_per_instance = np.array([
+            F.mse_loss(pred.feats[x_0.layout[i]], target.feats[x_0.layout[i]]).item()
+            for i in range(x_0.shape[0])
+        ])
+        time_bin = np.digitize(t.cpu().numpy(), np.linspace(0, 1, 11)) - 1
+        for i in range(10):
+            if (time_bin == i).sum() != 0:
+                terms[f"bin_{i}"] = {"mse": mse_per_instance[time_bin == i].mean()}
+
+        return terms, {}
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+        )
+
+        # inference
+        sampler = self.get_sampler()
+        sample_gt = []
+        sample = []
+        cond_vis = []
+        for i in range(0, num_samples, batch_size):
+            batch = min(batch_size, num_samples - i)
+            data = next(iter(dataloader))
+            data = {k: v[:batch].cuda() if not isinstance(v, list) else v[:batch] for k, v in data.items()}
+            noise = data['x_0'].replace(torch.randn_like(data['x_0'].feats))
+            sample_gt.append(data['x_0'])
+            cond_vis.append(self.vis_cond(**data))
+            del data['x_0']
+            args = self.get_inference_cond(**data)
+            res = sampler.sample(
+                self.models['denoiser'],
+                noise=noise,
+                **args,
+                steps=50, cfg_strength=3.0, verbose=verbose,
+            )
+            sample.append(res.samples)
+
+        sample_gt = sp.sparse_cat(sample_gt)
+        sample = sp.sparse_cat(sample)
+        sample_dict = {
+            'sample_gt': {'value': sample_gt, 'type': 'sample'},
+            'sample': {'value': sample, 'type': 'sample'},
+        }
+        sample_dict.update(dict_reduce(cond_vis, None, {
+            'value': lambda x: torch.cat(x, dim=0),
+            'type': lambda x: x[0],
+        }))
+        
+        return sample_dict
+
+
+class SparseFlowMatchingCFGTrainer(ClassifierFreeGuidanceMixin, SparseFlowMatchingTrainer):
+    """
+    Trainer for sparse diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+    """
+    pass
+
+
+class TextConditionedSparseFlowMatchingCFGTrainer(TextConditionedMixin, SparseFlowMatchingCFGTrainer):
+    """
+    Trainer for sparse text-conditioned diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        text_cond_model(str): Text conditioning model.
+    """
+    pass
+
+
+class ImageConditionedSparseFlowMatchingCFGTrainer(ImageConditionedMixin, SparseFlowMatchingCFGTrainer):
+    """
+    Trainer for sparse image-conditioned diffusion model with flow matching objective and classifier-free guidance.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+
+        t_schedule (dict): Time schedule for flow matching.
+        sigma_min (float): Minimum noise level.
+        p_uncond (float): Probability of dropping conditions.
+        image_cond_model (str): Image conditioning model.
+    """
+    pass

trellis/trainers/utils.py

@@ -0,0 +1,77 @@
+import torch.nn as nn
+
+
+# FP16 utils
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+def make_master_params(model_params):
+    """
+    Copy model parameters into a inflated tensor of full-precision parameters.
+    """
+    master_params = _flatten_dense_tensors(
+        [param.detach().float() for param in model_params]
+    )
+    master_params = nn.Parameter(master_params)
+    master_params.requires_grad = True
+    return [master_params]
+
+
+def unflatten_master_params(model_params, master_params):
+    """
+    Unflatten the master parameters to look like model_params.
+    """
+    return _unflatten_dense_tensors(master_params[0].detach(), model_params)
+
+
+def model_params_to_master_params(model_params, master_params):
+    """
+    Copy the model parameter data into the master parameters.
+    """
+    master_params[0].detach().copy_(
+        _flatten_dense_tensors([param.detach().float() for param in model_params])
+    )
+
+
+def master_params_to_model_params(model_params, master_params):
+    """
+    Copy the master parameter data back into the model parameters.
+    """
+    for param, master_param in zip(
+        model_params, _unflatten_dense_tensors(master_params[0].detach(), model_params)
+    ):
+        param.detach().copy_(master_param)
+
+
+def model_grads_to_master_grads(model_params, master_params):
+    """
+    Copy the gradients from the model parameters into the master parameters
+    from make_master_params().
+    """
+    master_params[0].grad = _flatten_dense_tensors(
+        [param.grad.data.detach().float() for param in model_params]
+    )
+    
+
+def zero_grad(model_params):
+    for param in model_params:
+       if param.grad is not None:
+            if param.grad.grad_fn is not None:
+                param.grad.detach_()
+            else:
+                param.grad.requires_grad_(False)
+            param.grad.zero_()
+            
+
+# LR Schedulers
+from torch.optim.lr_scheduler import LambdaLR
+
+class LinearWarmupLRScheduler(LambdaLR):
+    def __init__(self, optimizer, warmup_steps, last_epoch=-1):
+        self.warmup_steps = warmup_steps
+        super(LinearWarmupLRScheduler, self).__init__(optimizer, self.lr_lambda, last_epoch=last_epoch)
+        
+    def lr_lambda(self, current_step):
+        if current_step < self.warmup_steps:
+            return float(current_step + 1) / self.warmup_steps
+        return 1.0
+        

trellis/trainers/vae/sparse_structure_vae.py

@@ -0,0 +1,130 @@
+from typing import *
+import copy
+import torch
+import torch.nn.functional as F
+from torch.utils.data import DataLoader
+from easydict import EasyDict as edict
+
+from ..basic import BasicTrainer
+
+
+class SparseStructureVaeTrainer(BasicTrainer):
+    """
+    Trainer for Sparse Structure VAE.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+        
+        loss_type (str): Loss type. 'bce' for binary cross entropy, 'l1' for L1 loss, 'dice' for Dice loss.
+        lambda_kl (float): KL divergence loss weight.
+    """
+    
+    def __init__(
+        self,
+        *args,
+        loss_type='bce',
+        lambda_kl=1e-6,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.loss_type = loss_type
+        self.lambda_kl = lambda_kl
+    
+    def training_losses(
+        self,
+        ss: torch.Tensor,
+        **kwargs
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses.
+
+        Args:
+            ss: The [N x 1 x H x W x D] tensor of binary sparse structure.
+
+        Returns:
+            a dict with the key "loss" containing a scalar tensor.
+            may also contain other keys for different terms.
+        """
+        z, mean, logvar = self.training_models['encoder'](ss.float(), sample_posterior=True, return_raw=True)
+        logits = self.training_models['decoder'](z)
+
+        terms = edict(loss = 0.0)
+        if self.loss_type == 'bce':
+            terms["bce"] = F.binary_cross_entropy_with_logits(logits, ss.float(), reduction='mean')
+            terms["loss"] = terms["loss"] + terms["bce"]
+        elif self.loss_type == 'l1':
+            terms["l1"] = F.l1_loss(F.sigmoid(logits), ss.float(), reduction='mean')
+            terms["loss"] = terms["loss"] + terms["l1"]
+        elif self.loss_type == 'dice':
+            logits = F.sigmoid(logits)
+            terms["dice"] = 1 - (2 * (logits * ss.float()).sum() + 1) / (logits.sum() + ss.float().sum() + 1)
+            terms["loss"] = terms["loss"] + terms["dice"]
+        else:
+            raise ValueError(f'Invalid loss type {self.loss_type}')
+        terms["kl"] = 0.5 * torch.mean(mean.pow(2) + logvar.exp() - logvar - 1)
+        terms["loss"] = terms["loss"] + self.lambda_kl * terms["kl"]
+            
+        return terms, {}
+    
+    @torch.no_grad()
+    def snapshot(self, suffix=None, num_samples=64, batch_size=1, verbose=False):
+        super().snapshot(suffix=suffix, num_samples=num_samples, batch_size=batch_size, verbose=verbose)
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+        )
+
+        # inference
+        gts = []
+        recons = []
+        for i in range(0, num_samples, batch_size):
+            batch = min(batch_size, num_samples - i)
+            data = next(iter(dataloader))
+            args = {k: v[:batch].cuda() if isinstance(v, torch.Tensor) else v[:batch] for k, v in data.items()}
+            z = self.models['encoder'](args['ss'].float(), sample_posterior=False)
+            logits = self.models['decoder'](z)
+            recon = (logits > 0).long()
+            gts.append(args['ss'])
+            recons.append(recon)
+
+        sample_dict = {
+            'gt': {'value': torch.cat(gts, dim=0), 'type': 'sample'},
+            'recon': {'value': torch.cat(recons, dim=0), 'type': 'sample'},
+        }
+        return sample_dict

trellis/trainers/vae/structured_latent_vae_gaussian.py

@@ -0,0 +1,275 @@
+from typing import *
+import copy
+import torch
+from torch.utils.data import DataLoader
+import numpy as np
+from easydict import EasyDict as edict
+import utils3d.torch
+
+from ..basic import BasicTrainer
+from ...representations import Gaussian
+from ...renderers import GaussianRenderer
+from ...modules.sparse import SparseTensor
+from ...utils.loss_utils import l1_loss, l2_loss, ssim, lpips
+
+
+class SLatVaeGaussianTrainer(BasicTrainer):
+    """
+    Trainer for structured latent VAE.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+        
+        loss_type (str): Loss type. Can be 'l1', 'l2'
+        lambda_ssim (float): SSIM loss weight.
+        lambda_lpips (float): LPIPS loss weight.
+        lambda_kl (float): KL loss weight.
+        regularizations (dict): Regularization config.
+    """
+    
+    def __init__(
+        self,
+        *args,
+        loss_type: str = 'l1',
+        lambda_ssim: float = 0.2,
+        lambda_lpips: float = 0.2,
+        lambda_kl: float = 1e-6,
+        regularizations: Dict = {},
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.loss_type = loss_type
+        self.lambda_ssim = lambda_ssim
+        self.lambda_lpips = lambda_lpips
+        self.lambda_kl = lambda_kl
+        self.regularizations = regularizations
+        
+        self._init_renderer()
+        
+    def _init_renderer(self):
+        rendering_options = {"near" : 0.8,
+                             "far" : 1.6,
+                             "bg_color" : 'random'}
+        self.renderer = GaussianRenderer(rendering_options)
+        self.renderer.pipe.kernel_size = self.models['decoder'].rep_config['2d_filter_kernel_size']
+        
+    def _render_batch(self, reps: List[Gaussian], extrinsics: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor:
+        """
+        Render a batch of representations.
+
+        Args:
+            reps: The dictionary of lists of representations.
+            extrinsics: The [N x 4 x 4] tensor of extrinsics.
+            intrinsics: The [N x 3 x 3] tensor of intrinsics.
+        """
+        ret = None
+        for i, representation in enumerate(reps):
+            render_pack = self.renderer.render(representation, extrinsics[i], intrinsics[i])
+            if ret is None:
+                ret = {k: [] for k in list(render_pack.keys()) + ['bg_color']}
+            for k, v in render_pack.items():
+                ret[k].append(v)
+            ret['bg_color'].append(self.renderer.bg_color)
+        for k, v in ret.items():
+            ret[k] = torch.stack(v, dim=0) 
+        return ret
+        
+    @torch.no_grad()
+    def _get_status(self, z: SparseTensor, reps: List[Gaussian]) -> Dict:
+        xyz = torch.cat([g.get_xyz for g in reps], dim=0)
+        xyz_base = (z.coords[:, 1:].float() + 0.5) / self.models['decoder'].resolution - 0.5
+        offset = xyz - xyz_base.unsqueeze(1).expand(-1, self.models['decoder'].rep_config['num_gaussians'], -1).reshape(-1, 3)
+        status = {
+            'xyz': xyz,
+            'offset': offset,
+            'scale': torch.cat([g.get_scaling for g in reps], dim=0),
+            'opacity': torch.cat([g.get_opacity for g in reps], dim=0),
+        }
+
+        for k in list(status.keys()):
+            status[k] = {
+                'mean': status[k].mean().item(),
+                'max': status[k].max().item(),
+                'min': status[k].min().item(),
+            }
+            
+        return status
+    
+    def _get_regularization_loss(self, reps: List[Gaussian]) -> Tuple[torch.Tensor, Dict]:
+        loss = 0.0
+        terms = {}
+        if 'lambda_vol' in self.regularizations:
+            scales = torch.cat([g.get_scaling for g in reps], dim=0)   # [N x 3]
+            volume = torch.prod(scales, dim=1)  # [N]
+            terms[f'reg_vol'] = volume.mean()
+            loss = loss + self.regularizations['lambda_vol'] * terms[f'reg_vol']
+        if 'lambda_opacity' in self.regularizations:
+            opacity = torch.cat([g.get_opacity for g in reps], dim=0)
+            terms[f'reg_opacity'] = (opacity - 1).pow(2).mean()
+            loss = loss + self.regularizations['lambda_opacity'] * terms[f'reg_opacity']
+        return loss, terms
+    
+    def training_losses(
+        self,
+        feats: SparseTensor,
+        image: torch.Tensor,
+        alpha: torch.Tensor,
+        extrinsics: torch.Tensor,
+        intrinsics: torch.Tensor,
+        return_aux: bool = False,
+        **kwargs
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses.
+
+        Args:
+            feats: The [N x * x C] sparse tensor of features.
+            image: The [N x 3 x H x W] tensor of images.
+            alpha: The [N x H x W] tensor of alpha channels.
+            extrinsics: The [N x 4 x 4] tensor of extrinsics.
+            intrinsics: The [N x 3 x 3] tensor of intrinsics.
+            return_aux: Whether to return auxiliary information.
+
+        Returns:
+            a dict with the key "loss" containing a scalar tensor.
+            may also contain other keys for different terms.
+        """
+        z, mean, logvar = self.training_models['encoder'](feats, sample_posterior=True, return_raw=True)
+        reps = self.training_models['decoder'](z)
+        self.renderer.rendering_options.resolution = image.shape[-1]
+        render_results = self._render_batch(reps, extrinsics, intrinsics)     
+        
+        terms = edict(loss = 0.0, rec = 0.0)
+        
+        rec_image = render_results['color']
+        gt_image = image * alpha[:, None] + (1 - alpha[:, None]) * render_results['bg_color'][..., None, None]
+                
+        if self.loss_type == 'l1':
+            terms["l1"] = l1_loss(rec_image, gt_image)
+            terms["rec"] = terms["rec"] + terms["l1"]
+        elif self.loss_type == 'l2':
+            terms["l2"] = l2_loss(rec_image, gt_image)
+            terms["rec"] = terms["rec"] + terms["l2"]
+        else:
+            raise ValueError(f"Invalid loss type: {self.loss_type}")
+        if self.lambda_ssim > 0:
+            terms["ssim"] = 1 - ssim(rec_image, gt_image)
+            terms["rec"] = terms["rec"] + self.lambda_ssim * terms["ssim"]
+        if self.lambda_lpips > 0:
+            terms["lpips"] = lpips(rec_image, gt_image)
+            terms["rec"] = terms["rec"] + self.lambda_lpips * terms["lpips"]
+        terms["loss"] = terms["loss"] + terms["rec"]
+
+        terms["kl"] = 0.5 * torch.mean(mean.pow(2) + logvar.exp() - logvar - 1)
+        terms["loss"] = terms["loss"] + self.lambda_kl * terms["kl"]
+        
+        reg_loss, reg_terms = self._get_regularization_loss(reps)
+        terms.update(reg_terms)
+        terms["loss"] = terms["loss"] + reg_loss
+        
+        status = self._get_status(z, reps)
+        
+        if return_aux:
+            return terms, status, {'rec_image': rec_image, 'gt_image': gt_image}       
+        return terms, status
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+        )
+
+        # inference
+        ret_dict = {}
+        gt_images = []
+        exts = []
+        ints = []
+        reps = []
+        for i in range(0, num_samples, batch_size):
+            batch = min(batch_size, num_samples - i)
+            data = next(iter(dataloader))
+            args = {k: v[:batch].cuda() for k, v in data.items()}
+            gt_images.append(args['image'] * args['alpha'][:, None])
+            exts.append(args['extrinsics'])
+            ints.append(args['intrinsics'])
+            z = self.models['encoder'](args['feats'], sample_posterior=True, return_raw=False)
+            reps.extend(self.models['decoder'](z))
+        gt_images = torch.cat(gt_images, dim=0)
+        ret_dict.update({f'gt_image': {'value': gt_images, 'type': 'image'}})
+
+        # render single view
+        exts = torch.cat(exts, dim=0)
+        ints = torch.cat(ints, dim=0)
+        self.renderer.rendering_options.bg_color = (0, 0, 0)
+        self.renderer.rendering_options.resolution = gt_images.shape[-1]
+        render_results = self._render_batch(reps, exts, ints)
+        ret_dict.update({f'rec_image': {'value': render_results['color'], 'type': 'image'}})
+
+        # render multiview
+        self.renderer.rendering_options.resolution = 512
+        ## 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)]
+
+        ## render each view
+        miltiview_images = []
+        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)
+            extrinsics = extrinsics.unsqueeze(0).expand(num_samples, -1, -1)
+            intrinsics = intrinsics.unsqueeze(0).expand(num_samples, -1, -1)
+            render_results = self._render_batch(reps, extrinsics, intrinsics)
+            miltiview_images.append(render_results['color'])
+
+        ## Concatenate views
+        miltiview_images = torch.cat([
+            torch.cat(miltiview_images[:2], dim=-2),
+            torch.cat(miltiview_images[2:], dim=-2),
+        ], dim=-1)
+        ret_dict.update({f'miltiview_image': {'value': miltiview_images, 'type': 'image'}})
+
+        self.renderer.rendering_options.bg_color = 'random'
+                                    
+        return ret_dict

trellis/trainers/vae/structured_latent_vae_mesh_dec.py

@@ -0,0 +1,382 @@
+from typing import *
+import copy
+import torch
+from torch.utils.data import DataLoader
+import numpy as np
+from easydict import EasyDict as edict
+import utils3d.torch
+
+from ..basic import BasicTrainer
+from ...representations import MeshExtractResult
+from ...renderers import MeshRenderer
+from ...modules.sparse import SparseTensor
+from ...utils.loss_utils import l1_loss, smooth_l1_loss, ssim, lpips
+from ...utils.data_utils import recursive_to_device
+
+
+class SLatVaeMeshDecoderTrainer(BasicTrainer):
+    """
+    Trainer for structured latent VAE Mesh Decoder.
+    
+    Args:
+        models (dict[str, nn.Module]): Models to train.
+        dataset (torch.utils.data.Dataset): Dataset.
+        output_dir (str): Output directory.
+        load_dir (str): Load directory.
+        step (int): Step to load.
+        batch_size (int): Batch size.
+        batch_size_per_gpu (int): Batch size per GPU. If specified, batch_size will be ignored.
+        batch_split (int): Split batch with gradient accumulation.
+        max_steps (int): Max steps.
+        optimizer (dict): Optimizer config.
+        lr_scheduler (dict): Learning rate scheduler config.
+        elastic (dict): Elastic memory management config.
+        grad_clip (float or dict): Gradient clip config.
+        ema_rate (float or list): Exponential moving average rates.
+        fp16_mode (str): FP16 mode.
+            - None: No FP16.
+            - 'inflat_all': Hold a inflated fp32 master param for all params.
+            - 'amp': Automatic mixed precision.
+        fp16_scale_growth (float): Scale growth for FP16 gradient backpropagation.
+        finetune_ckpt (dict): Finetune checkpoint.
+        log_param_stats (bool): Log parameter stats.
+        i_print (int): Print interval.
+        i_log (int): Log interval.
+        i_sample (int): Sample interval.
+        i_save (int): Save interval.
+        i_ddpcheck (int): DDP check interval.
+        
+        loss_type (str): Loss type. Can be 'l1', 'l2'
+        lambda_ssim (float): SSIM loss weight.
+        lambda_lpips (float): LPIPS loss weight.
+    """
+    
+    def __init__(
+        self,
+        *args,
+        depth_loss_type: str = 'l1',
+        lambda_depth: int = 1,
+        lambda_ssim: float = 0.2,
+        lambda_lpips: float = 0.2,
+        lambda_tsdf: float = 0.01,
+        lambda_color: float = 0.1,
+        **kwargs
+    ):
+        super().__init__(*args, **kwargs)
+        self.depth_loss_type = depth_loss_type
+        self.lambda_depth = lambda_depth
+        self.lambda_ssim = lambda_ssim
+        self.lambda_lpips = lambda_lpips
+        self.lambda_tsdf = lambda_tsdf
+        self.lambda_color = lambda_color
+        self.use_color = self.lambda_color > 0
+        
+        self._init_renderer()
+        
+    def _init_renderer(self):
+        rendering_options = {"near" : 1,
+                             "far" : 3}
+        self.renderer = MeshRenderer(rendering_options, device=self.device)
+        
+    def _render_batch(self, reps: List[MeshExtractResult], extrinsics: torch.Tensor, intrinsics: torch.Tensor,
+                      return_types=['mask', 'normal', 'depth']) -> Dict[str, torch.Tensor]:
+        """
+        Render a batch of representations.
+
+        Args:
+            reps: The dictionary of lists of representations.
+            extrinsics: The [N x 4 x 4] tensor of extrinsics.
+            intrinsics: The [N x 3 x 3] tensor of intrinsics.
+            return_types: vary in ['mask', 'normal', 'depth', 'normal_map', 'color']
+            
+        Returns: 
+            a dict with
+                reg_loss : [N] tensor of regularization losses
+                mask : [N x 1 x H x W] tensor of rendered masks
+                normal : [N x 3 x H x W] tensor of rendered normals
+                depth : [N x 1 x H x W] tensor of rendered depths
+        """
+        ret = {k : [] for k in return_types}
+        for i, rep in enumerate(reps):
+            out_dict = self.renderer.render(rep, extrinsics[i], intrinsics[i], return_types=return_types)
+            for k in out_dict:
+                ret[k].append(out_dict[k][None] if k in ['mask', 'depth'] else out_dict[k])
+        for k in ret:
+            ret[k] = torch.stack(ret[k])
+        return ret
+    
+    @staticmethod
+    def _tsdf_reg_loss(rep: MeshExtractResult, depth_map: torch.Tensor, extrinsics: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor:
+        # Calculate tsdf
+        with torch.no_grad():
+            # Project points to camera and calculate pseudo-sdf as difference between gt depth and projected depth
+            projected_pts, pts_depth = utils3d.torch.project_cv(extrinsics=extrinsics, intrinsics=intrinsics, points=rep.tsdf_v)
+            projected_pts = (projected_pts - 0.5) * 2.0
+            depth_map_res = depth_map.shape[1]
+            gt_depth = torch.nn.functional.grid_sample(depth_map.reshape(1, 1, depth_map_res, depth_map_res), 
+            projected_pts.reshape(1, 1, -1, 2), mode='bilinear', padding_mode='border', align_corners=True)
+            pseudo_sdf = gt_depth.flatten() - pts_depth.flatten()
+            # Truncate pseudo-sdf
+            delta = 1 / rep.res * 3.0
+            trunc_mask = pseudo_sdf > -delta
+        
+        # Loss
+        gt_tsdf = pseudo_sdf[trunc_mask]
+        tsdf = rep.tsdf_s.flatten()[trunc_mask]
+        gt_tsdf = torch.clamp(gt_tsdf, -delta, delta)
+        return torch.mean((tsdf - gt_tsdf) ** 2)
+    
+    def _calc_tsdf_loss(self, reps : list[MeshExtractResult], depth_maps, extrinsics, intrinsics) -> torch.Tensor:
+        tsdf_loss = 0.0
+        for i, rep in enumerate(reps):
+            tsdf_loss += self._tsdf_reg_loss(rep, depth_maps[i], extrinsics[i], intrinsics[i])
+        return tsdf_loss / len(reps)
+    
+    @torch.no_grad()
+    def _flip_normal(self, normal: torch.Tensor, extrinsics: torch.Tensor, intrinsics: torch.Tensor) -> torch.Tensor:
+        """
+        Flip normal to align with camera.
+        """
+        normal = normal * 2.0 - 1.0
+        R = torch.zeros_like(extrinsics)
+        R[:, :3, :3] = extrinsics[:, :3, :3]
+        R[:, 3, 3] = 1.0
+        view_dir = utils3d.torch.unproject_cv(
+            utils3d.torch.image_uv(*normal.shape[-2:], device=self.device).reshape(1, -1, 2),
+            torch.ones(*normal.shape[-2:], device=self.device).reshape(1, -1),
+            R, intrinsics
+        ).reshape(-1, *normal.shape[-2:], 3).permute(0, 3, 1, 2)
+        unflip = (normal * view_dir).sum(1, keepdim=True) < 0
+        normal *= unflip * 2.0 - 1.0
+        return (normal + 1.0) / 2.0
+    
+    def _perceptual_loss(self, gt: torch.Tensor, pred: torch.Tensor, name: str) -> Dict[str, torch.Tensor]:
+        """
+        Combination of L1, SSIM, and LPIPS loss.
+        """
+        if gt.shape[1] != 3:
+            assert gt.shape[-1] == 3
+            gt = gt.permute(0, 3, 1, 2)
+        if pred.shape[1] != 3:
+            assert pred.shape[-1] == 3
+            pred = pred.permute(0, 3, 1, 2)
+        terms = {
+            f"{name}_loss" : l1_loss(gt, pred),
+            f"{name}_loss_ssim" : 1 - ssim(gt, pred),
+            f"{name}_loss_lpips" : lpips(gt, pred)
+        }
+        terms[f"{name}_loss_perceptual"] = terms[f"{name}_loss"] + terms[f"{name}_loss_ssim"] * self.lambda_ssim + terms[f"{name}_loss_lpips"] * self.lambda_lpips
+        return terms
+    
+    def geometry_losses(
+        self,
+        reps: List[MeshExtractResult],
+        mesh: List[Dict],
+        normal_map: torch.Tensor,
+        extrinsics: torch.Tensor,
+        intrinsics: torch.Tensor,
+    ):
+        with torch.no_grad():
+            gt_meshes = []
+            for i in range(len(reps)):
+                gt_mesh = MeshExtractResult(mesh[i]['vertices'].to(self.device), mesh[i]['faces'].to(self.device))
+                gt_meshes.append(gt_mesh)
+            target = self._render_batch(gt_meshes, extrinsics, intrinsics, return_types=['mask', 'depth', 'normal'])
+            target['normal'] = self._flip_normal(target['normal'], extrinsics, intrinsics)
+                
+        terms = edict(geo_loss = 0.0)
+        if self.lambda_tsdf > 0:
+            tsdf_loss = self._calc_tsdf_loss(reps, target['depth'], extrinsics, intrinsics)
+            terms['tsdf_loss'] = tsdf_loss
+            terms['geo_loss'] += tsdf_loss * self.lambda_tsdf
+        
+        return_types = ['mask', 'depth', 'normal', 'normal_map'] if self.use_color else ['mask', 'depth', 'normal']
+        buffer = self._render_batch(reps, extrinsics, intrinsics, return_types=return_types)
+        
+        success_mask = torch.tensor([rep.success for rep in reps], device=self.device)
+        if success_mask.sum() != 0:
+            for k, v in buffer.items():
+                buffer[k] = v[success_mask]
+            for k, v in target.items():
+                target[k] = v[success_mask]
+            
+            terms['mask_loss'] = l1_loss(buffer['mask'], target['mask']) 
+            if self.depth_loss_type == 'l1':
+                terms['depth_loss'] = l1_loss(buffer['depth'] * target['mask'], target['depth'] * target['mask'])
+            elif self.depth_loss_type == 'smooth_l1':
+                terms['depth_loss'] = smooth_l1_loss(buffer['depth'] * target['mask'], target['depth'] * target['mask'], beta=1.0 / (2 * reps[0].res))
+            else:
+                raise ValueError(f"Unsupported depth loss type: {self.depth_loss_type}")
+            terms.update(self._perceptual_loss(buffer['normal'] * target['mask'], target['normal'] * target['mask'], 'normal'))
+            terms['geo_loss'] = terms['geo_loss'] + terms['mask_loss'] + terms['depth_loss'] * self.lambda_depth + terms['normal_loss_perceptual']
+            if self.use_color and normal_map is not None:
+                terms.update(self._perceptual_loss(normal_map[success_mask], buffer['normal_map'], 'normal_map'))
+                terms['geo_loss'] = terms['geo_loss'] + terms['normal_map_loss_perceptual'] * self.lambda_color
+                
+        return terms
+      
+    def color_losses(self, reps, image, alpha, extrinsics, intrinsics):
+        terms = edict(color_loss = torch.tensor(0.0, device=self.device))
+        buffer = self._render_batch(reps, extrinsics, intrinsics, return_types=['color'])
+        success_mask = torch.tensor([rep.success for rep in reps], device=self.device)
+        if success_mask.sum() != 0:
+            terms.update(self._perceptual_loss(image * alpha[:, None][success_mask], buffer['color'][success_mask], 'color'))
+            terms['color_loss'] = terms['color_loss'] + terms['color_loss_perceptual'] * self.lambda_color
+        return terms
+    
+    def training_losses(
+        self,
+        latents: SparseTensor,
+        image: torch.Tensor,
+        alpha: torch.Tensor,
+        mesh: List[Dict],
+        extrinsics: torch.Tensor,
+        intrinsics: torch.Tensor,
+        normal_map: torch.Tensor = None,
+    ) -> Tuple[Dict, Dict]:
+        """
+        Compute training losses.
+
+        Args:
+            latents: The [N x * x C] sparse latents
+            image: The [N x 3 x H x W] tensor of images.
+            alpha: The [N x H x W] tensor of alpha channels.
+            mesh: The list of dictionaries of meshes.
+            extrinsics: The [N x 4 x 4] tensor of extrinsics.
+            intrinsics: The [N x 3 x 3] tensor of intrinsics.
+
+        Returns:
+            a dict with the key "loss" containing a scalar tensor.
+            may also contain other keys for different terms.
+        """
+        reps = self.training_models['decoder'](latents)
+        self.renderer.rendering_options.resolution = image.shape[-1]
+        
+        terms = edict(loss = 0.0, rec = 0.0)
+        
+        terms['reg_loss'] = sum([rep.reg_loss for rep in reps]) / len(reps)
+        terms['loss'] = terms['loss'] + terms['reg_loss']
+        
+        geo_terms = self.geometry_losses(reps, mesh, normal_map, extrinsics, intrinsics)
+        terms.update(geo_terms)
+        terms['loss'] = terms['loss'] + terms['geo_loss']
+                
+        if self.use_color:
+            color_terms = self.color_losses(reps, image, alpha, extrinsics, intrinsics)
+            terms.update(color_terms)
+            terms['loss'] = terms['loss'] + terms['color_loss']
+             
+        return terms, {}
+    
+    @torch.no_grad()
+    def run_snapshot(
+        self,
+        num_samples: int,
+        batch_size: int,
+        verbose: bool = False,
+    ) -> Dict:
+        dataloader = DataLoader(
+            copy.deepcopy(self.dataset),
+            batch_size=batch_size,
+            shuffle=True,
+            num_workers=0,
+            collate_fn=self.dataset.collate_fn if hasattr(self.dataset, 'collate_fn') else None,
+        )
+
+        # inference
+        ret_dict = {}
+        gt_images = []
+        gt_normal_maps = []
+        gt_meshes = []
+        exts = []
+        ints = []
+        reps = []
+        for i in range(0, num_samples, batch_size):
+            batch = min(batch_size, num_samples - i)
+            data = next(iter(dataloader))
+            args = recursive_to_device(data, 'cuda')
+            gt_images.append(args['image'] * args['alpha'][:, None])
+            if self.use_color and 'normal_map' in data:
+                gt_normal_maps.append(args['normal_map'])
+            gt_meshes.extend(args['mesh'])
+            exts.append(args['extrinsics'])
+            ints.append(args['intrinsics'])
+            reps.extend(self.models['decoder'](args['latents']))
+        gt_images = torch.cat(gt_images, dim=0)
+        ret_dict.update({f'gt_image': {'value': gt_images, 'type': 'image'}})
+        if self.use_color and gt_normal_maps:
+            gt_normal_maps = torch.cat(gt_normal_maps, dim=0)
+            ret_dict.update({f'gt_normal_map': {'value': gt_normal_maps, 'type': 'image'}})
+
+        # render single view
+        exts = torch.cat(exts, dim=0)
+        ints = torch.cat(ints, dim=0)
+        self.renderer.rendering_options.bg_color = (0, 0, 0)
+        self.renderer.rendering_options.resolution = gt_images.shape[-1]
+        gt_render_results = self._render_batch([
+            MeshExtractResult(vertices=mesh['vertices'].to(self.device), faces=mesh['faces'].to(self.device))
+            for mesh in gt_meshes
+        ], exts, ints, return_types=['normal'])
+        ret_dict.update({f'gt_normal': {'value': self._flip_normal(gt_render_results['normal'], exts, ints), 'type': 'image'}})
+        return_types = ['normal']
+        if self.use_color:
+            return_types.append('color')
+            if 'normal_map' in data:
+                return_types.append('normal_map')
+        render_results = self._render_batch(reps, exts, ints, return_types=return_types)
+        ret_dict.update({f'rec_normal': {'value': render_results['normal'], 'type': 'image'}})
+        if 'color' in return_types:
+            ret_dict.update({f'rec_image': {'value': render_results['color'], 'type': 'image'}})
+        if 'normal_map' in return_types:
+            ret_dict.update({f'rec_normal_map': {'value': render_results['normal_map'], 'type': 'image'}})
+
+        # render multiview
+        self.renderer.rendering_options.resolution = 512
+        ## 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)]
+
+        ## render each view
+        multiview_normals = []
+        multiview_normal_maps = []
+        miltiview_images = []
+        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)
+            extrinsics = extrinsics.unsqueeze(0).expand(num_samples, -1, -1)
+            intrinsics = intrinsics.unsqueeze(0).expand(num_samples, -1, -1)
+            render_results = self._render_batch(reps, extrinsics, intrinsics, return_types=return_types)
+            multiview_normals.append(render_results['normal'])
+            if 'color' in return_types:
+                miltiview_images.append(render_results['color'])
+            if 'normal_map' in return_types:
+                multiview_normal_maps.append(render_results['normal_map'])
+
+        ## Concatenate views
+        multiview_normals = torch.cat([
+            torch.cat(multiview_normals[:2], dim=-2),
+            torch.cat(multiview_normals[2:], dim=-2),
+        ], dim=-1)
+        ret_dict.update({f'multiview_normal': {'value': multiview_normals, 'type': 'image'}})
+        if 'color' in return_types:
+            miltiview_images = torch.cat([
+                torch.cat(miltiview_images[:2], dim=-2),
+                torch.cat(miltiview_images[2:], dim=-2),
+            ], dim=-1)
+            ret_dict.update({f'multiview_image': {'value': miltiview_images, 'type': 'image'}})
+        if 'normal_map' in return_types:
+            multiview_normal_maps = torch.cat([
+                torch.cat(multiview_normal_maps[:2], dim=-2),
+                torch.cat(multiview_normal_maps[2:], dim=-2),
+            ], dim=-1)
+            ret_dict.update({f'multiview_normal_map': {'value': multiview_normal_maps, 'type': 'image'}})
+                            
+        return ret_dict