cube/cube3d/model/autoencoder/one_d_autoencoder.py

import logging
import sys
from dataclasses import dataclass, field
from functools import partial
from typing import List, Optional, Tuple

import numpy as np
import torch
import torch.nn as nn
from skimage import measure
from torch.nn import functional as F
from tqdm import tqdm

from cube3d.model.autoencoder.embedder import PhaseModulatedFourierEmbedder
from cube3d.model.autoencoder.grid import (
    generate_dense_grid_points,
    marching_cubes_with_warp,
)
from cube3d.model.autoencoder.spherical_vq import SphericalVectorQuantizer
from cube3d.model.transformers.attention import (
    EncoderCrossAttentionLayer,
    EncoderLayer,
    init_linear,
    init_tfixup,
)
from cube3d.model.transformers.norm import LayerNorm


def init_sort(x):
    """
    Sorts the input tensor `x` based on its pairwise distances to the first element.
    This function computes the pairwise distances between all elements in `x` and the
    first element of `x`. It then sorts the elements of `x` in ascending order of
    their distances to the first element.
    Args:
        x (torch.Tensor): A 2D tensor where each row represents a data point.
    Returns:
        torch.Tensor: A tensor containing the rows of `x` sorted by their distances
        to the first row of `x`.
    """

    distances = torch.cdist(x, x[:1])
    _, indices = torch.sort(distances.squeeze(), dim=0)
    x = x[indices]
    return x


class MLPEmbedder(nn.Module):
    def __init__(self, in_dim: int, embed_dim: int, bias: bool = True):
        super().__init__()
        self.in_layer = nn.Linear(in_dim, embed_dim, bias=bias)
        self.silu = nn.SiLU()
        self.out_layer = nn.Linear(embed_dim, embed_dim, bias=bias)

        self.apply(partial(init_linear, embed_dim=embed_dim))

    def forward(self, x: torch.Tensor) -> torch.Tensor:
        return self.out_layer(self.silu(self.in_layer(x)))


class OneDEncoder(nn.Module):
    def __init__(
        self,
        embedder,
        num_latents: int,
        point_feats: int,
        embed_point_feats: bool,
        width: int,
        num_heads: int,
        num_layers: int,
        with_cls_token: bool = False,
        cross_attention_levels: Optional[List[int]] = None,
        eps: float = 1e-6,
    ) -> None:
        """
        Initializes the OneDEncoder model.
        Args:
            embedder: An embedding module that provides the input embedding functionality.
            num_latents (int): The number of latent variables.
            point_feats (int): The number of point features.
            embed_point_feats (bool): Whether to embed point features or not.
            width (int): The width of the embedding dimension.
            num_heads (int): The number of attention heads.
            num_layers (int): The number of encoder layers.
            with_cls_token (bool, optional): Whether to include a classification token like in Vision Transformers (ViT). Defaults to False.
            cross_attention_levels (Optional[List[int]], optional): The indices of layers where cross-attention is applied. Defaults to None.
            eps (float, optional): A small value added for numerical stability in normalization layers. Defaults to 1e-6.
        Returns:
            None
        """
        super().__init__()

        self.embedder = embedder

        # add cls token like ViT
        self.with_cls_token = with_cls_token
        if self.with_cls_token:
            query = torch.empty((1 + num_latents, width))
        else:
            query = torch.empty((num_latents, width))

        # initialize then sort query to potentially get better ordering
        query.uniform_(-1.0, 1.0)
        query = init_sort(query)

        # set parameter
        self.query = nn.Parameter(query)

        self.embed_point_feats = embed_point_feats
        in_dim = (
            self.embedder.out_dim * 2
            if self.embed_point_feats
            else self.embedder.out_dim + point_feats
        )
        self.feat_in = MLPEmbedder(in_dim, embed_dim=width)

        if cross_attention_levels is None:
            cross_attention_levels = [0]

        self.blocks = nn.ModuleList()
        for i in range(num_layers):
            if i in cross_attention_levels:
                self.blocks.append(
                    EncoderCrossAttentionLayer(
                        embed_dim=width,
                        num_heads=num_heads,
                        eps=eps,
                    )
                )
            else:
                self.blocks.append(
                    EncoderLayer(embed_dim=width, num_heads=num_heads, eps=eps)
                )
        self.ln_f = LayerNorm(width, eps=eps)

        init_tfixup(self, num_layers)

    def _forward(self, h, data, attn_mask=None):
        """
        Forward pass for the autoencoder model.

        Args:
            h (torch.Tensor): The input tensor to be processed, typically representing
                the hidden state or intermediate representation.
            data (torch.Tensor): The input data tensor to be transformed by the feature
                extraction layer and used in cross-attention layers.
            attn_mask (torch.Tensor, optional): An optional attention mask tensor to be
                used in attention layers for masking specific positions. Defaults to None.
        Returns:
            torch.Tensor: The output tensor after processing through the layers and
            applying final normalization.
        """

        data = self.feat_in(data)

        for block in self.blocks:
            if isinstance(block, EncoderCrossAttentionLayer):
                h = block(h, data)
            else:
                h = block(h, attn_mask=attn_mask)

        h = self.ln_f(h)
        return h

    def forward(
        self, pts: torch.Tensor, feats: torch.Tensor
    ) -> Tuple[torch.Tensor, list[torch.Tensor]]:
        """
        Forward pass of the 1D autoencoder model.
        Args:
            pts (torch.Tensor): Input tensor representing points with shape (batch_size, num_points, point_dim).
            feats (torch.Tensor): Input tensor representing features with shape (batch_size, num_points, feature_dim).
                                  Can be None if no features are provided.
        Returns:
            Tuple[torch.Tensor, list[torch.Tensor]]:
                - The output tensor after processing the input data.
                - A list of intermediate tensors (if applicable) generated during the forward pass.
        """

        b = pts.shape[0]
        data = self.embedder(pts)

        if feats is not None:
            if self.embed_point_feats:
                feats = self.embedder(feats)
            data = torch.cat([data, feats], dim=-1)

        # prepare query and data
        h = self.query.unsqueeze(0).expand(b, -1, -1)
        return self._forward(h, data, attn_mask=None)


class OneDBottleNeck(nn.Module):
    def __init__(
        self,
        block,
    ) -> None:
        """
        Initializes the OneDBottleNeck class.
        Args:
            block: The building block or module used within the autoencoder.
        """
        super().__init__()

        self.block = block

    def forward(self, h: torch.Tensor) -> Tuple[torch.Tensor, dict]:
        """
        Forward pass of the OneDBottleNeck function.
        Args:
            h (torch.Tensor): Input tensor to the model.
        Returns:
            Tuple[torch.Tensor, dict]: A tuple containing:
                - The transformed tensor `z` after passing through the block (if applicable).
                - A dictionary `ret_dict` containing additional information:
                    - "indices": Indices from the block output (if present).
                    - "z_q": Quantized tensor from the block output (if present).

        """

        z = h
        ret_dict = {}
        if self.block is not None:
            z, d = self.block(z)

            key_mappings = {
                "q": "indices",
                "z_q": "z_q",
            }
            for in_key, out_key in key_mappings.items():
                if in_key in d:
                    ret_dict[out_key] = d[in_key]

        return z, ret_dict


class OneDDecoder(nn.Module):
    def __init__(
        self,
        num_latents: int,
        width: int,
        num_heads: int,
        num_layers: int,
        eps: float = 1e-6,
    ) -> None:
        """
        Initializes the OneDDecoder class.
        Args:
            num_latents (int): The number of latent variables.
            width (int): The width of the embedding dimension.
            num_heads (int): The number of attention heads in each encoder layer.
            num_layers (int): The number of encoder layers.
            eps (float, optional): A small value added for numerical stability. Defaults to 1e-6.
        """
        super().__init__()

        self.register_buffer("query", torch.empty([0, width]), persistent=False)
        self.positional_encodings = nn.Parameter(
            init_sort(F.normalize(torch.empty(num_latents, width).normal_()))
        )
        self.blocks = nn.ModuleList(
            [
                EncoderLayer(embed_dim=width, num_heads=num_heads, eps=eps)
                for _ in range(num_layers)
            ]
        )

        init_tfixup(self, num_layers)

    def _forward(self, h):
        """
        Applies a sequence of operations to the input tensor `h` using the blocks
        defined in the model.
        Args:
            h (torch.Tensor): The input tensor to be processed by the blocks.
        Returns:
            torch.Tensor: The output tensor after applying all blocks sequentially.
        """

        for block in self.blocks:
            h = block(h)
        return h

    def forward(self, z):
        """
        This method processes the input tensor `z` by padding it to a fixed length,
        adding positional encodings, and then passing it through the `_forward` method.

        Args:
            z (torch.Tensor): Input tensor.
        Returns:
            torch.Tensor: Output tensor after processing through the autoencoder.
        Notes:
            - If the `query` attribute has a non-zero shape, the input tensor `z` is padded
              to match the required length using slices of `query`.
            - Positional encodings are added to the padded input tensor before passing it
              to the `_forward` method.
        """

        # pad input to fixed length
        if self.query.shape[0] > 0:
            pad_len = self.query.shape[0] + 1 - z.shape[1]
            paddings = self.query[:pad_len, ...].unsqueeze(0).expand(z.shape[0], -1, -1)
            z = torch.cat([paddings, z], dim=1)
        h = z + self.positional_encodings[: z.shape[1], :].unsqueeze(0).expand(
            z.shape[0], -1, -1
        )

        return self._forward(h)


class OneDOccupancyDecoder(nn.Module):
    def __init__(
        self, embedder, out_features: int, width: int, num_heads: int, eps=1e-6
    ) -> None:
        """
        Initializes the OneDOccupancyDecoder module.
        Args:
            embedder: An embedding module that provides input embeddings.
            out_features (int): The number of output features for the final linear layer.
            width (int): The width of the intermediate layers.
            num_heads (int): The number of attention heads for the cross-attention layer.
            eps (float, optional): A small value added for numerical stability in layer normalization. Defaults to 1e-6.
        """
        super().__init__()

        self.embedder = embedder
        self.query_in = MLPEmbedder(self.embedder.out_dim, width)

        self.attn_out = EncoderCrossAttentionLayer(embed_dim=width, num_heads=num_heads)
        self.ln_f = LayerNorm(width, eps=eps)
        self.c_head = nn.Linear(width, out_features)

    def query(self, queries: torch.Tensor):
        """
        Processes the input tensor through the embedder and query_in layers.
        Args:
            queries (torch.Tensor): A tensor containing the input data to be processed.
        Returns:
            torch.Tensor: The output tensor after being processed by the embedder and query_in layers.
        """

        return self.query_in(self.embedder(queries))

    def forward(self, queries: torch.Tensor, latents: torch.Tensor):
        """
        Defines the forward pass of the model.
        Args:
            queries (torch.Tensor): Input tensor representing the queries.
            latents (torch.Tensor): Input tensor representing the latent representations.
        Returns:
            torch.Tensor: Output tensor after applying the query transformation,
                          attention mechanism, and final processing layers.
        """
        queries = self.query(queries)
        x = self.attn_out(queries, latents)
        x = self.c_head(self.ln_f(x))
        return x


class OneDAutoEncoder(nn.Module):
    @dataclass
    class Config:
        checkpoint_path: str = ""

        # network params
        num_encoder_latents: int = 256
        num_decoder_latents: int = 256
        embed_dim: int = 12
        width: int = 768
        num_heads: int = 12
        out_dim: int = 1
        eps: float = 1e-6

        # grid features embedding
        num_freqs: int = 128
        point_feats: int = 0
        embed_point_feats: bool = False

        num_encoder_layers: int = 1
        encoder_cross_attention_levels: list[int] = field(default_factory=list)
        num_decoder_layers: int = 23

        encoder_with_cls_token: bool = True
        num_codes: int = 16384

    def __init__(self, cfg: Config) -> None:
        """
        Initializes the OneDAutoencoder model.
        Args:
            cfg (Config): Configuration object containing the parameters for the model.
        Attributes:
            cfg (Config): Stores the configuration object.
            embedder (PhaseModulatedFourierEmbedder): Embeds input data using phase-modulated Fourier features.
            encoder (OneDEncoder): Encodes the input data into latent representations.
            bottleneck (OneDBottleNeck): Bottleneck layer containing a spherical vector quantizer for dimensionality reduction.
            decoder (OneDDecoder): Decodes latent representations back into the original data space.
            occupancy_decoder (OneDOccupancyDecoder): Decodes occupancy information from latent representations.
        """

        super().__init__()

        self.cfg = cfg

        self.embedder = PhaseModulatedFourierEmbedder(
            num_freqs=self.cfg.num_freqs, input_dim=3
        )

        self.encoder = OneDEncoder(
            embedder=self.embedder,
            num_latents=self.cfg.num_encoder_latents,
            with_cls_token=self.cfg.encoder_with_cls_token,
            point_feats=self.cfg.point_feats,
            embed_point_feats=self.cfg.embed_point_feats,
            width=self.cfg.width,
            num_heads=self.cfg.num_heads,
            num_layers=self.cfg.num_encoder_layers,
            cross_attention_levels=self.cfg.encoder_cross_attention_levels,
            eps=self.cfg.eps,
        )

        block = SphericalVectorQuantizer(
            self.cfg.embed_dim,
            self.cfg.num_codes,
            self.cfg.width,
            codebook_regularization="kl",
        )
        self.bottleneck = OneDBottleNeck(block=block)

        self.decoder = OneDDecoder(
            num_latents=self.cfg.num_encoder_latents,
            width=self.cfg.width,
            num_heads=self.cfg.num_heads,
            num_layers=self.cfg.num_decoder_layers,
            eps=self.cfg.eps,
        )

        self.occupancy_decoder = OneDOccupancyDecoder(
            embedder=self.embedder,
            out_features=self.cfg.out_dim,
            width=self.cfg.width,
            num_heads=self.cfg.num_heads,
            eps=self.cfg.eps,
        )

    @torch.no_grad()
    def decode_indices(self, shape_ids: torch.Tensor):
        """
        Decodes the given shape indices into latent representations.
        Args:
            shape_ids (torch.Tensor): A tensor containing the shape indices to be decoded.
        Returns:
            torch.Tensor: The decoded latent representations corresponding to the input shape indices.
        """

        z_q = self.bottleneck.block.lookup_codebook(shape_ids)
        latents = self.decode(z_q)
        return latents

    @torch.no_grad()
    def query_embeds(self, shape_ids: torch.Tensor):
        """
        Retrieves the latent embeddings corresponding to the given shape IDs.
        Args:
            shape_ids (torch.Tensor): A tensor containing the IDs of the shapes
                for which the latent embeddings are to be queried.
        Returns:
            torch.Tensor: A tensor containing the latent embeddings retrieved
                from the codebook for the provided shape IDs.
        """

        z_q = self.bottleneck.block.lookup_codebook_latents(shape_ids)
        return z_q

    @torch.no_grad()
    def query_indices(self, shape_embs: torch.Tensor):
        """
        Queries the indices of the quantized embeddings from the bottleneck layer.
        Args:
            shape_embs (torch.Tensor): The input tensor containing shape embeddings
                to be quantized.
        Returns:
            torch.Tensor: A tensor containing the quantized indices.
        """

        _, ret_dict = self.bottleneck.block.quantize(shape_embs)
        return ret_dict["q"]

    def encode(self, x: torch.Tensor, **kwargs):
        """
        Encodes the input tensor using the encoder and bottleneck layers.
        Args:
            x (torch.Tensor): Input tensor with shape (..., N), where the first 3
                dimensions represent points (pts) and the remaining dimensions
                represent features (feats).
            **kwargs: Additional keyword arguments.
        Returns:
            Tuple[torch.Tensor, torch.Tensor, None, dict]: A tuple containing:
                - z_e (torch.Tensor): Encoded tensor before bottleneck processing.
                - z (torch.Tensor): Encoded tensor after bottleneck processing.
                - None: Placeholder for compatibility with other methods.
                - d (dict): Dictionary containing additional information, including:
                    - "z_cls" (torch.Tensor, optional): Class token if
                      `self.cfg.encoder_with_cls_token` is True.
        """

        pts, feats = x[..., :3], x[..., 3:]
        z_e = self.encoder(pts, feats)

        # split class token
        if self.cfg.encoder_with_cls_token:
            z_cls = z_e[:, 0, ...]
            z_e = z_e[:, 1:, ...]

        # quantize or kl
        z, d = self.bottleneck(z_e)

        if self.cfg.encoder_with_cls_token:
            d["z_cls"] = z_cls
        return z_e, z, None, d

    def decode(self, z: torch.Tensor):
        """
        Decodes the latent representation `z` using the decoder network.
        Args:
            z (torch.Tensor): The latent representation tensor to be decoded.
        Returns:
            torch.Tensor: The decoded output tensor.
        """

        h = self.decoder(z)
        return h

    def query(self, queries: torch.Tensor, latents: torch.Tensor):
        """
        Computes the logits by decoding the given queries and latent representations.
        Args:
            queries (torch.Tensor): A tensor containing the query points to be decoded.
            latents (torch.Tensor): A tensor containing the latent representations corresponding to the queries.
        Returns:
            torch.Tensor: A tensor containing the decoded logits for the given queries and latents.
        """

        logits = self.occupancy_decoder(queries, latents).squeeze(-1)
        return logits

    def forward(self, surface, queries, **kwargs):
        """
        Perform a forward pass through the autoencoder model.
        Args:
            surface (torch.Tensor): The input surface tensor to be encoded.
            queries (torch.Tensor): The query tensor used for generating logits.
            **kwargs: Additional keyword arguments.
        Returns:
            tuple: A tuple containing:
                - z (torch.Tensor): The latent representation of the input surface.
                - latents (torch.Tensor): The decoded output from the latent representation.
                - None: Placeholder for a potential future return value.
                - logits (torch.Tensor): The logits generated from the queries and latents.
                - d (torch.Tensor): Additional output from the encoding process.
        """

        _, z, _, d = self.encode(surface)

        latents = self.decode(z)
        logits = self.query(queries, latents)

        return z, latents, None, logits, d

    @torch.no_grad()
    def extract_geometry(
        self,
        latents: torch.FloatTensor,
        bounds: list[float] = [
            -1.05,
            -1.05,
            -1.05,
            1.05,
            1.05,
            1.05,
        ],
        resolution_base: float = 9.0,
        chunk_size: int = 2_000_000,
        use_warp: bool = False,
    ):
        """
        Extracts 3D geometry from latent representations using a dense grid sampling
        and marching cubes algorithm.
        Args:
            latents (torch.FloatTensor): A tensor of latent representations with shape
                (batch_size, latent_dim).
            bounds (list[float], optional): A list of six floats defining the bounding box
                for the 3D grid in the format [xmin, ymin, zmin, xmax, ymax, zmax].
                Defaults to [-1.05, -1.05, -1.05, 1.05, 1.05, 1.05].
            resolution_base (float, optional): The base resolution for the grid. Higher
                values result in finer grids. Defaults to 9.0.
            chunk_size (int, optional): The number of grid points to process in a single
                chunk. Defaults to 2,000,000.
            use_warp (bool, optional): Whether to use a GPU-accelerated marching cubes
                implementation. If False, falls back to a CPU implementation. Defaults to False.
        Returns:
            tuple:
                - mesh_v_f (list[tuple]): A list of tuples containing vertices and faces
                  for each batch element. Each tuple is of the form
                  (vertices, faces), where:
                    - vertices (np.ndarray): Array of vertex coordinates with shape
                      (num_vertices, 3).
                    - faces (np.ndarray): Array of face indices with shape
                      (num_faces, 3).
                  If geometry extraction fails for a batch element, the tuple will be
                  (None, None).
                - has_surface (np.ndarray): A boolean array indicating whether a surface
                  was successfully extracted for each batch element.
        Raises:
            Exception: Logs warnings or errors if geometry extraction fails for any
            batch element or if the marching cubes algorithm encounters issues.
        """
        bbox_min = np.array(bounds[0:3])
        bbox_max = np.array(bounds[3:6])
        bbox_size = bbox_max - bbox_min

        xyz_samples, grid_size, length = generate_dense_grid_points(
            bbox_min=bbox_min,
            bbox_max=bbox_max,
            resolution_base=resolution_base,
            indexing="ij",
        )
        xyz_samples = torch.FloatTensor(xyz_samples)
        batch_size = latents.shape[0]

        batch_logits = []

        progress_bar = tqdm(
            range(0, xyz_samples.shape[0], chunk_size),
            desc=f"extracting geometry",
            unit="chunk",
        )
        for start in progress_bar:
            queries = xyz_samples[start : start + chunk_size, :]

            num_queries = queries.shape[0]
            if start > 0 and num_queries < chunk_size:
                queries = F.pad(queries, [0, 0, 0, chunk_size - num_queries])
            batch_queries = queries.unsqueeze(0).expand(batch_size, -1, -1).to(latents)

            logits = self.query(batch_queries, latents)[:, :num_queries]
            batch_logits.append(logits)

        grid_logits = (
            torch.cat(batch_logits, dim=1)
            .detach()
            .view((batch_size, grid_size[0], grid_size[1], grid_size[2]))
            .float()
        )

        mesh_v_f = []
        has_surface = np.zeros((batch_size,), dtype=np.bool_)
        for i in range(batch_size):
            try:
                warp_success = False
                if use_warp:
                    try:
                        vertices, faces = marching_cubes_with_warp(
                            grid_logits[i],
                            level=0.0,
                            device=grid_logits.device,
                        )
                        warp_success = True
                    except Exception as e:
                        logging.warning(
                            f"Warning: error in marching cubes with warp: {e}"
                        )
                        warp_success = False  # Fall back to CPU version

                if not warp_success:
                    logging.warning(
                        "Warning: falling back to CPU version of marching cubes using skimage measure"
                    )
                    vertices, faces, _, _ = measure.marching_cubes(
                        grid_logits[i].cpu().numpy(), 0, method="lewiner"
                    )

                vertices = vertices / grid_size * bbox_size + bbox_min
                faces = faces[:, [2, 1, 0]]
                mesh_v_f.append(
                    (vertices.astype(np.float32), np.ascontiguousarray(faces))
                )
                has_surface[i] = True
            except Exception as e:
                logging.error(f"Error: error in extract_geometry: {e}")
                mesh_v_f.append((None, None))
                has_surface[i] = False

        return mesh_v_f, has_surface