# chaos96/pointnsp architecture

> Architecture diagram of the GitHub repository chaos96/pointnsp, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/chaos96/pointnsp
- Repository: https://github.com/chaos96/pointnsp

## Overview

PointNSP generates 3D point clouds through a two-stage coarse-to-fine workflow: train a multi-scale VQVAE on ShapeNet point clouds, then train a causal transformer to predict its scale-wise code tokens, and combine the trained models for generation. The VQVAE encodes points with a point-voxel network, forms an FPS-based level-of-detail hierarchy, quantizes features, and decodes point clouds. Transformer training can tokenize online or consume pre-tokenized data. The README documents these capabilities; several implementation details, especially generation and fast-tokenization wiring, were not sampled, so those links follow documented entry points only.

Main entry points: [`model/vqvae_model.py`](https://github.com/chaos96/pointnsp/blob/main/model/vqvae_model.py) (Multi-Scale VQVAE), [`model/transformer.py`](https://github.com/chaos96/pointnsp/blob/main/model/transformer.py) (Causal Transformer), [`generate.py`](https://github.com/chaos96/pointnsp/blob/main/generate.py) (Point Cloud Generation), [`datasets/shapenet_data_pc.py`](https://github.com/chaos96/pointnsp/blob/main/datasets/shapenet_data_pc.py) (Point Loader)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Point Cloud Data"]
  node_shapenet[("ShapeNet Data")]
  node_pc_loader["Point Loader"]
  node_partnet_loader["PartNet Loader<br/>[partnet.py]"]
end

subgraph group_vqvae["VQVAE Stage"]
  node_vq_train["VQVAE Training<br/>[train_vqvae.py]"]
  node_vqvae_model["Multi-Scale VQVAE<br/>[vqvae_model.py]"]
  node_pvcnn["PVCNN Encoder<br/>[pvcnn.py]"]
  node_voxel["Voxel Branch<br/>[voxelization.py]"]
  node_fps["FPS LoD<br/>[fps.py]"]
  node_upsampling["Point Upsampling<br/>[upsampling.py]"]
  node_vq_state["VQ Codebook<br/>[vqvae_model.py]"]
end

subgraph group_transformer["Transformer Stage"]
  node_online_train["Online Transformer Training"]
  node_fast_train["Fast Transformer Training"]
  node_transformer_model["Causal Transformer<br/>[transformer.py]"]
  node_masking["Causal Masking<br/>[masking.py]"]
  node_position["Scale Position Encoding"]
end

subgraph group_generation["Generation"]
  node_generate["Point Cloud Generation<br/>[generate.py]"]
  node_visualize["Point Cloud Export<br/>[visualize.py]"]
  node_generated["Generated Samples"]
end

node_operator(("Operator"))

node_operator -->|"starts training"| node_vq_train
node_operator -->|"starts training"| node_online_train
node_operator -.->|"starts training"| node_fast_train
node_operator -->|"requests samples"| node_generate
node_shapenet -->|"loads points"| node_pc_loader
node_pc_loader -->|"batches points"| node_vq_train
node_partnet_loader -.->|"supplies batches"| node_vq_train
node_vq_train -->|"optimizes"| node_vqvae_model
node_vqvae_model -->|"encodes points"| node_pvcnn
node_pvcnn -->|"processes voxel branch"| node_voxel
node_vqvae_model -->|"builds LoD"| node_fps
node_vqvae_model -->|"quantizes features"| node_vq_state
node_vqvae_model -->|"refines points"| node_upsampling
node_shapenet -->|"loads points"| node_pc_loader
node_pc_loader -->|"batches points"| node_online_train
node_online_train -->|"tokenizes with"| node_vqvae_model
node_online_train -->|"trains tokens"| node_transformer_model
node_fast_train -.->|"trains tokens"| node_transformer_model
node_transformer_model -->|"applies causal mask"| node_masking
node_transformer_model -->|"uses scale positions"| node_position
node_generate -->|"decodes tokens"| node_vqvae_model
node_generate -->|"predicts tokens"| node_transformer_model
node_generate -->|"exports samples"| node_visualize
node_visualize -->|"writes files"| node_generated

click node_pc_loader "https://github.com/chaos96/pointnsp/blob/main/datasets/shapenet_data_pc.py"
click node_partnet_loader "https://github.com/chaos96/pointnsp/blob/main/datasets/partnet.py"
click node_vq_train "https://github.com/chaos96/pointnsp/blob/main/train_vqvae.py"
click node_vqvae_model "https://github.com/chaos96/pointnsp/blob/main/model/vqvae_model.py"
click node_pvcnn "https://github.com/chaos96/pointnsp/blob/main/model/pvcnn/pvcnn.py"
click node_voxel "https://github.com/chaos96/pointnsp/blob/main/model/pvcnn/voxelization.py"
click node_fps "https://github.com/chaos96/pointnsp/blob/main/model/fps.py"
click node_upsampling "https://github.com/chaos96/pointnsp/blob/main/model/upsampling.py"
click node_vq_state "https://github.com/chaos96/pointnsp/blob/main/model/vqvae_model.py"
click node_online_train "https://github.com/chaos96/pointnsp/blob/main/train_transformer.py"
click node_fast_train "https://github.com/chaos96/pointnsp/blob/main/train_transformer_fast.py"
click node_transformer_model "https://github.com/chaos96/pointnsp/blob/main/model/transformer.py"
click node_masking "https://github.com/chaos96/pointnsp/blob/main/model/masking.py"
click node_position "https://github.com/chaos96/pointnsp/blob/main/model/positional_encoding.py"
click node_generate "https://github.com/chaos96/pointnsp/blob/main/generate.py"
click node_visualize "https://github.com/chaos96/pointnsp/blob/main/utils/visualize.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_shapenet,node_pc_loader,node_partnet_loader toneBlue
class node_vq_train,node_vqvae_model,node_pvcnn,node_voxel,node_fps,node_upsampling,node_vq_state toneAmber
class node_online_train,node_fast_train,node_transformer_model,node_masking,node_position toneMint
class node_generate,node_visualize,node_generated toneRose
class node_operator toneIndigo
```

## Components

### Point Cloud Data

- **ShapeNet Data**
- [Point Loader](https://github.com/chaos96/pointnsp/blob/main/datasets/shapenet_data_pc.py): `datasets/shapenet_data_pc.py`
- [PartNet Loader](https://github.com/chaos96/pointnsp/blob/main/datasets/partnet.py): `datasets/partnet.py`

### VQVAE Stage

- [VQVAE Training](https://github.com/chaos96/pointnsp/blob/main/train_vqvae.py): `train_vqvae.py`
- [Multi-Scale VQVAE](https://github.com/chaos96/pointnsp/blob/main/model/vqvae_model.py): `model/vqvae_model.py`
- [PVCNN Encoder](https://github.com/chaos96/pointnsp/blob/main/model/pvcnn/pvcnn.py): `model/pvcnn/pvcnn.py`
- [Voxel Branch](https://github.com/chaos96/pointnsp/blob/main/model/pvcnn/voxelization.py): `model/pvcnn/voxelization.py`
- [FPS LoD](https://github.com/chaos96/pointnsp/blob/main/model/fps.py): `model/fps.py`
- [Point Upsampling](https://github.com/chaos96/pointnsp/blob/main/model/upsampling.py): `model/upsampling.py`
- [VQ Codebook](https://github.com/chaos96/pointnsp/blob/main/model/vqvae_model.py): `model/vqvae_model.py`

### Transformer Stage

- [Online Transformer Training](https://github.com/chaos96/pointnsp/blob/main/train_transformer.py): `train_transformer.py`
- [Fast Transformer Training](https://github.com/chaos96/pointnsp/blob/main/train_transformer_fast.py): `train_transformer_fast.py`
- [Causal Transformer](https://github.com/chaos96/pointnsp/blob/main/model/transformer.py): `model/transformer.py`
- [Causal Masking](https://github.com/chaos96/pointnsp/blob/main/model/masking.py): `model/masking.py`
- [Scale Position Encoding](https://github.com/chaos96/pointnsp/blob/main/model/positional_encoding.py): `model/positional_encoding.py`

### Generation

- [Point Cloud Generation](https://github.com/chaos96/pointnsp/blob/main/generate.py): `generate.py`
- [Point Cloud Export](https://github.com/chaos96/pointnsp/blob/main/utils/visualize.py): `utils/visualize.py`
- **Generated Samples**

### Other components

- **Operator**

## Connections

- Operator → VQVAE Training: starts training
- Operator → Online Transformer Training: starts training
- Operator → Fast Transformer Training: starts training
- Operator → Point Cloud Generation: requests samples
- ShapeNet Data → Point Loader: loads points
- Point Loader → VQVAE Training: batches points
- PartNet Loader → VQVAE Training: supplies batches
- VQVAE Training → Multi-Scale VQVAE: optimizes
- Multi-Scale VQVAE → PVCNN Encoder: encodes points
- PVCNN Encoder → Voxel Branch: processes voxel branch
- Multi-Scale VQVAE → FPS LoD: builds LoD
- Multi-Scale VQVAE → VQ Codebook: quantizes features
- Multi-Scale VQVAE → Point Upsampling: refines points
- ShapeNet Data → Point Loader: loads points
- Point Loader → Online Transformer Training: batches points
- Online Transformer Training → Multi-Scale VQVAE: tokenizes with
- Online Transformer Training → Causal Transformer: trains tokens
- Fast Transformer Training → Causal Transformer: trains tokens
- Causal Transformer → Causal Masking: applies causal mask
- Causal Transformer → Scale Position Encoding: uses scale positions
- Point Cloud Generation → Multi-Scale VQVAE: decodes tokens
- Point Cloud Generation → Causal Transformer: predicts tokens
- Point Cloud Generation → Point Cloud Export: exports samples
- Point Cloud Export → Generated Samples: writes files