Model.py

"""NVGS/Model.py: Implementation of the model for 3D Gaussian Splatting."""

from pathlib import Path
from typing import Any
import torch
import numpy as np

import Framework
from Logging import Logger
from Methods.Base.Model import BaseModel
from Methods.NVGSViewer.utils import extract_upper_triangular_matrix, build_covariances, load_ply
from CudaUtils.MortonEncoding import morton_encode
from Methods.NVGSViewer.NVGSCudaBackend import initialize_tcnn_models
import Thirdparty.TinyCudaNN as tcnn




class EmbeddingMLP(torch.nn.Module):
    def __init__(self,
        n_outputs_embed: int = 6, # embedding dims for gaussian encoding
        n_inputs_embed: int = 8, # takes last n inputs and embed them using secondary mlp
        n_hidden_layers_embed: int = 2,
        n_neurons_embed: int = 32
    ):
        super().__init__()

        self.embedding_mlp = tcnn.NetworkWithInputEncoding(
            n_input_dims=n_inputs_embed,
            n_output_dims=n_outputs_embed,
            encoding_config={"otype": "Identity"},
            network_config={
                'otype': 'FullyFusedMLP',
                'activation': 'ReLU',
                'output_activation': 'None',
                'n_neurons': n_neurons_embed,
                'n_hidden_layers': n_hidden_layers_embed,
            }   
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.embedding_mlp(inputs)
        


class VisibilityMLP(torch.nn.Module):
    """A module that uses an MLP to predict visibility based on frequency-encoded features and latent embeddings."""

    def __init__(
        self,
        n_inputs_main: int = 16,
        n_hidden_layers_main: int = 2,
        n_neurons_main: int = 32,
    ) -> None:
        super().__init__()
        # setup tcnn model
        self.visibility_mlp = tcnn.NetworkWithInputEncoding(
            n_input_dims=n_inputs_main,
            n_output_dims=1,
            encoding_config={"otype": "Identity" },
            network_config={
                'otype': 'FullyFusedMLP',
                'activation': 'ReLU',
                'output_activation': 'None',
                'n_neurons': n_neurons_main,
                'n_hidden_layers': n_hidden_layers_main,
            }   
        )

    def get_tcnn_config(self, ) -> tuple[int, int, dict[str, Any], dict[str, Any]]:
        return (
            self.visibility_mlp.n_input_dims,
            self.visibility_mlp.n_output_dims,
            self.visibility_mlp.encoding_config,
            self.visibility_mlp.network_config
        )

    def forward(self, inputs: torch.Tensor) -> torch.Tensor:
        return self.visibility_mlp(inputs)


class ComposedGaussians(torch.nn.Module):
    """Stores a set of points with 3D Gaussian extent."""

    def __init__(self, use_tcnn: bool, debug: bool) -> None:
        super().__init__()
        self.use_tcnn = use_tcnn
        self.tcnn_initialized = False
        if self.use_tcnn:
            Logger.log_info('using tinycudann mlp implementation')
        else:
            Logger.log_info('using naive mlp implementation')
        self.debug = debug
        if self.debug:
            Logger.log_warning('enabling debug mode: only a single asset and instance will be used, and morton ordering will be disabled')
        # general data
        self.register_buffer('gaussian_counts', None)
        self.register_buffer('instance_counts', None)
        self.register_buffer('virtual_gaussian_counts', None)
        self.register_buffer('gaussian_offsets', None)
        self.register_buffer('instance_offsets', None)
        self.register_buffer('virtual_gaussian_offsets', None)
        # per asset data
        self.register_buffer('base_distances', None)
        self.register_buffer('mlp_weights', None)
        self.register_buffer('scene_centers', None)
        self.register_buffer('scene_extents', None)
        # per Gaussian data
        self.register_buffer('means', None)
        self.register_buffer('covariances', None)
        self.register_buffer('opacities', None)
        self.register_buffer('sh_coefficients', None)
        self.register_buffer('embeddings', None)
        # per instance data
        self.register_buffer('instance_transforms', None)
        # constants
        self.default_minimum_distance = 0.5
        self.default_maximum_distance = 10.0

    @property
    def n_assets(self) -> int:
        """Returns the number of assets."""
        return self.gaussian_counts.numel() if self.gaussian_counts is not None else 0

    @property
    def n_instances(self) -> int:
        """Returns the number of instances."""
        return self.instance_offsets[-1].item() if self.instance_offsets is not None else 0

    @property
    def n_gaussians(self) -> int:
        """Returns the number of Gaussians."""
        return self.gaussian_offsets[-1].item() if self.gaussian_offsets is not None else 0

    @property
    def n_virtual_gaussians(self) -> int:
        """Returns the number of virtual Gaussians."""
        return self.virtual_gaussian_offsets[-1].item() if self.virtual_gaussian_offsets is not None else 0

    @property
    def available_sh_bases(self) -> int:
        """Returns the number of available spherical harmonics bases."""
        return self.sh_coefficients.shape[1]

    def add_asset(self, base_path: Path, ply_suffix: str, embedding_suffix:str, tcnn_lod_model_suffix: str, instances: list[dict[str, Any]], sh_degree=0) -> None:
        """Appends a point cloud asset to the model."""
        asset_name = base_path.name
        if self.debug and self.n_assets == 1:
            Logger.log_warning(f'skipping asset "{asset_name}" because debug mode only supports a single asset')
            return

        ply_path = base_path / ply_suffix
        if not ply_path.is_file():
            Logger.log_warning(f'skipping asset "{asset_name}" as the ply file "{ply_path}" does not exist')
            return
    
        embedding_path = base_path / embedding_suffix
        tcnn_lod_model_path = base_path / tcnn_lod_model_suffix
          
        use_tcnn_for_asset = True
        if not embedding_path.is_file() or not tcnn_lod_model_path.is_file():
            Logger.log_warning(f'skipping asset "{asset_name}" as the embedding file "{embedding_path}" or the tcnn checkpoint does not exist')
            Logger.log_warning('falling back to non nvgs-enabled assets')
            use_tcnn_for_asset = False
           
        # loads ply + embedding_file
        means, scales, rotations, opacities, sh_coefficients_0 = load_ply(ply_path)
  
        # convert to torch
        means = torch.as_tensor(means, dtype=torch.float32, device='cuda').contiguous()
        scales = torch.as_tensor(scales, dtype=torch.float32, device='cuda').contiguous()
        rotations = torch.as_tensor(rotations, dtype=torch.float32, device='cuda').contiguous()
        opacities = torch.as_tensor(opacities, dtype=torch.float32, device='cuda').contiguous()
        sh_coefficients_0 = torch.as_tensor(sh_coefficients_0, dtype=torch.float32, device='cuda')

        # Calculate scene extent for normalization
        scene_center, scene_extent = self.get_scene_stats(means)
        scene_center = scene_center[None].cpu()
        scene_extent = scene_extent[None].cpu()
        
        # Adds additional SH coefficients -> we add zeros as dummy values for higher SH degrees as assets were trained with 0 degrees
        # This helps us test compute for higher degrees
        if sh_degree > 0:
            n_bases_rest = (sh_degree + 1) ** 2 - 1
            sh_coefficients_rest = torch.zeros((means.shape[0], n_bases_rest, 3), dtype=torch.float32, device='cuda')
            sh_coefficients = torch.cat((sh_coefficients_0[:, None, :], sh_coefficients_rest), dim=1)
            del sh_coefficients_rest
        else:
            sh_coefficients = sh_coefficients_0[:, None, :].contiguous()
        del sh_coefficients_0

        # Apply activations
        scales = torch.exp(scales)
        rotations = torch.nn.functional.normalize(rotations, dim=-1)
        covariances = extract_upper_triangular_matrix(build_covariances(scales, rotations))
        opacities = torch.sigmoid(opacities)



       

        n_gaussians_asset = torch.tensor((means.shape[0],), dtype=torch.int32, device='cpu')
        Logger.log_info(f'adding asset "{asset_name}" with {n_gaussians_asset.item():,} Gaussians')
       
        # setup visibility configuration
        visibility_config_path = base_path / 'asset_config.pt'
        visibility_config = None
        if visibility_config_path.is_file():
            visibility_config = torch.load(visibility_config_path, map_location='cpu')
            if visibility_config['n_gaussians'] != n_gaussians_asset.item():
                Logger.log_warning(f'visibility configuration for asset "{asset_name}" does not match the number of Gaussians ({visibility_config["n_gaussians"]:,} != {n_gaussians_asset.item():,}) -> using default values instead')
                visibility_config = None
        if visibility_config is None:
            Logger.log_info(f'using default visibility configuration for asset "{asset_name}" (base distance {self.default_minimum_distance:.3f}) (maximum distance {self.default_maximum_distance:.3f})')
            minimum_distance = torch.tensor((self.default_minimum_distance,), dtype=torch.float32, device='cpu')[None]
            maximum_distance = torch.tensor((self.default_maximum_distance,), dtype=torch.float32, device='cpu')[None]
        else:
            minimum_distance = torch.tensor((visibility_config['minimum_distance'].item(),), dtype=torch.float32, device='cpu')[None]
            maximum_distance = torch.tensor((visibility_config['maximum_distance'].item(),), dtype=torch.float32, device='cpu')[None]
            Logger.log_info(f'found visibility configuration for asset "{asset_name}" (base distance {minimum_distance.item():.3f})')

        # Checks if tcnn should be used in general, and for asset specificially
        use_nvgs = torch.tensor([True], dtype=torch.bool, device='cpu')[None]
        if self.use_tcnn and use_tcnn_for_asset:
            # Load the tiny-cuda-nn visibility network and keep only runtime
            # parameters needed during inference.
    
            tcnn_lod_model_params = torch.load(tcnn_lod_model_path, map_location='cpu')

            # Calculate embeddings from the checkpointed embedding MLP.
            embedding_mlp = EmbeddingMLP().to(opacities.device)
            embedding_mlp_params = {
                'embedding_mlp.params': tcnn_lod_model_params['embedding_mlp.params']
            }
            embedding_mlp.load_state_dict(embedding_mlp_params, strict=True)
            embedding_mlp.eval()

            opacities_mlp = opacities.reshape(-1, 1) * 2.0 - 1.0
            rotations_mlp = rotations.clone()
            rotations_mlp[rotations[...,0] < 0] *= -1

            # Embedding MLP expects 8 channels: 3 scales, 4 rotations, 1 opacity.
            embedding_inputs = torch.cat((
                opacities_mlp, #normalize
                rotations_mlp,
                scales    
            ), dim=-1).to(torch.float32).contiguous()
            del opacities_mlp, rotations_mlp

            with torch.no_grad():
                embeddings = embedding_mlp(embedding_inputs)
            del embedding_inputs, embedding_mlp, embedding_mlp_params


            # This removes embedding mlp from network 
            if 'embedding_mlp.params' in tcnn_lod_model_params:
                del tcnn_lod_model_params['embedding_mlp.params']
            tcnn_lod_model = VisibilityMLP()
            tcnn_lod_model.load_state_dict(tcnn_lod_model_params, strict=True)
            Logger.log_info(tcnn_lod_model)

            mlp_weights = tcnn_lod_model.visibility_mlp.params.data.to(opacities).contiguous()
            Logger.log_info(f'loaded visibility MLP weights for asset "{asset_name}" with shape {tuple(mlp_weights.shape)}')
            del tcnn_lod_model, tcnn_lod_model_params
       
            
            
        else:
           use_nvgs = torch.tensor([False], dtype=torch.bool, device='cpu')[None]
           embeddings = torch.zeros(means.shape[0], 6, dtype=torch.float16, device='cuda')
           mlp_weights = torch.zeros(2048, dtype=torch.float32, device='cuda')

        # Morton ordering improves cache locality for downstream rasterization.
        if not self.debug:
            morton_encoding = morton_encode(means)
            order = torch.argsort(morton_encoding)
            means = means[order].contiguous()
            # scales = scales[order].contiguous()
            # rotations = rotations[order].contiguous()
            covariances = covariances[order].contiguous()
            opacities = opacities[order].contiguous()
            sh_coefficients = sh_coefficients[order].contiguous()
            embeddings = embeddings[order].contiguous()
            del morton_encoding, order

        # covariances are precomputed and only instance transforms are needed afterwards.
        del scales, rotations

        # setup instance data
        n_instances_asset = torch.tensor((len(instances),), dtype=torch.int32, device='cpu')
        instance_transforms = self.extract_instance_transforms(instances)

        # register asset and instance data
        current_asset_idx = self.n_assets
        if current_asset_idx == 0:
            # first asset, initialize buffers
            self.gaussian_counts = n_gaussians_asset
            self.instance_counts = n_instances_asset
            self.virtual_gaussian_counts = n_instances_asset * n_gaussians_asset
            self.gaussian_offsets = torch.tensor((0, self.gaussian_counts.item()), dtype=torch.int32, device='cpu')
            self.instance_offsets = torch.tensor((0, self.instance_counts.item()), dtype=torch.int32, device='cpu')
            self.virtual_gaussian_offsets = torch.tensor((0, self.virtual_gaussian_counts.item()), dtype=torch.int32, device='cpu')
            self.minimum_distances = minimum_distance
            self.maximum_distances = maximum_distance
            self.use_neural_culling_for_asset = use_nvgs
            self.mlp_weights = mlp_weights[None]
            self.means = means
            # self.scalings = scales
            # self.rotations = rotations
            self.scene_center = scene_center
            self.scene_extent = scene_extent
            self.covariances = covariances
            self.opacities = opacities
            self.sh_coefficients = sh_coefficients
            self.embeddings = embeddings
            self.instance_transforms = instance_transforms
        else:
            # add to existing buffers
            self.gaussian_counts = torch.cat((self.gaussian_counts, n_gaussians_asset), dim=0)
            self.instance_counts = torch.cat((self.instance_counts, n_instances_asset), dim=0)
            self.virtual_gaussian_counts = torch.cat((self.virtual_gaussian_counts, n_instances_asset * n_gaussians_asset), dim=0)
            self.gaussian_offsets = torch.cat((self.gaussian_offsets, self.gaussian_offsets[-1:] + self.gaussian_counts[-1:]), dim=0)
            self.instance_offsets = torch.cat((self.instance_offsets, self.instance_offsets[-1:] + self.instance_counts[-1:]), dim=0)
            self.virtual_gaussian_offsets = torch.cat((self.virtual_gaussian_offsets, self.virtual_gaussian_offsets[-1:] + self.virtual_gaussian_counts[-1:]), dim=0)
            self.minimum_distances = torch.cat((self.minimum_distances, minimum_distance), dim=0)
            self.maximum_distances = torch.cat((self.maximum_distances, maximum_distance), dim=0)
            self.use_neural_culling_for_asset = torch.cat(( self.use_neural_culling_for_asset , use_nvgs), dim=0)
            self.mlp_weights = torch.cat((self.mlp_weights, mlp_weights[None]), dim=0)
            self.means = torch.cat((self.means, means), dim=0)
            # self.scalings = torch.cat((self.scalings, scales), dim=0)
            # self.rotations = torch.cat((self.rotations, rotations), dim=0)
            self.scene_center = torch.cat((self.scene_center, scene_center), dim=0)
            self.scene_extent = torch.cat((self.scene_extent, scene_extent), dim=0)
            self.covariances = torch.cat((self.covariances, covariances,), dim=0)
            self.opacities = torch.cat((self.opacities, opacities), dim=0)
            self.sh_coefficients = torch.cat((self.sh_coefficients, sh_coefficients), dim=0)
            self.embeddings = torch.cat((self.embeddings, embeddings), dim=0)
            self.instance_transforms = torch.cat((self.instance_transforms, instance_transforms), dim=0)

        torch.cuda.empty_cache()

    @staticmethod
    def get_scene_stats(gaussian_xyz):
        """Returns both scene center and scene extent"""
        scene_min = torch.min(gaussian_xyz, dim=0).values
        scene_max = torch.max(gaussian_xyz, dim=0).values
        scene_center = (scene_max + scene_min) / 2.0
        scene_extent = torch.max(scene_max - scene_min)
        return scene_center, scene_extent


    @staticmethod
    def extract_instance_transforms(instances: list[dict[str, Any]]) -> torch.Tensor:
        """ Extracts instance transforms from a list of instances."""
        instance_transforms = []
        for instance in instances:
            # setup model transformation data
            rotation = instance['rotation']
            if len(rotation) != 4:
                raise Framework.ModelError(f'rotation must be a list of 4 floats, got "{rotation}" instead')
            # Blender to Colmap q = [sqrt(0.5), sqrt(0.5), 0, 0], rest is simplification of q_mult(q, rotation)
            w, x, y, z = rotation
            rotation = np.sqrt(0.5) * np.array([w - x, x + w, y - z, z + y])
            rotation = torch.tensor(rotation, dtype=torch.float32, device='cuda')
            rotation = torch.nn.functional.normalize(rotation, dim=-1)
            scale = instance['scale']
            if not isinstance(scale, (int, float)):
                raise Framework.ModelError(f'scale must be a float or int, got "{scale}" instead')
            scale = torch.tensor((scale,), dtype=torch.float32, device='cuda')
            position = instance['position']
            if len(position) != 3:
                raise Framework.ModelError(f'position must be a list of 3 floats, got "{position}" instead')
            x, y, z = position
            position = [x, -z, y]  # Blender to Colmap
            position = torch.tensor(position, dtype=torch.float32, device='cuda')
            instance_transform = torch.cat((rotation, scale, position), dim=0)
            instance_transforms.append(instance_transform)
        # stack all instance transforms
        instance_transforms = torch.stack(instance_transforms, dim=0)
        return instance_transforms

    def tcnn_initialization(self) -> None:
        """Initializes the tinycudann models."""
        # The backend keeps one shared TCNN module per asset, initialized from
        # the packed per-asset weight tensor.
        initialize_tcnn_models(self.mlp_weights, *VisibilityMLP().get_tcnn_config())
        # The C++ backend keeps a compact half-precision copy of parameters.
        # Drop the Python-side float32 tensor to reduce persistent VRAM.
        self.mlp_weights = torch.empty(1, 1, dtype=torch.float32, device='cuda')
        self.tcnn_initialized = True


@Framework.Configurable.configure(
    USE_TCNN=True,
    DEBUG=False,
)
class NVGSModel(BaseModel):
    """Defines the 3DGS model."""

    def __init__(self, name: str = None) -> None:
        super().__init__(name)
        self.composed_gaussians: ComposedGaussians | None = None

    def build(self) -> 'NVGSModel':
        """Builds the model."""
        self.composed_gaussians = ComposedGaussians(self.USE_TCNN, self.DEBUG)
        return self