model/vqvae.py

import torch
import torch.nn as nn
from copy import deepcopy  
import torch.nn.functional as F
from einops import rearrange
from torch_cluster import fps
from os import sys, path
import math
import numpy as np
from collections import (defaultdict, OrderedDict, deque)
import os

from SoftGroup.softgroup.ops import (voxelization_idx, voxelization)
from model.basic_vae import (Encoder, Decoder)
from model.NGC import get_NGC_structure
from model.common import (ResidualConv, CappedLayerNorm)
import pyvista as pv
import matplotlib.pyplot as plt

# Chamfer distance calculation


class VQVAE(nn.Module):
    def __init__(self, input_dim,
                       hidden_dim,
                       codebook_size, 
                       embedding_dim,
                       num_points, 
                       voxel_size, 
                       spconv_channels, 
                       blocks_num,
                       smooth_end_epoch,
                       beta,
                       lambda_chamfer,
                       lambda_vq_start,
                       lambda_vq_final,
                       lambda_construction,
                       lambda_usage,
                       warmup_steps,
                       ema_decay,
                       token_N1_stage,
                       token_N2_stage,
                       token_N3_stage,
                       decoder_layer,
                       buffer_stage_num):
        super(VQVAE, self).__init__()
        self.device = device = torch.device("cuda")
        self.blocks_num = blocks_num
        self.spconv_channels = spconv_channels
        self.encoder = Encoder(input_dim, hidden_dim, self.spconv_channels, self.blocks_num)
        self.hidden_dim = hidden_dim
        self.embedding_dim = embedding_dim
        self.codebook_size = codebook_size
        self.num_points = num_points
        self.voxel_size = voxel_size
        self.side_length_patchs = [(2 ** i) for i in range(self.voxel_size)]

        self.token_N3_stage = token_N3_stage
        
        self.token_N2_stage = token_N2_stage
        self.token_N1_stage = token_N1_stage
        self.token_N1_nums = 2 ** (token_N1_stage - 1)
        self.token_N2_nums = 2 ** (token_N2_stage - 1)
        self.token_N3_nums = 2 ** (token_N3_stage - 1)

        self.token_stage = self.token_N2_stage + buffer_stage_num
        self.token_stage_nums_list = [(2 ** i) for i in range(self.token_stage)]

        self.if_normal_compress = True
        self.num_freqs = 6
        self.pos_dim = 3 + 2 * 3 * self.num_freqs
        self.pos_linear = nn.Linear(self.pos_dim, self.spconv_channels)
        self.after_spconv = nn.Sequential(
            nn.Linear(spconv_channels, hidden_dim),
            nn.LayerNorm(hidden_dim),    
            nn.ReLU(),

            nn.Linear(hidden_dim, hidden_dim),
            nn.LayerNorm(hidden_dim),     
            nn.ReLU(),

            nn.Linear(hidden_dim, embedding_dim),
            CappedLayerNorm(embedding_dim)  
        )
        nn.init.xavier_uniform_(self.pos_linear.weight)  
        nn.init.zeros_(self.pos_linear.bias)

        self.embeddings = nn.Embedding(codebook_size, embedding_dim)
        self.embeddings.weight.data.uniform_(-1.0 / embedding_dim , 1.0 / embedding_dim)

        self.smooth_end_epoch = smooth_end_epoch
        self.phi = nn.ModuleList([
            ResidualConv(embedding_dim, embedding_dim)
            for _ in range(self.token_stage)  
        ])
        self.beta = beta
        self.decoder = Decoder(embedding_dim, hidden_dim, input_dim, decoder_layer)  # 6 layers
        self.if_freeze_spconv = False

        self.lambda_chamfer = lambda_chamfer
        self.warmup_steps = warmup_steps
        self.lambda_vq_start = lambda_vq_start
        self.lambda_vq_final = lambda_vq_final
        self.store_step = 200
        self.epoch_step = 20
        self.lambda_construction = lambda_construction

        self.alpha = nn.Parameter(torch.tensor(0.1))
        self.act = nn.GELU()
        self.log_data = torch.zeros(8, dtype=torch.float32)
        self.usage_weight = lambda_usage
        self.ema_decay = ema_decay

        self.vq_window = deque(maxlen=10)  # window_size e.g. 10
        self.triggers = 0
        self.patience_counter = 0
        self.safe_vq_pct_thr = 0.40
        self.lowest_vq_weight = 0.03
        
        self.init_ema()

    def init_ema(self):
        # Ensure embeddings won't be touched by optimizer
        self.embeddings.weight.requires_grad = False

        # register or create EMA buffers if not exist
        # Use register_buffer if inside nn.Module; otherwise set attributes
        if not hasattr(self, 'cluster_size'):
            self.register_buffer('cluster_size', torch.zeros(self.codebook_size, dtype=torch.float32))
        if not hasattr(self, 'embed_avg'):
            self.register_buffer('embed_avg', torch.zeros(self.codebook_size, self.embedding_dim, dtype=torch.float32))
        # move to device if user passed one
        self.cluster_size = self.cluster_size.to(self.device)
        self.embed_avg = self.embed_avg.to(self.device)
        self.ema_cluster_size = None
        self.ema_embed_sum = None
        self.ema_weights_list = [((i + 1) / self.token_stage) for i in range(self.token_stage)]

    def init_from_ckpt(self, path, ignore_keys, if_only_init_param):
        ignore_keys=list()
        sd = torch.load(path, map_location=self.device)
        if not if_only_init_param:  # return opt & curr_epoch
            curr_epoch = sd['epoch']
            params = sd['model']
            opt = sd['optimizer']
            missing, unexpected = self.load_state_dict(params, strict=False)
            print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
            if len(missing) > 0:
                print(f"Missing keys: {missing}")
            if len(unexpected) > 0:
                print(f"Unexpected keys: {unexpected}")
            return opt, curr_epoch
        else:   # start with 0 epoch
            params = sd['model']
            keys = list(params.keys())
            missing, unexpected = self.load_state_dict(params, strict=False)
            print(f"Restored from {path} with {len(missing)} missing and {len(unexpected)} unexpected keys")
            if len(missing) > 0:
                print(f"Missing keys: {missing}")
            if len(unexpected) > 0:
                print(f"Unexpected keys: {unexpected}")
            return None, None


    @torch.no_grad()
    def update_ema(self, decay, eps=1e-5):
        """
        z_e: (B, N, D) encoder outputs (float)
        codes: (B, N) long tensor of indices in [0, num_embeddings)
        This updates self.cluster_size and self.embed_avg, then writes to embeddings.
        """
        cluster_size_batch = self.ema_cluster_size
        embed_sum_batch = self.ema_embed_sum

        # EMA update (in-place)
        self.cluster_size.mul_(decay).add_(cluster_size_batch, alpha=1.0 - decay)   # (K,)
        self.embed_avg.mul_(decay).add_(embed_sum_batch, alpha=1.0 - decay)        # (K, D)

        # compute normalized cluster sizes to avoid divide-by-zero
        n = self.cluster_size.sum()
        # stabilized cluster_size for division
        cluster_size_normalized = ((self.cluster_size + eps) / (n + self.codebook_size * eps)) * n  # (K,)

        # new embeddings = embed_avg / cluster_size_normalized.unsqueeze(1)
        # guard against zeros/nans
        denom = cluster_size_normalized.unsqueeze(1)  # (K,1)
        # safe division
        new_emb = self.embed_avg / denom
        # replace any NaN/Inf entries with current embedding values (safe fallback)
        if torch.isnan(new_emb).any() or torch.isinf(new_emb).any():
            cur = self.embeddings.weight.data.to(device)
            nan_mask = (~torch.isfinite(new_emb))
            new_emb[nan_mask] = cur[nan_mask]

        # write back to embedding weights (in-place)
        self.embeddings.weight.data.copy_(new_emb)

    def point_cloud_normalize(self, x):  
        x_normalized = None
        if self.if_normal_compress:
            x_max_minus_min = x.max(dim=1, keepdim=True)[0] - x.min(dim=1, keepdim=True)[0] + 1e-8
            max_val = x_max_minus_min.max(dim=-1, keepdim=True)[0]
            x_norm = x_max_minus_min / (max_val + 1e-8)
            x_normalized = ((x - x.min(dim=1, keepdim=True)[0]) / x_max_minus_min) * x_norm
            pass
        else:
            x_normalized = (x - x.min(dim=1, keepdim=True)[0]) / (x.max(dim=1, keepdim=True)[0] - x.min(dim=1, keepdim=True)[0] + 1e-8)  # [0, 1]
        scaled_x_normalized = x_normalized * (self.side_length_patchs[-1] - 1) # [0, 63]
        x_normalized = 2. * x_normalized - 1.  # [-1, 1]
        return x_normalized, scaled_x_normalized

    def point_cloud_to_voxel(self, x):  # smn
        B, N, C = x.shape
        
        # Normalize point cloud to [0, resolution-1]
        x_normalized, scaled_x_normalized = self.point_cloud_normalize(x)  # x_normalized [-1, 1], scaled_x_normalized [0, resolution-1]
        x_normalized_long = scaled_x_normalized.long().cpu()

        batch_ids = torch.arange(B, dtype=torch.long).view(B, 1, 1).expand(B, N, 1)
        coords_long = torch.cat([batch_ids, x_normalized_long], dim=-1)
        coords_long = rearrange(coords_long, 'b n c -> (b n) c')

        batch_ids = batch_ids.to(x_normalized.device)
        coords_float = torch.cat([batch_ids, scaled_x_normalized], dim=-1)
        coords_float = rearrange(coords_float, 'b n c -> (b n) c')

        voxel_coords, v2p_map, p2v_map = voxelization_idx(coords_long, B)
        spatial_shape = [self.side_length_patchs[-1]] * 3
        p2v_map = p2v_map.to(x_normalized.device)

        voxel_feats = voxelization(coords_float, p2v_map)  # mode=4 avg
        batch_ids = voxel_feats[:, 0].long()
        voxel_feats = voxel_feats[:, 1:]
        voxel_feats_split = [voxel_feats[batch_ids == b] for b in range(B)]

        voxel_coords_split = []
        for i in range(B):
            voxel_coords_split.append(voxel_feats_split[i].long())

        return  voxel_feats_split, voxel_coords_split, spatial_shape, x_normalized

    def fourier_embed(self, x, num_freqs):
        x_norm = self.normalize_to_minus1_1(x)
        freqs = 2 ** torch.arange(num_freqs, device=x_norm.device) * torch.pi
        x_proj = x_norm[..., None] * freqs
        embed = torch.cat([x_norm, torch.sin(x_proj).flatten(-2), torch.cos(x_proj).flatten(-2)], dim=-1)
        embed = embed * (self.side_length_patchs[-1] / 2.0)
        return embed

    def flat_tokens(self, voxel_feats_total, voxel_batch_id_total, voxel_floatcoords_total):
        
        # first FPS to 8192 token nums
        L = self.token_N3_nums
        B = len(voxel_feats_total)
        sampled_feats_total = []
        sampled_floatcoords_total = []

        for batch_id in range(B):
            voxel_batch_id = voxel_batch_id_total[batch_id].squeeze(1).long()
            voxel_floatcoords = voxel_floatcoords_total[batch_id]
            voxel_feats = voxel_feats_total[batch_id].features
            points_num = voxel_batch_id.shape[0]
            if points_num < L:
                sample_idx = np.random.choice(points_num, L, replace=True)
            else:
                sample_idx = fps(voxel_floatcoords, voxel_batch_id, ratio=L/points_num)
            sampled_feats = voxel_feats[sample_idx]
            sampled_floatcoords = voxel_floatcoords[sample_idx]
            sampled_feats_total.append(sampled_feats)
            sampled_floatcoords_total.append(sampled_floatcoords)    

        sampled_feats_total = torch.stack(sampled_feats_total, dim=0)
        sampled_floatcoords_total = torch.stack(sampled_floatcoords_total, dim=0)
        floatcoords_total = sampled_floatcoords_total

        # pos embbedding & after-encode token mlp
        pos_embedding = self.act(self.pos_linear(self.fourier_embed(sampled_floatcoords_total, self.num_freqs)))
        token_feats = sampled_feats_total + self.alpha * pos_embedding   # [-1 1]
        token_feats = self.after_spconv(token_feats)

        # prepare for the construction loss
        norm_floatcoords_total = self.normalize_to_minus1_1(floatcoords_total)

        return token_feats, norm_floatcoords_total
    
    

    def encode(self, voxel_feats, voxel_coords, spatial_shape, epoch):

        # first voxelize & spconv encode pc
        voxel_feats_total, voxel_coords_total, voxel_batch_id_total = self.encoder(voxel_feats, voxel_coords, spatial_shape) # [0, resolution-1]


        # next flat voxel_feats into N3 tokens
        token_feats_N3, floatcoords_N3 = self.flat_tokens(voxel_feats_total, voxel_batch_id_total, voxel_feats) # [-1 1]

        # get structure gt & N3 -> N2
        structure_gt02_bin, structure_gt02_dec, token_feats_N2, floatcoords_N2 = get_NGC_structure(token_feats_N3, 
                                                                                                   floatcoords_N3, 
                                                                                                   self.token_N2_nums,
                                                                                                   self.token_N1_nums,
                                                                                                   self.token_stage)
        B, L, C = token_feats_N2.shape
        for b in range(B):
            structure_gt02_dec[b] =  structure_gt02_dec[b][:, :self.token_stage]
        structure_gt02_dec = torch.stack(structure_gt02_dec, dim=0).permute(0, 2, 1)

        # codebook_size = K, embedding_dim = D
        self.ema_cluster_size = torch.zeros(self.codebook_size, device=self.device)
        self.ema_embed_sum    = torch.zeros(self.codebook_size, C, device=self.device) 

        # start quantizing
        f_BLC = token_feats_N2
        f_no_grad = f_BLC.detach()
        f_rest = f_no_grad.clone()
        f_hat = torch.zeros_like(f_rest) 
        embedding = self.embeddings.weight
        min_encoding_stages_indices = []
        SN = len(self.token_stage_nums_list)
        mean_vq_loss: torch.Tensor = 0.0
        vq_loss_dict = defaultdict(float)

        for si in range(SN):
            # find the nearest embedding
            if si == SN - 1:  # last stage
               structure_map = torch.arange(L, device=self.device, dtype=torch.long).unsqueeze(0).repeat(B, 1)
            else:
               structure_map = structure_gt02_dec[:, si, :] # [B, L]
            h_BLC_list = []
            min_encoding_indices = []
            for b in range(B):
                structure_b = structure_map[b]      # (L,)
                f_rest_b = f_rest[b]                 # (L, C)
                uniq, inv = torch.unique(structure_b, return_inverse=True)
                M = uniq.numel()                     # unique 
                class_sums = torch.zeros(M, C, device=f_rest.device)
                class_sums.scatter_add_(0, inv.unsqueeze(-1).expand(-1, C), f_rest_b)    # (M, C)
                counts = torch.zeros(M, device=f_rest.device)
                counts.scatter_add_(0, inv, torch.ones_like(inv, dtype=counts.dtype))    # (M,)
                rest_NC_b = class_sums / counts.unsqueeze(-1)  # (M, C)
                d = (
                    rest_NC_b.pow(2).sum(1, keepdim=True)
                    + embedding.pow(2).sum(1)
                    - 2 * rest_NC_b @ embedding.T
                )  # (M, K)
                if epoch is not None and self.smooth_end_epoch != -1 and epoch < self.smooth_end_epoch:
                    logit = F.softmax(d.max(dim=-1, keepdim=True)[0] - d, dim=-1)
                    idx = torch.argmax(logit, dim=-1)
                    one_hot = F.one_hot(idx, self.codebook_size).type_as(logit)
                    one_hot = one_hot - logit.detach() + logit
                    h_NC_b = one_hot @ embedding
                else:
                    idx = torch.argmin(d, dim=1) # M
                    h_NC_b = embedding[idx]

                # store some ema information  # rest_NC: (Mi, C) # idx_N:   (Mi,)
                # count
                self.ema_cluster_size.scatter_add_(0, idx, torch.ones_like(idx, dtype=self.ema_cluster_size.dtype))
                # sum
                self.ema_embed_sum.scatter_add_(0, idx.unsqueeze(-1).expand(-1, C), rest_NC_b.detach() * self.ema_weights_list[si])
                
                min_encoding_indices.append(idx)
                h_LC_b = h_NC_b[inv]  
                h_BLC_list.append(h_LC_b)

            h_BLC = torch.stack(h_BLC_list, dim=0)  # (B, L, C)
            h_BLC = self.phi[int(si/SN)](h_BLC)

            f_hat = f_hat + h_BLC
            f_rest -= h_BLC
            mean_vq_loss_i = F.mse_loss(f_hat.detach(), f_BLC).mul_(self.beta) + F.mse_loss(f_hat, f_no_grad)
            
            vq_loss_dict[f'vq_loss_{si}'] = mean_vq_loss_i.item()
            mean_vq_loss += mean_vq_loss_i

            min_encoding_stages_indices.append(min_encoding_indices)

        mean_vq_loss *= 1. / SN
        f_hat = f_hat.detach() - f_no_grad + f_BLC

        return f_hat, f_BLC, mean_vq_loss, min_encoding_stages_indices, vq_loss_dict, floatcoords_N2

    def normalize_to_minus1_1(self, x):

        x_max_minus_min = x.max(dim=1, keepdim=True)[0] - x.min(dim=1, keepdim=True)[0] + 1e-8
        max_val = x_max_minus_min.max(dim=-1, keepdim=True)[0]
        x_norm = x_max_minus_min / (max_val + 1e-8)
        x_normalized = ((x - x.min(dim=1, keepdim=True)[0]) / x_max_minus_min) * x_norm
        x_normalized = 2. * x_normalized - 1.  # [-1, 1]
        return x_normalized

    def downsample(self, points):
        B, N, C = points.shape
        device = points.device
        out = []
        for b in range(B):
            pc = points[b]  # (N, 3)
            ratio = self.token_N2_nums / N

            idx = fps(pc, ratio=ratio)  # (target_n,)
            pc_out = pc[idx]  # (target_n, 3)
            out.append(pc_out)
        return torch.stack(out, dim=0)  # (B, target_n, 3)
    
    def chamfer_distance(self, x, y):
        xx = torch.sum(x**2, dim=2)
        yy = torch.sum(y**2, dim=2)
        zz = torch.matmul(x, y.transpose(2, 1))
        rx = xx.unsqueeze(2).expand(-1, -1, y.size(1))
        ry = yy.unsqueeze(1).expand(-1, x.size(1), -1)
        P = rx + ry - 2*zz
        return torch.mean(torch.min(P, dim=2)[0]) + torch.mean(torch.min(P, dim=1)[0])
    
    def get_vq_weight(self, step):
        vq_w_start = self.lambda_vq_start
        vq_w_end = self.lambda_vq_final
        if step >= self.warmup_steps:
            return vq_w_end
        # linear decay
        alpha = step / self.warmup_steps
        return vq_w_start *(1 - alpha) + vq_w_end * alpha
    
    def compute_usage_loss(self, eps=1e-12):
        """
        Use EMA cluster_size to compute global usage entropy loss
        """
        cluster_size = self.ema_cluster_size  # (K,)
        p = cluster_size / (cluster_size.sum() + eps)

        entropy = - (p * (p + eps).log()).sum()
        # maximize entropy → minimize -entropy
        usage_loss = -entropy
        return usage_loss
    

    def forward(self, x, epoch, global_step):

        # voxelize point cloud
        voxel_feats, voxel_coords, spatial_shape, x_normalized = self.point_cloud_to_voxel(x)
        x_normalized = self.downsample(x_normalized)

        # encode & quantize
         
        f_hat, f_gt, vq_loss, token_label_codebook_idxs, vq_loss_dict, floatcoords = self.encode(voxel_feats, voxel_coords, spatial_shape, epoch)
        vq_weight = self.get_vq_weight(epoch)
        vq_loss = vq_weight * vq_loss

        # decode & reconstruct 
        reconstructed = self.normalize_to_minus1_1(self.decoder(f_hat))
        
        # Calculate Chamfer distance loss
        chamfer_loss = self.chamfer_distance(x_normalized, reconstructed) # [-1 1]
        chamfer_loss = self.lambda_chamfer * chamfer_loss

        # Calculate Construction loss
        construction_loss = F.smooth_l1_loss(reconstructed, floatcoords, reduction='mean') # [-1 1]
        construction_loss = self.lambda_construction * construction_loss

        # Calculate usage loss
        usage_loss = self.compute_usage_loss()
        usage_loss = self.usage_weight * usage_loss
        
        # Total loss is the sum of Chamfer distance and VQ losses
        total_loss = construction_loss + vq_loss + chamfer_loss + usage_loss

        # ema update codebook
        self.update_ema(self.ema_decay)

        # stateful controller (init once)
        current_vq_pct = vq_loss.detach() / (vq_loss.detach() + construction_loss.detach() + chamfer_loss.detach())
        self.vq_stateful_controller(current_vq_pct, epoch)


        # print monitor information
        if global_step % 20 == 0:
            print("[",epoch,"/",global_step,"/",x.shape[0],"]")
            print("[Monitor] Token_feats(z_e) std: ", f_gt.std(dim=1).mean().item())
            print("[Monitor] Recon_feats(z_q) std: ", f_hat.std(dim=1).mean().item())
            active_codes = (self.ema_cluster_size > 1e-3).sum().item()
            print("[Monitor] Active embeddings:", active_codes, "/", self.codebook_size)
            print("[Monitor] Embedding weight std:",self.embeddings.weight.std().item())
            # ===== Perplexity computation =====
            probs = self.cluster_size / (self.cluster_size.sum() + 1e-10)  # (K,)
            perplexity = torch.exp(-torch.sum(probs * torch.log(probs + 1e-10)))

            print("[Monitor] Perplexity:", perplexity.item())
        
        return reconstructed, total_loss, construction_loss, vq_loss, chamfer_loss, usage_loss

    def vq_stateful_controller(self, current_vq_pct, epoch):
        if epoch < 100:
            return 
        self.vq_window.append(current_vq_pct)  # current_vq_pct = vq_loss / (recon + chamfer + vq)
        vq_pct_ma = sum(self.vq_window) / len(self.vq_window)

        if vq_pct_ma > self.safe_vq_pct_thr:
            self.patience_counter += 1
        else:
            self.patience_counter = 0

        if self.patience_counter >= 10 and self.triggers < 100:
            self.lambda_vq_final = max(self.lambda_vq_final * 0.99, self.lowest_vq_weight)
            self.triggers += 1
            self.patience_counter = 0