Spaces:
Sleeping
Sleeping
| """Inference-only VQ-VAE using the released three-level SPAC checkpoint.""" | |
| import json | |
| from pathlib import Path | |
| from typing import Mapping | |
| import torch | |
| import torch.nn as nn | |
| from .sparse import SparseTensor | |
| from .nn.encoder import SparseSDFEncoder | |
| from .nn.decoder import SparseSDFDecoder | |
| class SparseVectorQuantizer(nn.Module): | |
| def __init__(self, num_embeddings: int, embedding_dim: int, use_ema_update: bool = True): | |
| super().__init__() | |
| self.embeddings = nn.Embedding(num_embeddings, embedding_dim) | |
| self.num_embeddings = num_embeddings | |
| self.embedding_dim = embedding_dim | |
| if use_ema_update: | |
| # Retain checkpoint buffers while omitting all training updates. | |
| self.register_buffer("ema_cluster_size", torch.zeros(num_embeddings)) | |
| self.register_buffer("ema_w", torch.zeros(num_embeddings, embedding_dim)) | |
| self.register_buffer("_ema_initialized", torch.tensor(False)) | |
| def forward(self, value: SparseTensor) -> SparseTensor: | |
| distances = torch.cdist(value.feats.to(self.embeddings.weight.dtype), self.embeddings.weight) | |
| indices = torch.argmin(distances, dim=1) | |
| return value.replace(indices.unsqueeze(-1).float()) | |
| class SparseSDFVQVAE(nn.Module): | |
| def __init__( | |
| self, resolution: int = 32, model_channels: int = 512, latent_channels: int = 16, | |
| num_blocks: int = 8, num_embeddings: int = 8192, vq_group_size: int = 1, | |
| num_heads: int = 8, num_head_channels: int = 64, mlp_ratio: float = 4.0, | |
| attn_mode: str = "swin", window_size: int = 8, pe_mode: str = "ape", | |
| use_ema_update: bool = True, legacy_three_level: bool = True, | |
| ): | |
| super().__init__() | |
| if not legacy_three_level or vq_group_size != 1: | |
| raise ValueError("This package contains only the released three-level, single-voxel SPAC VQ-VAE") | |
| self.resolution = resolution | |
| self.embed_dim = latent_channels | |
| self.vq_group_size = 1 | |
| self.vq_block_side = 1 | |
| arguments = dict( | |
| resolution=resolution, model_channels=model_channels, | |
| latent_channels=latent_channels, num_blocks=num_blocks, | |
| num_heads=num_heads, num_head_channels=num_head_channels, | |
| mlp_ratio=mlp_ratio, attn_mode=attn_mode, window_size=window_size, pe_mode=pe_mode, | |
| ) | |
| self.encoder = SparseSDFEncoder(**arguments) | |
| self.decoder = SparseSDFDecoder(**arguments) | |
| self.vq = SparseVectorQuantizer(num_embeddings, latent_channels, use_ema_update) | |
| def Encode(self, batch) -> SparseTensor: | |
| if isinstance(batch, Mapping): | |
| features = batch["sparse_sdf"] | |
| if features.ndim == 1: | |
| features = features.unsqueeze(-1) | |
| coordinates = torch.cat([batch["batch_idx"].unsqueeze(-1), batch["sparse_index"]], dim=-1).int() | |
| value = SparseTensor(features.float(), coordinates) | |
| factor = batch.get("factor") | |
| else: | |
| value = batch | |
| factor = None | |
| return self.vq(self.encoder(value, factor)) | |
| def Decode(self, indices: SparseTensor, band_prune=None, gt_prune=None) -> SparseTensor: | |
| ids = indices.feats.long().squeeze(-1) | |
| if not ids.numel() or ids.min() < 0 or ids.max() >= self.vq.num_embeddings: | |
| raise ValueError("A nonempty sequence of valid VQ codebook indices is required") | |
| latent = indices.replace(self.vq.embeddings(ids)) | |
| return self.decoder(latent, band_prune=band_prune, gt_prune=gt_prune) | |
| def load_vae(checkpoint, config, device) -> SparseSDFVQVAE: | |
| """Load every released encoder, decoder and codebook tensor strictly in FP32.""" | |
| configuration = json.loads(Path(config).read_text()) if not isinstance(config, Mapping) else dict(config) | |
| arguments = configuration.get("args", configuration) | |
| model = SparseSDFVQVAE(**arguments) | |
| state = torch.load(checkpoint, map_location="cpu", weights_only=True) | |
| if isinstance(state, Mapping) and isinstance(state.get("state_dict"), Mapping): | |
| state = state["state_dict"] | |
| state = {key.removeprefix("module."): value for key, value in state.items()} | |
| model.load_state_dict(state, strict=True) | |
| return model.float().eval().requires_grad_(False).to(device) | |