nodes/MrForExample-ComfyUI-3D-Pack
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Integrated Triplane Gaussian Transformers & Better 3DGS Import & Export

Browse Source
This commit is contained in:
MrForExample
2024-01-25 15:09:07 +01:00
parent 553416ac15
commit b95a0e475a
67 changed files with 9212 additions and 1008 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+10 -1
View File
@@ -155,4 +155,13 @@ cython_debug/
diff-gaussian-rasterization/
# Nvdiffrast
nvdiffrast/
nvdiffrast/
# Model files
*.ckpt
*.bin
*.pth
*.safetensors
# Ignore test outputs
outputs/
README.md
+53 -41
View File
@@ -14,6 +14,50 @@ ps: I'll show some generated result soon
- I tried install this package with ComfyUI embed python env first, but I can't find a way to build CUDA related libraries, e.g. diff-gaussian-rasterization, nvdiffrast, simple-knn.
- Torch version: 2.1.2+cu121/V.2.1.2+cu118
## Currently support:
- For use case please check [example workflows](./_Example_Workflows/)
- **Note:** you need to put [Example Inputs Files & Folders](_Example_Workflows/_Example_Inputs_Files/) under ComfyUI Root Directory\ComfyUI\input folder before you can run the example workflow
- Triplane Gaussian Transformers:
- Enable single image to 3D Gaussian, then you can use Deep Marching Tetrahedrons node to convert it to mesh
- Stack Orbit Camera Poses node to automatically generate all range of camera pose combinations
- You can use it to conditioning the [StableZero123 (You need to Download the checkpoint first)](https://comfyanonymous.github.io/ComfyUI_examples/3d/), with full range of camera poses in one prompt pass
- You can use it to generate the orbit camera poses and directly input to other 3D process node (e.g. GaussianSplatting and BakeTextureToMesh)
- Example usage:
<img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Clockwise_Camposes.png" width="256"/> <img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Counter_Clockwise_Camposes.png" width="256"/>
<img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Clockwise.gif" width="256"/> <img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Counter_Clockwise.gif" width="256"/>
- Load 3D file (.obj, .ply, .glb)
- Also have node that support save and load 3DGS .ply file
- 3D Gaussian Splatting, with:
- [Improved Differential Gaussian Rasterization](https://github.com/ashawkey/diff-gaussian-rasterization)
- Better Compactness-based Densification method from [Gsgen](https://gsgen3d.github.io/),
- Support initialize gaussians from given 3D mesh (Optional)
- Support mini-batch optimazation
- Multi-View images as inputs
- Export to .ply support
- Bake Multi-View images into UVTexture of given 3D mesh using [Nvdiffrast](https://github.com/NVlabs/nvdiffrast), supports:
- Export to .obj, .ply, .glb
- Deep Marching Tetrahedrons
- Allow convert 3DGS .ply file to 3D mesh <br>
*Note: I didn't spent time to turn the hyperprameters yet, the result will be improved in the future!*
## Roadmap:
- [x] Add DMTet algorithm to allow convertion from points cloud(Gaussian/.ply) to mesh (.obj, .ply, .glb)
- [x] Integrate [Triplane Meets Gaussian Splatting: Fast and Generalizable Single-View 3D Reconstruction with Transformers](https://github.com/VAST-AI-Research/TriplaneGaussian)
- [ ] Add interactive 3D UI inside ComfuUI to visulaize training and generated results for 3D representations
- [ ] Improve DMTet result & add support to training it with images(RGB, Alpha, Normal Map)
- [ ] Add a general SDS/VSD Optimization algorithm to allow training 3D representations with diffusion model, *The real fun begins here* ;)
- [ ] Add a few best Nerf algorithms (No idea yet, [instant-ngp](https://github.com/NVlabs/instant-ngp) maybe?)
## Install:
Assume you have already downloaded [ComfyUI](https://github.com/comfyanonymous/ComfyUI)
@@ -56,53 +100,21 @@ install.bat
git clone --recursive https://github.com/NVlabs/nvdiffrast/`
pip install ./nvdiffrast
# Install pointnet2_ops
cd tgs/models/snowflake/pointnet2_ops_lib && python setup.py install && cd ../../../../
# Install pytorch_scatter
pip install git+https://github.com/rusty1s/pytorch_scatter.git
# Install pytorch3d
pip install git+https://github.com/facebookresearch/pytorch3d.git@stable
```
## Run:
Copy the files inside folder [__New_ComfyUI_Bats](./_New_ComfyUI_Bats/) to your ComfyUI root directory, and double click run_nvidia_gpu_miniconda.bat to start ComfyUI!
- Alternatively you can just activate the Conda env: python_miniconda_env\ComfyUI, and go to your ComfyUI root directory then run command `python ./ComfyUI/main.py`
## Currently support:
- For use case please check [example workflows](./_Example_Workflows/)
- **Note:** you need to put [Example Inputs Files & Folders](_Example_Workflows/_Example_Inputs_Files/) under ComfyUI Root Directory\ComfyUI\input folder before you can run the example workflow
- Stack Orbit Camera Poses node to automatically generate all range of camera pose combinations
- You can use it to conditioning the [StableZero123 (You need to Download the checkpoint first)](https://comfyanonymous.github.io/ComfyUI_examples/3d/), with full range of camera poses in one prompt pass
- You can use it to generate the orbit camera poses and directly input to other 3D process node (e.g. GaussianSplatting and BakeTextureToMesh)
- Example usage:
<img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Clockwise_Camposes.png" width="256"/> <img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Counter_Clockwise_Camposes.png" width="256"/>
<img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Clockwise.gif" width="256"/> <img src="_Example_Workflows/_Example_Outputs/Cammy_Cam_Rotate_Counter_Clockwise.gif" width="256"/>
- Load 3D file (.obj, .ply, .glb)
- Also have node that support save and load 3DGS .ply file
- 3D Gaussian Splatting, with:
- [Improved Differential Gaussian Rasterization](https://github.com/ashawkey/diff-gaussian-rasterization)
- Better Compactness-based Densification method from [Gsgen](https://gsgen3d.github.io/),
- Support initialize gaussians from given 3D mesh (Optional)
- Support mini-batch optimazation
- Multi-View images as inputs
- Export to .ply support
- Bake Multi-View images into UVTexture of given 3D mesh using [Nvdiffrast](https://github.com/NVlabs/nvdiffrast), supports:
- Export to .obj, .ply, .glb
- Deep Marching Tetrahedrons
- Allow convert 3DGS .ply file to 3D mesh <br>
*Note: I didn't spent time to turn the hyperprameters yet, the result will be improved in the future!*
## Roadmap:
- [x] Add DMTet algorithm to allow convertion from points cloud(Gaussian/.ply) to mesh (.obj, .ply, .glb)
- [ ] Integrate [Triplane Meets Gaussian Splatting: Fast and Generalizable Single-View 3D Reconstruction with Transformers](https://github.com/VAST-AI-Research/TriplaneGaussian)
- [ ] Add interactive 3D UI inside ComfuUI to visulaize training and generated results for 3D representations
- [ ] Improve DMTet result & add support to training it with images(RGB, Alpha, Normal Map)
- [ ] Add a general SDS/VSD Optimization algorithm to allow training 3D representations with diffusion model, *The real fun begins here* ;)
- [ ] Add a few best Nerf algorithms (No idea yet, [instant-ngp](https://github.com/NVlabs/instant-ngp) maybe?)
## Tips
* The world & camera coordinate system is the same as OpenGL:
```
_Example_Workflows/3DGS_Multi-View_Renders_to_ply.json
+238 -230
View File
@@ -1,6 +1,6 @@
{
"last_node_id": 19,
"last_link_id": 22,
"last_node_id": 21,
"last_link_id": 26,
"nodes": [
{
"id": 10,
@@ -63,191 +63,6 @@
"Node name for S&R": "MaskToImage"
}
},
{
"id": 17,
"type": "[ComfyUI-3D] Save 3DGS",
"pos": [
2122,
54
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "raw_3DGS",
"type": "GS_RAW",
"link": 21
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Save 3DGS"
},
"widgets_values": [
"MeshTest\\3DGS_01.ply"
]
},
{
"id": 8,
"type": "InvertMask",
"pos": [
746,
-14
],
"size": {
"0": 210,
"1": 26
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 2
}
],
"outputs": [
{
"name": "MASK",
"type": "MASK",
"links": [
8,
18
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "InvertMask"
}
},
{
"id": 4,
"type": "[ComfyUI-3D] Generate Orbit Camera Poses",
"pos": [
830,
132
],
"size": {
"0": 400,
"1": 200
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 1
}
],
"outputs": [
{
"name": "orbit_camposes",
"type": "ORBIT_CAMPOSES",
"links": [
19
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Generate Orbit Camera Poses"
},
"widgets_values": [
"([0:120], 1.75, 0, 0, 360)"
]
},
{
"id": 18,
"type": "[ComfyUI-3D] Gaussian Splatting",
"pos": [
1486,
53
],
"size": {
"0": 468.5999755859375,
"1": 718
},
"flags": {
"collapsed": false
},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 17
},
{
"name": "reference_masks",
"type": "MASK",
"link": 18
},
{
"name": "reference_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 19
},
{
"name": "mesh_to_initialize_gaussian",
"type": "MESH",
"link": null
}
],
"outputs": [
{
"name": "raw_3DGS",
"type": "GS_RAW",
"links": [
21
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Gaussian Splatting"
},
"widgets_values": [
49.1,
1000,
1,
10000,
0.2,
3,
0,
0,
0.5,
0.01,
0.05,
0.005,
0.005,
0.001,
0.00002,
0.02,
500,
5000,
3,
0.01,
100,
100000,
100,
700,
0.01,
0
]
},
{
"id": 13,
"type": "[ComfyUI-3D] Load 3D Mesh",
@@ -283,6 +98,31 @@
false
]
},
{
"id": 19,
"type": "PreviewImage",
"pos": [
420,
-320
],
"size": [
210,
246
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 22
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 7,
"type": "VHS_LoadImages",
@@ -303,8 +143,8 @@
"type": "IMAGE",
"links": [
1,
17,
22
22,
23
],
"shape": 3,
"slot_index": 0
@@ -349,29 +189,197 @@
}
},
{
"id": 19,
"type": "PreviewImage",
"id": 8,
"type": "InvertMask",
"pos": [
420,
-320
746,
-14
],
"size": {
"0": 210,
"1": 26
},
"flags": {},
"order": 3,
"order": 4,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 22
"name": "mask",
"type": "MASK",
"link": 2
}
],
"outputs": [
{
"name": "MASK",
"type": "MASK",
"links": [
8,
24
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "PreviewImage"
"Node name for S&R": "InvertMask"
}
},
{
"id": 4,
"type": "[ComfyUI-3D] Generate Orbit Camera Poses",
"pos": [
830,
132
],
"size": {
"0": 400,
"1": 200
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 1
}
],
"outputs": [
{
"name": "orbit_camposes",
"type": "ORBIT_CAMPOSES",
"links": [
25
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Generate Orbit Camera Poses"
},
"widgets_values": [
"([0:120], 1.75, 0, 0, 360)"
]
},
{
"id": 20,
"type": "[ComfyUI-3D] Gaussian Splatting",
"pos": [
1381,
54
],
"size": {
"0": 443.4000244140625,
"1": 758
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 23
},
{
"name": "reference_masks",
"type": "MASK",
"link": 24
},
{
"name": "reference_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 25
},
{
"name": "points_cloud",
"type": "POINTCLOUD",
"link": null
},
{
"name": "ply_to_initialize_gaussian",
"type": "GS_PLY",
"link": null
},
{
"name": "mesh_to_initialize_gaussian",
"type": "MESH",
"link": null
}
],
"outputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"links": [
26
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Gaussian Splatting"
},
"widgets_values": [
49.1,
1000,
3,
10000,
0.2,
3,
0,
0,
0.5,
0.01,
0.05,
0.005,
0.005,
0.001,
0.00002,
0.02,
500,
5000,
3,
0.01,
100,
100000,
100,
700,
0.01,
0
]
},
{
"id": 21,
"type": "[ComfyUI-3D] Save 3DGS",
"pos": [
1994,
55
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"link": 26
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Save 3DGS"
},
"widgets_values": [
"3DGS_[time(%Y-%m-%d)].ply"
]
}
],
"links": [
@@ -407,38 +415,6 @@
0,
"MASK"
],
[
17,
7,
0,
18,
0,
"IMAGE"
],
[
18,
8,
0,
18,
1,
"MASK"
],
[
19,
4,
0,
18,
2,
"ORBIT_CAMPOSES"
],
[
21,
18,
0,
17,
0,
"GS_RAW"
],
[
22,
7,
@@ -446,6 +422,38 @@
19,
0,
"IMAGE"
],
[
23,
7,
0,
20,
0,
"IMAGE"
],
[
24,
8,
0,
20,
1,
"MASK"
],
[
25,
4,
0,
20,
2,
"ORBIT_CAMPOSES"
],
[
26,
20,
0,
21,
0,
"GS_PLY"
]
],
"groups": [],
_Example_Workflows/3DGS_Multi-View_Renders_to_ply_with_Stack_Orbit_Cam_Poses_node.json
+213 -205
View File
@@ -1,6 +1,6 @@
{
"last_node_id": 22,
"last_link_id": 25,
"last_node_id": 24,
"last_link_id": 29,
"nodes": [
{
"id": 10,
@@ -64,150 +64,29 @@
}
},
{
"id": 17,
"type": "[ComfyUI-3D] Save 3DGS",
"id": 19,
"type": "PreviewImage",
"pos": [
2122,
54
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "raw_3DGS",
"type": "GS_RAW",
"link": 21
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Save 3DGS"
},
"widgets_values": [
"MeshTest\\3DGS_01.ply"
]
},
{
"id": 8,
"type": "InvertMask",
"pos": [
746,
-14
420,
-320
],
"size": {
"0": 210,
"1": 26
"1": 246
},
"flags": {},
"order": 4,
"order": 3,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 2
}
],
"outputs": [
{
"name": "MASK",
"type": "MASK",
"links": [
8,
18
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "InvertMask"
}
},
{
"id": 18,
"type": "[ComfyUI-3D] Gaussian Splatting",
"pos": [
1486,
53
],
"size": {
"0": 468.5999755859375,
"1": 718
},
"flags": {
"collapsed": false
},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"name": "images",
"type": "IMAGE",
"link": 17
},
{
"name": "reference_masks",
"type": "MASK",
"link": 18
},
{
"name": "reference_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 25
},
{
"name": "mesh_to_initialize_gaussian",
"type": "MESH",
"link": null
}
],
"outputs": [
{
"name": "raw_3DGS",
"type": "GS_RAW",
"links": [
21
],
"shape": 3,
"slot_index": 0
"link": 22
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Gaussian Splatting"
},
"widgets_values": [
49.1,
1000,
1,
10000,
0.2,
3,
0,
0,
0.5,
0.01,
0.05,
0.005,
0.005,
0.001,
0.00002,
0.02,
500,
5000,
3,
0.01,
100,
100000,
100,
700,
0.01,
0
]
"Node name for S&R": "PreviewImage"
}
},
{
"id": 7,
@@ -228,8 +107,8 @@
"name": "IMAGE",
"type": "IMAGE",
"links": [
17,
22
22,
26
],
"shape": 3,
"slot_index": 0
@@ -274,85 +153,62 @@
}
},
{
"id": 19,
"type": "PreviewImage",
"id": 8,
"type": "InvertMask",
"pos": [
420,
-320
746,
-14
],
"size": {
"0": 210,
"1": 246
"1": 26
},
"flags": {},
"order": 3,
"order": 4,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 22
"name": "mask",
"type": "MASK",
"link": 2
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 13,
"type": "[ComfyUI-3D] Load 3D Mesh",
"pos": [
1100,
315
],
"size": {
"0": 315,
"1": 154
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "mesh",
"type": "MESH",
"links": [],
"name": "MASK",
"type": "MASK",
"links": [
8,
27
],
"shape": 3,
"slot_index": 0
}
],
"title": "Load 3D Mesh [Optional]",
"properties": {
"Node name for S&R": "[ComfyUI-3D] Load 3D Mesh"
},
"widgets_values": [
"Meshes\\Cammy_Uniform.obj",
false,
true,
false,
false
]
"Node name for S&R": "InvertMask"
}
},
{
"id": 22,
"type": "[ComfyUI-3D] Stack Orbit Camera Poses",
"pos": [
726,
160
696,
174
],
"size": {
"0": 315,
"1": 586
},
"flags": {},
"order": 2,
"order": 1,
"mode": 0,
"outputs": [
{
"name": "orbit_camposes",
"type": "ORBIT_CAMPOSES",
"links": [
25
28
],
"shape": 3,
"slot_index": 0
@@ -417,6 +273,158 @@
0,
0.1
]
},
{
"id": 23,
"type": "[ComfyUI-3D] Save 3DGS",
"pos": [
2178,
54
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"link": 29
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Save 3DGS"
},
"widgets_values": [
"MeshTest\\3DGS_01.ply"
]
},
{
"id": 13,
"type": "[ComfyUI-3D] Load 3D Mesh",
"pos": [
1109,
327
],
"size": {
"0": 315,
"1": 154
},
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{
"name": "mesh",
"type": "MESH",
"links": [],
"shape": 3,
"slot_index": 0
}
],
"title": "Load 3D Mesh [Optional]",
"properties": {
"Node name for S&R": "[ComfyUI-3D] Load 3D Mesh"
},
"widgets_values": [
"Meshes\\Cammy_Uniform.obj",
false,
true,
false,
false
]
},
{
"id": 24,
"type": "[ComfyUI-3D] Gaussian Splatting",
"pos": [
1516,
54
],
"size": {
"0": 443.4000244140625,
"1": 758
},
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 26
},
{
"name": "reference_masks",
"type": "MASK",
"link": 27
},
{
"name": "reference_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 28
},
{
"name": "points_cloud",
"type": "POINTCLOUD",
"link": null
},
{
"name": "ply_to_initialize_gaussian",
"type": "GS_PLY",
"link": null
},
{
"name": "mesh_to_initialize_gaussian",
"type": "MESH",
"link": null
}
],
"outputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"links": [
29
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Gaussian Splatting"
},
"widgets_values": [
49.1,
1000,
3,
10000,
0.2,
3,
0,
0,
0.5,
0.01,
0.05,
0.005,
0.005,
0.001,
0.00002,
0.02,
500,
5000,
3,
0.01,
100,
100000,
100,
700,
0.01,
0
]
}
],
"links": [
@@ -444,30 +452,6 @@
0,
"MASK"
],
[
17,
7,
0,
18,
0,
"IMAGE"
],
[
18,
8,
0,
18,
1,
"MASK"
],
[
21,
18,
0,
17,
0,
"GS_RAW"
],
[
22,
7,
@@ -477,12 +461,36 @@
"IMAGE"
],
[
25,
26,
7,
0,
24,
0,
"IMAGE"
],
[
27,
8,
0,
24,
1,
"MASK"
],
[
28,
22,
0,
18,
24,
2,
"ORBIT_CAMPOSES"
],
[
29,
24,
0,
23,
0,
"GS_PLY"
]
],
"groups": [],
_Example_Workflows/Images_to_3DGS_to_DMTetMesh.json
+325 -287
View File
@@ -1,6 +1,6 @@
{
"last_node_id": 22,
"last_link_id": 26,
"last_node_id": 27,
"last_link_id": 36,
"nodes": [
{
"id": 10,
@@ -14,7 +14,7 @@
"1": 246
},
"flags": {},
"order": 7,
"order": 8,
"mode": 0,
"inputs": [
{
@@ -39,7 +39,7 @@
"1": 26
},
"flags": {},
"order": 5,
"order": 6,
"mode": 0,
"inputs": [
{
@@ -64,79 +64,63 @@
}
},
{
"id": 8,
"type": "InvertMask",
"id": 19,
"type": "PreviewImage",
"pos": [
746,
-14
420,
-320
],
"size": {
"0": 210,
"1": 26
"1": 246
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 2
}
],
"outputs": [
{
"name": "MASK",
"type": "MASK",
"links": [
8,
18
],
"shape": 3,
"slot_index": 0
"name": "images",
"type": "IMAGE",
"link": 22
}
],
"properties": {
"Node name for S&R": "InvertMask"
"Node name for S&R": "PreviewImage"
}
},
{
"id": 4,
"type": "[ComfyUI-3D] Generate Orbit Camera Poses",
"id": 13,
"type": "[ComfyUI-3D] Load 3D Mesh",
"pos": [
830,
132
1053,
415
],
"size": {
"0": 400,
"1": 200
"0": 315,
"1": 154
},
"flags": {},
"order": 2,
"order": 0,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 1
}
],
"outputs": [
{
"name": "orbit_camposes",
"type": "ORBIT_CAMPOSES",
"links": [
19
],
"name": "mesh",
"type": "MESH",
"links": [],
"shape": 3,
"slot_index": 0
}
],
"title": "Load 3D Mesh [Optional]",
"properties": {
"Node name for S&R": "[ComfyUI-3D] Generate Orbit Camera Poses"
"Node name for S&R": "[ComfyUI-3D] Load 3D Mesh"
},
"widgets_values": [
"([0:120], 1.75, 0, 0, 360)"
"Meshes\\Cammy_Uniform.obj",
false,
true,
false,
false
]
},
{
@@ -151,7 +135,7 @@
194
],
"flags": {},
"order": 0,
"order": 1,
"mode": 0,
"outputs": [
{
@@ -159,8 +143,8 @@
"type": "IMAGE",
"links": [
1,
17,
22
22,
27
],
"shape": 3,
"slot_index": 0
@@ -205,258 +189,128 @@
}
},
{
"id": 19,
"type": "PreviewImage",
"id": 8,
"type": "InvertMask",
"pos": [
420,
-320
746,
-14
],
"size": {
"0": 210,
"1": 246
"1": 26
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 2
}
],
"outputs": [
{
"name": "MASK",
"type": "MASK",
"links": [
8,
28
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "InvertMask"
}
},
{
"id": 4,
"type": "[ComfyUI-3D] Generate Orbit Camera Poses",
"pos": [
830,
132
],
"size": {
"0": 400,
"1": 200
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 22
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 18,
"type": "[ComfyUI-3D] Gaussian Splatting",
"pos": [
1486,
53
],
"size": {
"0": 468.5999755859375,
"1": 718
},
"flags": {
"collapsed": false
},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 17
},
{
"name": "reference_masks",
"type": "MASK",
"link": 18
},
{
"name": "reference_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 19
},
{
"name": "mesh_to_initialize_gaussian",
"type": "MESH",
"link": null
"link": 1
}
],
"outputs": [
{
"name": "raw_3DGS",
"type": "GS_RAW",
"name": "orbit_camposes",
"type": "ORBIT_CAMPOSES",
"links": [
21,
24
29
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Gaussian Splatting"
"Node name for S&R": "[ComfyUI-3D] Generate Orbit Camera Poses"
},
"widgets_values": [
49.1,
1000,
1,
10000,
0.2,
3,
0,
0,
0.5,
0.01,
0.05,
0.005,
0.005,
0.001,
0.00002,
0.02,
500,
5000,
3,
0.01,
100,
100000,
100,
700,
0.01,
0
"([0:120], 1.75, 0, 0, 360)"
]
},
{
"id": 17,
"id": 24,
"type": "[ComfyUI-3D] Save 3DGS",
"pos": [
2243,
52
2122,
-131
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 8,
"order": 9,
"mode": 0,
"inputs": [
{
"name": "raw_3DGS",
"type": "GS_RAW",
"link": 21
"name": "gs_ply",
"type": "GS_PLY",
"link": 30
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Save 3DGS"
},
"widgets_values": [
"MeshTest\\3DGS_01.ply"
]
},
{
"id": 20,
"type": "[ComfyUI-3D] Convert 3DGS To Pointcloud",
"pos": [
2242,
251
],
"size": {
"0": 267,
"1": 26
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "raw_3DGS",
"type": "GS_RAW",
"link": 24
}
],
"outputs": [
{
"name": "points_cloud",
"type": "POINTCLOUD",
"links": [
25
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Convert 3DGS To Pointcloud"
}
},
{
"id": 13,
"type": "[ComfyUI-3D] Load 3D Mesh",
"pos": [
1053,
415
],
"size": {
"0": 315,
"1": 154
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "mesh",
"type": "MESH",
"links": [],
"shape": 3,
"slot_index": 0
}
],
"title": "Load 3D Mesh [Optional]",
"properties": {
"Node name for S&R": "[ComfyUI-3D] Load 3D Mesh"
},
"widgets_values": [
"Meshes\\Cammy_Uniform.obj",
false,
true,
false,
false
]
},
{
"id": 22,
"type": "[ComfyUI-3D] Save 3D Mesh",
"pos": [
3261,
55
],
"size": {
"0": 353.38916015625,
"1": 61.879417419433594
},
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESH",
"link": 26
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Save 3D Mesh"
},
"widgets_values": [
"MeshTest\\Cammy_Uniform_DMTet_01.obj"
"MeshTest/3DGS_01.ply"
]
},
{
"id": 21,
"type": "[ComfyUI-3D] Deep Marching Tetrahedrons",
"pos": [
2710,
56
2963,
55
],
"size": {
"0": 443.4000244140625,
"1": 314
},
"flags": {},
"order": 10,
"order": 11,
"mode": 0,
"inputs": [
{
"name": "reference_points_cloud",
"type": "POINTCLOUD",
"link": 25
"link": 35
},
{
"name": "reference_images",
@@ -495,6 +349,190 @@
128,
5
]
},
{
"id": 22,
"type": "[ComfyUI-3D] Save 3D Mesh",
"pos": [
3534,
56
],
"size": {
"0": 353.38916015625,
"1": 61.879417419433594
},
"flags": {},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESH",
"link": 26
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Save 3D Mesh"
},
"widgets_values": [
"MeshTest\\Cammy_Uniform_DMTet_01.obj"
]
},
{
"id": 26,
"type": "[ComfyUI-3D] Load 3DGS",
"pos": [
2107,
194
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 2,
"mode": 0,
"outputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"links": [],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Load 3DGS"
},
"widgets_values": [
"Meshes/3DGS_01.ply"
]
},
{
"id": 23,
"type": "[ComfyUI-3D] Gaussian Splatting",
"pos": [
1501,
53
],
"size": {
"0": 443.4000244140625,
"1": 758
},
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 27
},
{
"name": "reference_masks",
"type": "MASK",
"link": 28
},
{
"name": "reference_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 29
},
{
"name": "points_cloud",
"type": "POINTCLOUD",
"link": null
},
{
"name": "ply_to_initialize_gaussian",
"type": "GS_PLY",
"link": null
},
{
"name": "mesh_to_initialize_gaussian",
"type": "MESH",
"link": null
}
],
"outputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"links": [
30,
36
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Gaussian Splatting"
},
"widgets_values": [
49.1,
1000,
3,
10000,
0.2,
3,
0,
0,
0.5,
0.01,
0.05,
0.005,
0.005,
0.001,
0.00002,
0.02,
500,
5000,
3,
0.01,
100,
100000,
100,
700,
0.01,
0
]
},
{
"id": 27,
"type": "[ComfyUI-3D] Convert 3DGS To Pointcloud",
"pos": [
2555,
53
],
"size": {
"0": 241.79998779296875,
"1": 26
},
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"link": 36
}
],
"outputs": [
{
"name": "points_cloud",
"type": "POINTCLOUD",
"links": [
35
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Convert 3DGS To Pointcloud"
}
}
],
"links": [
@@ -530,38 +568,6 @@
0,
"MASK"
],
[
17,
7,
0,
18,
0,
"IMAGE"
],
[
18,
8,
0,
18,
1,
"MASK"
],
[
19,
4,
0,
18,
2,
"ORBIT_CAMPOSES"
],
[
21,
18,
0,
17,
0,
"GS_RAW"
],
[
22,
7,
@@ -570,22 +576,6 @@
0,
"IMAGE"
],
[
24,
18,
0,
20,
0,
"GS_RAW"
],
[
25,
20,
0,
21,
0,
"POINTCLOUD"
],
[
26,
21,
@@ -593,6 +583,54 @@
22,
0,
"MESH"
],
[
27,
7,
0,
23,
0,
"IMAGE"
],
[
28,
8,
0,
23,
1,
"MASK"
],
[
29,
4,
0,
23,
2,
"ORBIT_CAMPOSES"
],
[
30,
23,
0,
24,
0,
"GS_PLY"
],
[
35,
27,
0,
21,
0,
"POINTCLOUD"
],
[
36,
23,
0,
27,
0,
"GS_PLY"
]
],
"groups": [],
_Example_Workflows/Load_3DGS_to_DMTetMesh.json
+122 -122
View File
@@ -1,109 +1,13 @@
{
"last_node_id": 4,
"last_link_id": 3,
"last_node_id": 28,
"last_link_id": 37,
"nodes": [
{
"id": 2,
"type": "[ComfyUI-3D] Convert 3DGS To Pointcloud",
"pos": [
561,
484
],
"size": {
"0": 267,
"1": 26
},
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"name": "raw_3DGS",
"type": "GS_RAW",
"link": 1
}
],
"outputs": [
{
"name": "points_cloud",
"type": "POINTCLOUD",
"links": [
2
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Convert 3DGS To Pointcloud"
}
},
{
"id": 1,
"type": "[ComfyUI-3D] Load 3DGS",
"pos": [
166,
484
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "raw_3DGS",
"type": "GS_RAW",
"links": [
1
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Load 3DGS"
},
"widgets_values": [
"Meshes\\Generated\\Cammy_Uniform_3DGS_01.ply"
]
},
{
"id": 4,
"type": "[ComfyUI-3D] Save 3D Mesh",
"pos": [
1446,
485
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESH",
"link": 3
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Save 3D Mesh"
},
"widgets_values": [
"MeshTest\\Cammy_Uniform_DMTet_01.obj"
]
},
{
"id": 3,
"id": 21,
"type": "[ComfyUI-3D] Deep Marching Tetrahedrons",
"pos": [
913,
484
1450,
420
],
"size": {
"0": 443.4000244140625,
@@ -116,7 +20,7 @@
{
"name": "reference_points_cloud",
"type": "POINTCLOUD",
"link": 2
"link": 37
},
{
"name": "reference_images",
@@ -134,7 +38,7 @@
"name": "trained_mesh",
"type": "MESH",
"links": [
3
26
],
"shape": 3,
"slot_index": 0
@@ -155,32 +59,128 @@
128,
5
]
},
{
"id": 22,
"type": "[ComfyUI-3D] Save 3D Mesh",
"pos": [
2020,
420
],
"size": {
"0": 353.38916015625,
"1": 61.879417419433594
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESH",
"link": 26
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Save 3D Mesh"
},
"widgets_values": [
"MeshTest\\picacho.obj"
]
},
{
"id": 28,
"type": "[ComfyUI-3D] Convert 3DGS To Pointcloud",
"pos": [
1035,
422
],
"size": {
"0": 241.79998779296875,
"1": 26
},
"flags": {},
"order": 1,
"mode": 0,
"inputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"link": 36
}
],
"outputs": [
{
"name": "points_cloud",
"type": "POINTCLOUD",
"links": [
37
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Convert 3DGS To Pointcloud"
}
},
{
"id": 26,
"type": "[ComfyUI-3D] Load 3DGS",
"pos": [
592,
423
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"links": [
36
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Load 3DGS"
},
"widgets_values": [
"Meshes/3DGS_TGS.ply"
]
}
],
"links": [
[
1,
1,
26,
21,
0,
2,
0,
"GS_RAW"
],
[
2,
2,
0,
3,
0,
"POINTCLOUD"
],
[
3,
3,
0,
4,
22,
0,
"MESH"
],
[
36,
26,
0,
28,
0,
"GS_PLY"
],
[
37,
28,
0,
21,
0,
"POINTCLOUD"
]
],
"groups": [],
_Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS.json
+242
View File
@@ -0,0 +1,242 @@
{
"last_node_id": 9,
"last_link_id": 11,
"nodes": [
{
"id": 7,
"type": "[ComfyUI-3D] Run Triplane Gaussian Transformers",
"pos": [
1032,
627
],
"size": {
"0": 430.8000183105469,
"1": 98
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "reference_image",
"type": "IMAGE",
"link": 6
},
{
"name": "reference_mask",
"type": "MASK",
"link": 11
},
{
"name": "tgs_model",
"type": "TGS_MODEL",
"link": 8
}
],
"outputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"links": [
9
],
"shape": 3,
"slot_index": 0
},
{
"name": "gaussian_sh_degree",
"type": "FLOAT",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Run Triplane Gaussian Transformers"
},
"widgets_values": [
1.9
]
},
{
"id": 9,
"type": "InvertMask",
"pos": [
506,
647
],
"size": {
"0": 210,
"1": 26
},
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 10
}
],
"outputs": [
{
"name": "MASK",
"type": "MASK",
"links": [
11
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "InvertMask"
}
},
{
"id": 2,
"type": "[ComfyUI-3D] Load Triplane Gaussian Transformers",
"pos": [
579,
735
],
"size": {
"0": 369,
"1": 47.66668701171875
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "tgs_model",
"type": "TGS_MODEL",
"links": [
8
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Load Triplane Gaussian Transformers"
}
},
{
"id": 8,
"type": "[ComfyUI-3D] Save 3DGS",
"pos": [
1531,
627
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"link": 9
}
],
"properties": {
"Node name for S&R": "[ComfyUI-3D] Save 3DGS"
},
"widgets_values": [
"MeshTest/3DGS_TGS.ply"
]
},
{
"id": 4,
"type": "LoadImage",
"pos": [
87,
211
],
"size": {
"0": 354.32330322265625,
"1": 315.4895935058594
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
6
],
"shape": 3,
"slot_index": 0
},
{
"name": "MASK",
"type": "MASK",
"links": [
10
],
"shape": 3,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "LoadImage"
},
"widgets_values": [
"a_pikachu_with_smily_face-removebg.png",
"image"
]
}
],
"links": [
[
6,
4,
0,
7,
0,
"IMAGE"
],
[
8,
2,
0,
7,
2,
"TGS_MODEL"
],
[
9,
7,
0,
8,
0,
"GS_PLY"
],
[
10,
4,
1,
9,
0,
"MASK"
],
[
11,
9,
0,
7,
1,
"MASK"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
_Example_Workflows/_Example_Inputs_Files/a_pikachu_with_smily_face-removebg.png
BIN
View File
Binary file not shown.
Side by Side

After

Width:  |  Height:  |  Size: 393 KiB

algorithms/diff_texturing.py
+1
View File
@@ -29,6 +29,7 @@ class DiffTextureBaker:
self.batch_size = batch_size
def prepare_img(self, img):
# (H, W, C) -> (C, H, W)
img_new = img.permute(2, 0, 1).unsqueeze(0).to(self.device)
img_new = F.interpolate(img_new, (self.ref_size_H, self.ref_size_W), mode="bilinear", align_corners=False).contiguous()
return img_new
algorithms/main_3DGS_renderer.py
+15 -45
View File
@@ -31,7 +31,7 @@ from simple_knn._C import distCUDA2
from ..shared_utils.sh_utils import eval_sh, SH2RGB, RGB2SH
from ..mesh_processer.mesh import Mesh
from ..mesh_processer.mesh_utils import decimate_mesh, clean_mesh
from ..mesh_processer.mesh_utils import read_gs_ply
def inverse_sigmoid(x):
return torch.log(x/(1-x))
@@ -500,6 +500,11 @@ class GaussianModel:
lr = self.xyz_scheduler_args(iteration)
param_group['lr'] = lr
return lr
def reset_opacity(self):
opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))
optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity")
self._opacity = optimizable_tensors["opacity"]
def construct_list_of_attributes(self):
l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
@@ -525,9 +530,7 @@ class GaussianModel:
def get_mesh(self):
xyz = self._xyz.detach().cpu().numpy()
def save_ply(self, path):
os.makedirs(os.path.dirname(path), exist_ok=True)
def to_ply(self):
xyz = self._xyz.detach().cpu().numpy()
normals = np.zeros_like(xyz)
f_dc = self._features_dc.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
@@ -542,45 +545,10 @@ class GaussianModel:
attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
elements[:] = list(map(tuple, attributes))
el = PlyElement.describe(elements, 'vertex')
PlyData([el]).write(path)
return PlyData([el])
def reset_opacity(self):
opacities_new = inverse_sigmoid(torch.min(self.get_opacity, torch.ones_like(self.get_opacity)*0.01))
optimizable_tensors = self.replace_tensor_to_optimizer(opacities_new, "opacity")
self._opacity = optimizable_tensors["opacity"]
def load_ply(self, path):
plydata = PlyData.read(path)
xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
np.asarray(plydata.elements[0]["y"]),
np.asarray(plydata.elements[0]["z"])), axis=1)
opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
print("Number of points at loading : ", xyz.shape[0])
features_dc = np.zeros((xyz.shape[0], 3, 1))
features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
assert len(extra_f_names)==3*(self.max_sh_degree + 1) ** 2 - 3
features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
for idx, attr_name in enumerate(extra_f_names):
features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
# Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
features_extra = features_extra.reshape((features_extra.shape[0], 3, (self.max_sh_degree + 1) ** 2 - 1))
scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
scales = np.zeros((xyz.shape[0], len(scale_names)))
for idx, attr_name in enumerate(scale_names):
scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
rots = np.zeros((xyz.shape[0], len(rot_names)))
for idx, attr_name in enumerate(rot_names):
rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
def create_from_ply(self, plydata):
xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata)
self._xyz = nn.Parameter(torch.tensor(xyz, dtype=torch.float, device="cuda").requires_grad_(True))
self._features_dc = nn.Parameter(torch.tensor(features_dc, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
@@ -894,6 +862,11 @@ class Renderer:
if isinstance(input, Mesh):
# load from 3D mesh
self.gaussians.create_from_mesh(input, num_pts)
elif isinstance(input, PlyData):
self.gaussians.create_from_ply(input)
elif isinstance(input, PointCloud):
# load from a provided pcd
self.gaussians.create_from_pcd(input, 1)
elif input is None:
# init from random point cloud
@@ -913,9 +886,6 @@ class Renderer:
points=xyz, colors=SH2RGB(shs), normals=np.zeros((num_pts, 3))
)
self.gaussians.create_from_pcd(pcd, 10)
elif isinstance(input, PointCloud):
# load from a provided pcd
self.gaussians.create_from_pcd(input, 1)
else:
self.gaussians.create_from_uv_data(input)
algorithms/triplane_gaussian_transformers.py
+167
View File
@@ -0,0 +1,167 @@
import torch
from dataclasses import dataclass, field
from einops import rearrange
from ..tgs import find
from ..tgs.models.image_feature import ImageFeature
from ..tgs.utils.saving import SaverMixin
from ..tgs.utils.config import parse_structured
from ..tgs.utils.ops import points_projection
from ..tgs.utils.misc import load_module_weights
from ..tgs.utils.typing import *
class TGS(torch.nn.Module, SaverMixin):
@dataclass
class Config:
weights: Optional[str] = None
weights_ignore_modules: Optional[List[str]] = None
camera_embedder_cls: str = ""
camera_embedder: dict = field(default_factory=dict)
image_feature: dict = field(default_factory=dict)
image_tokenizer_cls: str = ""
image_tokenizer: dict = field(default_factory=dict)
tokenizer_cls: str = ""
tokenizer: dict = field(default_factory=dict)
backbone_cls: str = ""
backbone: dict = field(default_factory=dict)
post_processor_cls: str = ""
post_processor: dict = field(default_factory=dict)
renderer_cls: str = ""
renderer: dict = field(default_factory=dict)
pointcloud_generator_cls: str = ""
pointcloud_generator: dict = field(default_factory=dict)
pointcloud_encoder_cls: str = ""
pointcloud_encoder: dict = field(default_factory=dict)
cfg: Config
def load_weights(self, weights: str, ignore_modules: Optional[List[str]] = None):
state_dict = load_module_weights(
weights, ignore_modules=ignore_modules, map_location="cpu"
)
self.load_state_dict(state_dict, strict=False)
def __init__(self, cfg):
super().__init__()
self.cfg = parse_structured(self.Config, cfg)
self.image_tokenizer = find(self.cfg.image_tokenizer_cls)(
self.cfg.image_tokenizer
)
assert self.cfg.camera_embedder_cls == 'tgs.models.networks.MLP'
weights = self.cfg.camera_embedder.pop("weights") if "weights" in self.cfg.camera_embedder else None
self.camera_embedder = find(self.cfg.camera_embedder_cls)(**self.cfg.camera_embedder)
if weights:
from ..tgs.utils.misc import load_module_weights
weights_path, module_name = weights.split(":")
state_dict = load_module_weights(
weights_path, module_name=module_name, map_location="cpu"
)
self.camera_embedder.load_state_dict(state_dict)
self.image_feature = ImageFeature(self.cfg.image_feature)
self.tokenizer = find(self.cfg.tokenizer_cls)(self.cfg.tokenizer)
self.backbone = find(self.cfg.backbone_cls)(self.cfg.backbone)
self.post_processor = find(self.cfg.post_processor_cls)(
self.cfg.post_processor
)
self.renderer = find(self.cfg.renderer_cls)(self.cfg.renderer)
# pointcloud generator
self.pointcloud_generator = find(self.cfg.pointcloud_generator_cls)(self.cfg.pointcloud_generator)
self.point_encoder = find(self.cfg.pointcloud_encoder_cls)(self.cfg.pointcloud_encoder)
# load checkpoint
if self.cfg.weights is not None:
self.load_weights(self.cfg.weights, self.cfg.weights_ignore_modules)
def _forward(self, batch: Dict[str, Any]) -> Dict[str, Any]:
# generate point cloud
out = self.pointcloud_generator(batch)
pointclouds = out["points"]
batch_size, n_input_views = batch["rgb_cond"].shape[:2]
# Camera modulation
camera_extri = batch["c2w_cond"].view(*batch["c2w_cond"].shape[:-2], -1)
camera_intri = batch["intrinsic_normed_cond"].view(*batch["intrinsic_normed_cond"].shape[:-2], -1)
camera_feats = torch.cat([camera_intri, camera_extri], dim=-1)
camera_feats = self.camera_embedder(camera_feats)
input_image_tokens: Float[Tensor, "B Cit Nit"] = self.image_tokenizer(
rearrange(batch["rgb_cond"], 'B Nv H W C -> B Nv C H W'),
modulation_cond=camera_feats,
)
input_image_tokens = rearrange(input_image_tokens, 'B Nv C Nt -> B (Nv Nt) C', Nv=n_input_views)
# get image features for projection
image_features = self.image_feature(
rgb = batch["rgb_cond"],
mask = batch.get("mask_cond", None),
feature = input_image_tokens
)
# only support number of input view is one
c2w_cond = batch["c2w_cond"].squeeze(1)
intrinsic_cond = batch["intrinsic_cond"].squeeze(1)
proj_feats = points_projection(pointclouds, c2w_cond, intrinsic_cond, image_features)
point_cond_embeddings = self.point_encoder(torch.cat([pointclouds, proj_feats], dim=-1))
tokens: Float[Tensor, "B Ct Nt"] = self.tokenizer(batch_size, cond_embeddings=point_cond_embeddings)
tokens = self.backbone(
tokens,
encoder_hidden_states=input_image_tokens,
modulation_cond=None,
)
scene_codes = self.post_processor(self.tokenizer.detokenize(tokens))
rend_out = self.renderer(scene_codes,
query_points=pointclouds,
additional_features=proj_feats,
**batch)
return {**out, **rend_out}
def forward(self, batch):
plys = []
out = self._forward(batch)
batch_size = batch["index"].shape[0]
for b in range(batch_size):
if batch["view_index"][b, 0] == 0:
plys.append(out["3dgs"][b].to_ply())
#self.render_test_video(out, batch, b)
return plys
def render_test_video(self, out, batch, b):
for index, render_image in enumerate(out["comp_rgb"][b]):
view_index = batch["view_index"][b, index]
self.save_image_grid(
f"video/{batch['instance_id'][b]}/{view_index}.png",
[
{
"type": "rgb",
"img": render_image,
"kwargs": {"data_format": "HWC"},
}
]
)
configs/tgs_config.yaml
+150
View File
@@ -0,0 +1,150 @@
data:
background_color: [1.0, 1.0, 1.0]
cond_width: 252 # multiply of 14
cond_height: 252
relative_pose: true
num_workers: 16
eval_batch_size: 1
eval_height: 512
eval_width: 512
system:
camera_embedder_cls: tgs.models.networks.MLP
camera_embedder:
dim_in: 25 # c2w + [fx, fy, cx, cy]
dim_out: 768
n_neurons: 768
n_hidden_layers: 1
activation: silu
image_feature:
out_dim: 773
image_tokenizer_cls: tgs.models.tokenizers.image.DINOV2SingleImageTokenizer
image_tokenizer:
pretrained_model_name_or_path: "facebook/dinov2-base"
width: ${data.cond_width}
height: ${data.cond_height}
modulation: true
modulation_zero_init: true
modulation_single_layer: true
modulation_cond_dim: ${system.camera_embedder.dim_out} # c2w + intrinsic
freeze_backbone_params: false
enable_memory_efficient_attention: ${system.backbone.enable_memory_efficient_attention}
enable_gradient_checkpointing: ${system.backbone.gradient_checkpointing}
tokenizer_cls: tgs.models.tokenizers.triplane.TriplaneLearnablePositionalEmbedding
tokenizer:
plane_size: 32
num_channels: 512
backbone_cls: tgs.models.transformers.Transformer1D
backbone:
in_channels: ${system.tokenizer.num_channels}
num_attention_heads: 8
attention_head_dim: 64
num_layers: 10
cross_attention_dim: 768 # hard-code, =DINO feature dim
norm_type: "layer_norm"
enable_memory_efficient_attention: false
gradient_checkpointing: false
post_processor_cls: tgs.models.networks.TriplaneUpsampleNetwork
post_processor:
in_channels: ${system.tokenizer.num_channels}
out_channels: 80
pointcloud_generator_cls: tgs.models.pointclouds.simplepoint.SimplePointGenerator
pointcloud_generator:
camera_embedder_cls: tgs.models.networks.MLP
camera_embedder:
dim_in: 25 # c2w + [fx, fy, cx, cy]
dim_out: 768
n_neurons: 768
n_hidden_layers: 1
activation: silu
image_tokenizer_cls: tgs.models.tokenizers.image.DINOV2SingleImageTokenizer
image_tokenizer:
pretrained_model_name_or_path: "facebook/dinov2-base"
width: ${data.cond_width}
height: ${data.cond_height}
modulation: true
modulation_zero_init: true
modulation_single_layer: true
modulation_cond_dim: ${system.camera_embedder.dim_out} # c2w + intrinsic
freeze_backbone_params: true
enable_memory_efficient_attention: ${system.backbone.enable_memory_efficient_attention}
enable_gradient_checkpointing: false
tokenizer_cls: tgs.models.tokenizers.point.PointLearnablePositionalEmbedding
tokenizer:
num_pcl: 2048
num_channels: 512
backbone_cls: tgs.models.transformers.Transformer1D
backbone:
in_channels: ${system.pointcloud_generator.tokenizer.num_channels}
num_attention_heads: 8
attention_head_dim: 64
num_layers: 10
cross_attention_dim: 768 # hard-code, =DINO feature dim
norm_type: "layer_norm"
enable_memory_efficient_attention: ${system.backbone.enable_memory_efficient_attention}
gradient_checkpointing: ${system.backbone.gradient_checkpointing}
post_processor_cls: tgs.models.networks.PointOutLayer
post_processor:
in_channels: 512
out_channels: 3
pointcloud_upsampling_cls: tgs.models.snowflake.model_spdpp.SnowflakeModelSPDPP
pointcloud_upsampling:
input_channels: 768
dim_feat: 128
num_p0: 2048
radius: 1
bounding: true
use_fps: true
up_factors: [2,4]
token_type: "image_token"
pointcloud_encoder_cls: tgs.models.pointclouds.pointnet.LocalPoolPointnet
pointcloud_encoder:
input_channels: 776 # 3 + 3 + 768 + 1 + 1 [xyz, local features]
c_dim: ${system.tokenizer.num_channels}
hidden_dim: 128
plane_size: ${system.tokenizer.plane_size}
n_blocks: 5
radius: ${system.renderer.radius}
renderer_cls: tgs.models.renderer.GS3DRenderer
renderer:
sh_degree: 3
radius: 0.6
mlp_network_config:
n_neurons: ${system.renderer.gs_out.in_channels}
n_hidden_layers: 2
activation: silu
gs_out:
in_channels: 128
xyz_offset: true
restrict_offset: true
use_rgb: false
feature_channels:
xyz: 3
scaling: 3
rotation: 4
opacity: 1
shs: ${shsdim:${system.renderer.sh_degree}}
clip_scaling: 0.2
install.bat
+7 -2
View File
@@ -16,8 +16,13 @@ if exist "%python_miniconda_exec%" (
git clone --recursive https://github.com/NVlabs/nvdiffrast
"%python_miniconda_exec%" -s -m pip install ./nvdiffrast
"%python_miniconda_exec%" -s -m pip install git+https://github.com/facebookresearch/pytorch3d.git@stable
)
else (
"%python_miniconda_exec%" -s -m pip install git+https://github.com/rusty1s/pytorch_scatter.git
set "python_miniconda_exec=..\..\..\..\..\..\..\python_miniconda_env\ComfyUI\python.exe"
cd tgs/models/snowflake/pointnet2_ops_lib && "%python_miniconda_exec%" -s setup.py install && cd ../../../../
) else (
echo ERROR: Cannot find Miniconda Environment "%python_miniconda_exec%"
)
mesh_processer/mesh.py
-4
View File
@@ -240,10 +240,6 @@ class Mesh:
albedo = albedo.astype(np.float32) / 255
print(f"[load_obj] load texture: {albedo.shape}")
# import matplotlib.pyplot as plt
# plt.imshow(albedo)
# plt.show()
mesh.albedo = torch.tensor(albedo, dtype=torch.float32, device=device)
return mesh
mesh_processer/mesh_utils.py
+44
View File
@@ -2,6 +2,9 @@ import torch
import numpy as np
import pymeshlab as pml
from .mesh import PointCloud
from ..shared_utils.sh_utils import SH2RGB, RGB2SH
def _base_face_areas(face_vertices_0, face_vertices_1, face_vertices_2):
"""Base function to compute the face areas."""
x1, x2, x3 = torch.split(face_vertices_0 - face_vertices_1, 1, dim=-1)
@@ -297,3 +300,44 @@ def clean_mesh(
)
return verts, faces
def read_gs_ply(plydata):
xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
np.asarray(plydata.elements[0]["y"]),
np.asarray(plydata.elements[0]["z"])), axis=1)
opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
features_dc = np.zeros((xyz.shape[0], 3, 1))
features_dc[:, 0, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
features_dc[:, 1, 0] = np.asarray(plydata.elements[0]["f_dc_1"])
features_dc[:, 2, 0] = np.asarray(plydata.elements[0]["f_dc_2"])
extra_f_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("f_rest_")]
#assert len(extra_f_names)!=3*(max_sh_degree + 1) ** 2 - 3:
max_sh_degree = int(((len(extra_f_names) + 3) / 3) ** 0.5 - 1)
features_extra = np.zeros((xyz.shape[0], len(extra_f_names)))
for idx, attr_name in enumerate(extra_f_names):
features_extra[:, idx] = np.asarray(plydata.elements[0][attr_name])
# Reshape (P,F*SH_coeffs) to (P, F, SH_coeffs except DC)
features_extra = features_extra.reshape((features_extra.shape[0], 3, (max_sh_degree + 1) ** 2 - 1))
scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
scales = np.zeros((xyz.shape[0], len(scale_names)))
for idx, attr_name in enumerate(scale_names):
scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot")]
rots = np.zeros((xyz.shape[0], len(rot_names)))
for idx, attr_name in enumerate(rot_names):
rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
return xyz, features_dc, features_extra, opacities, scales, rots
def ply_to_points_cloud(plydata):
xyz, features_dc, features_extra, opacities, scales, rots = read_gs_ply(plydata)
features_dc = np.transpose(features_dc, (0, 2, 1)) # equivalent of torch.transpose(features_dc, 1, 2)
features_extra = np.transpose(features_extra, (0, 2, 1))
shs = np.concatenate((features_dc, features_extra), axis=1)
normals = np.zeros_like(xyz)
pcd = PointCloud(points=xyz, colors=SH2RGB(shs), normals=normals)
return pcd
nodes.py
+195 -66
View File
@@ -1,18 +1,31 @@
import os
import sys
import re
import math
import copy
from enum import Enum
import subprocess
import folder_paths as comfy_paths
import torch
from torch.utils.data import DataLoader
import numpy as np
from plyfile import PlyData
from .shared_utils.common_utils import cstr
from .mesh_processer.mesh import Mesh
from .mesh_processer.mesh_utils import ply_to_points_cloud
from .algorithms.main_3DGS import GaussianSplatting, GSParams
from .algorithms.diff_texturing import DiffTextureBaker
from .algorithms.dmtet import DMTetMesh
from .algorithms.triplane_gaussian_transformers import TGS
from .tgs.utils.config import ExperimentConfig, load_config
from .tgs.data import CustomImageOrbitDataset
from .tgs.utils.misc import todevice, get_device
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "ComfyUI-3D-Pack")
sys.path.append(ROOT_PATH)
MANIFEST = {
"name": "ComfyUI-3D-Pack",
@@ -83,38 +96,34 @@ class Load_3DGS:
def INPUT_TYPES(cls):
return {
"required": {
"raw_3DGS_file_path": ("STRING", {"default": '', "multiline": False}),
"gs_ply_file_path": ("STRING", {"default": '', "multiline": False}),
},
}
RETURN_TYPES = (
"GS_RAW",
"GS_PLY",
)
RETURN_NAMES = (
"raw_3DGS",
"gs_ply",
)
FUNCTION = "load_mesh"
FUNCTION = "load_gs"
CATEGORY = "ComfyUI3D/Import|Export"
def load_mesh(self, raw_3DGS_file_path):
raw_3DGS = None
def load_gs(self, gs_ply_file_path):
gs_ply = None
if not os.path.isabs(raw_3DGS_file_path):
raw_3DGS_file_path = os.path.join(comfy_paths.input_directory, raw_3DGS_file_path)
if not os.path.isabs(gs_ply_file_path):
gs_ply_file_path = os.path.join(comfy_paths.input_directory, gs_ply_file_path)
if os.path.exists(raw_3DGS_file_path):
folder, filename = os.path.split(raw_3DGS_file_path)
if os.path.exists(gs_ply_file_path):
folder, filename = os.path.split(gs_ply_file_path)
if filename.lower().endswith(SUPPORTED_3DGS_EXTENSIONS):
with torch.inference_mode(False):
raw_3DGS = GaussianSplatting()
raw_3DGS.renderer.gaussians.load_ply(raw_3DGS_file_path)
# TODO: better way of loading 3DGS, since this method omit render & training parameters
gs_ply = PlyData.read(gs_ply_file_path)
else:
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3DGS file extensions: {SUPPORTED_3DGS_EXTENSIONS}").error.print()
else:
cstr(f"[{self.__class__.__name__}] File {raw_3DGS_file_path} does not exist").error.print()
return (raw_3DGS, )
cstr(f"[{self.__class__.__name__}] File {gs_ply_file_path} does not exist").error.print()
return (gs_ply, )
class Save_3D_Mesh:
@@ -156,7 +165,7 @@ class Save_3DGS:
def INPUT_TYPES(cls):
return {
"required": {
"raw_3DGS": ("GS_RAW",),
"gs_ply": ("GS_PLY",),
"save_path": ("STRING", {"default": '3DGS_[time(%Y-%m-%d)].ply', "multiline": False}),
},
}
@@ -166,7 +175,7 @@ class Save_3DGS:
FUNCTION = "save_3DGS"
CATEGORY = "ComfyUI3D/Import|Export"
def save_3DGS(self, raw_3DGS, save_path):
def save_3DGS(self, gs_ply, save_path):
folder_path, filename = os.path.split(save_path)
@@ -177,8 +186,7 @@ class Save_3DGS:
os.makedirs(folder_path, exist_ok=True)
raw_3DGS.renderer.gaussians.save_ply(save_path)
# TODO: better way of saving 3DGS, since this method omit render & training parameters
gs_ply.write(save_path)
else:
cstr(f"[{self.__class__.__name__}] File name {filename} does not end with supported 3DGS file extensions: {SUPPORTED_3DGS_EXTENSIONS}").error.print()
@@ -190,7 +198,7 @@ class Convert_3DGS_To_Pointcloud:
def INPUT_TYPES(cls):
return {
"required": {
"raw_3DGS": ("GS_RAW",),
"gs_ply": ("GS_PLY",),
},
}
@@ -200,12 +208,12 @@ class Convert_3DGS_To_Pointcloud:
RETURN_NAMES = (
"points_cloud",
)
FUNCTION = "convert_3DGS"
FUNCTION = "convert_gs_ply"
CATEGORY = "ComfyUI3D/Preprocessor"
def convert_3DGS(self, raw_3DGS):
def convert_gs_ply(self, gs_ply):
points_cloud = raw_3DGS.renderer.gaussians.get_points_cloud()
points_cloud = ply_to_points_cloud(gs_ply)
return (points_cloud, )
@@ -510,7 +518,7 @@ class Gaussian_Splatting:
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
"reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
"training_iterations": ("INT", {"default": 1000, "min": 1, "max": 100000}),
"batch_size": ("INT", {"default": 5, "min": 1, "max": 0xffffffffffffffff}),
"batch_size": ("INT", {"default": 3, "min": 1, "max": 0xffffffffffffffff}),
"loss_value_scale": ("FLOAT", {"default": 10000.0, "min": 1.0}),
"ms_ssim_loss_weight": ("FLOAT", {"default": 0.2, "min": 0.0, "max": 1.0, }),
"alpha_loss_weight": ("FLOAT", {"default": 3, "min": 0.0, }),
@@ -537,53 +545,58 @@ class Gaussian_Splatting:
},
"optional": {
"points_cloud": ("POINTCLOUD",),
"ply_to_initialize_gaussian": ("GS_PLY",),
"mesh_to_initialize_gaussian": ("MESH",),
}
}
RETURN_TYPES = (
"GS_RAW",
"GS_PLY",
)
RETURN_NAMES = (
"raw_3DGS",
"gs_ply",
)
FUNCTION = "run_3DGS"
CATEGORY = "ComfyUI3D/Algorithm"
def run_3DGS(self,
reference_images,
reference_masks,
reference_orbit_camera_poses,
reference_orbit_camera_fovy,
training_iterations,
batch_size,
loss_value_scale,
ms_ssim_loss_weight,
alpha_loss_weight,
offset_loss_weight,
offset_opacity_loss_weight,
invert_background_probability,
feature_learning_rate,
opacity_learning_rate,
scaling_learning_rate,
rotation_learning_rate,
position_learning_rate_init,
position_learning_rate_final,
position_learning_rate_delay_mult,
position_learning_rate_max_steps,
initial_gaussians_num,
K_nearest_neighbors,
percent_dense,
density_start_iterations,
density_end_iterations,
densification_interval,
opacity_reset_interval,
densify_grad_threshold,
gaussian_sh_degree,
mesh_to_initialize_gaussian=None):
def run_3DGS(
self,
reference_images,
reference_masks,
reference_orbit_camera_poses,
reference_orbit_camera_fovy,
training_iterations,
batch_size,
loss_value_scale,
ms_ssim_loss_weight,
alpha_loss_weight,
offset_loss_weight,
offset_opacity_loss_weight,
invert_background_probability,
feature_learning_rate,
opacity_learning_rate,
scaling_learning_rate,
rotation_learning_rate,
position_learning_rate_init,
position_learning_rate_final,
position_learning_rate_delay_mult,
position_learning_rate_max_steps,
initial_gaussians_num,
K_nearest_neighbors,
percent_dense,
density_start_iterations,
density_end_iterations,
densification_interval,
opacity_reset_interval,
densify_grad_threshold,
gaussian_sh_degree,
points_cloud=None,
ply_to_initialize_gaussian=None,
mesh_to_initialize_gaussian=None,
):
raw_3DGS = None
gs_ply = None
ref_imgs_num = len(reference_images)
ref_masks_num = len(reference_masks)
@@ -624,18 +637,26 @@ class Gaussian_Splatting:
densify_grad_threshold,
gaussian_sh_degree)
gs = GaussianSplatting(gs_params, mesh_to_initialize_gaussian)
if points_cloud is not None:
gs_init_input = points_cloud
elif ply_to_initialize_gaussian is not None:
gs_init_input = ply_to_initialize_gaussian
else:
gs_init_input = mesh_to_initialize_gaussian
gs = GaussianSplatting(gs_params, gs_init_input)
gs.prepare_training(reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy)
gs.training()
raw_3DGS = gs
gs_ply = gs.renderer.gaussians.to_ply()
else:
cstr(f"[{self.__class__.__name__}] Number of reference images {ref_imgs_num} does not equal to number of reference camera poses {ref_cam_poses_num}").error.print()
else:
cstr(f"[{self.__class__.__name__}] Number of reference images {ref_imgs_num} does not equal to number of masks {ref_masks_num}").error.print()
return (raw_3DGS, )
return (gs_ply, )
class Bake_Texture_To_Mesh:
@@ -652,7 +673,7 @@ class Bake_Texture_To_Mesh:
"reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
"mesh": ("MESH",),
"training_iterations": ("INT", {"default": 1000, "min": 1, "max": 100000}),
"batch_size": ("INT", {"default": 5, "min": 1, "max": 0xffffffffffffffff}),
"batch_size": ("INT", {"default": 3, "min": 1, "max": 0xffffffffffffffff}),
"texture_learning_rate": ("FLOAT", {"default": 0.1, "min": 0.00001, "step": 0.00001}),
"train_mesh_geometry": ("BOOLEAN", {"default": False},),
"geometry_learning_rate": ("FLOAT", {"default": 0.01, "min": 0.00001, "step": 0.00001}),
@@ -761,4 +782,112 @@ class Deep_Marching_Tetrahedrons:
dmtet_mesh.training(reference_points_cloud, reference_images, reference_masks)
mesh = dmtet_mesh.get_mesh()
return (mesh, )
return (mesh, )
TGS_CONFIG_PATH = "configs/tgs_config.yaml"
class Load_Triplane_Gaussian_Transformers:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
},
"optional": {
}
}
RETURN_TYPES = (
"TGS_MODEL",
)
RETURN_NAMES = (
"tgs_model",
)
FUNCTION = "load_TGS"
CATEGORY = "ComfyUI3D/Import|Export"
def load_TGS(self):
device = get_device()
config_path = os.path.join(ROOT_PATH, TGS_CONFIG_PATH)
cfg: ExperimentConfig = load_config(config_path)
# Download pre-trained model if it not exist locally
model_name = "model_lvis_rel.ckpt"
model_path = os.path.join(ROOT_PATH, f"checkpoints/{model_name}")
if not os.path.isfile(model_path):
from huggingface_hub import hf_hub_download
hf_hub_download(repo_id="VAST-AI/TriplaneGaussian", local_dir="./checkpoints", filename=model_name, repo_type="model")
cfg.system.weights=model_path
tgs_model = TGS(cfg=cfg.system).to(device)
save_path = os.path.join(ROOT_PATH, "outputs")
tgs_model.set_save_dir(save_path)
cstr(f"[{self.__class__.__name__}] loaded model ckpt from {model_path}").msg.print()
return (tgs_model, )
class Run_Triplane_Gaussian_Transformers:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"reference_image": ("IMAGE", ),
"reference_mask": ("MASK",),
"tgs_model": ("TGS_MODEL", ),
"cam_dist": ("FLOAT", {"default": 1.9, "min": 0.01, "step": 0.01}),
},
}
OUTPUT_NODE = True
RETURN_TYPES = (
"GS_PLY",
"FLOAT"
)
RETURN_NAMES = (
"gs_ply",
"gaussian_sh_degree"
)
FUNCTION = "run_TGS"
CATEGORY = "ComfyUI3D/Algorithm"
def run_TGS(self, reference_image, reference_mask, tgs_model, cam_dist):
config_path = os.path.join(ROOT_PATH, TGS_CONFIG_PATH)
cfg: ExperimentConfig = load_config(config_path)
#save_path = os.path.join(ROOT_PATH, "outputs")
#tgs_model.set_save_dir(save_path)
cfg.data.cond_camera_distance = cam_dist
cfg.data.eval_camera_distance = cam_dist
dataset = CustomImageOrbitDataset(reference_image, reference_mask, cfg.data)
dataloader = DataLoader(dataset,
batch_size=cfg.data.eval_batch_size,
shuffle=False,
collate_fn=dataset.collate
)
gs_ply = []
for batch in dataloader:
batch = todevice(batch)
gs_ply.extend(tgs_model(batch))
"""
# Output rendered video, for testing this node only
tgs_model.save_img_sequences(
"video",
"(\d+)\.png",
save_format="mp4",
fps=30,
delete=True,
)
"""
return (gs_ply[0], tgs_model.renderer.cfg.sh_degree)
requirements.txt
+12 -1
View File
@@ -1,14 +1,24 @@
#base
cmake
ninja
kornia
OmegaConf
numpy
einops
scipy
opencv-python
kornia
# for use ML models
diffusers
transformers
# for training differentiable tensors
pytorch_msssim
# for process images & videos
imageio
imageio-ffmpeg
# for dmtet and mesh import & export
trimesh
plyfile
@@ -19,3 +29,4 @@ pymeshlab
# extra
torchtyping
tqdm
jaxtyping
test.bat
+2 -4
View File
@@ -1,5 +1,3 @@
set "python_miniconda_exec=..\..\..\python_miniconda_env\ComfyUI\python.exe"
set "python_miniconda_exec=..\..\..\..\..\..\..\python_miniconda_env\ComfyUI\python.exe"
echo ERROR: Cannot find Miniconda Environment "%python_miniconda_exec%"
pause
cd tgs/models/snowflake/pointnet2_ops_lib && "%python_miniconda_exec%" -s setup.py install && cd ../../../../
tgs/__init__.py
+11
View File
@@ -0,0 +1,11 @@
import importlib
from .utils.typing import *
def find(cls_string) -> Type:
cls_full_name = cls_string.split(".")
module_string = ".".join(cls_full_name[:-1])
cls_name = cls_full_name[-1]
print(module_string)
module = importlib.import_module(module_string, package=__name__)
cls = getattr(module, cls_name)
return cls
tgs/data.py
+237
View File
@@ -0,0 +1,237 @@
import json
import math
from dataclasses import dataclass, field
import os
import imageio
import numpy as np
import torch
import torch.nn.functional as F
from PIL import Image
from torch.utils.data import Dataset
from .utils.config import parse_structured
from .utils.ops import get_intrinsic_from_fov, get_ray_directions, get_rays
from .utils.typing import *
@dataclass
class CustomImageDataModuleConfig:
background_color: Tuple[float, float, float] = field(
default_factory=lambda: (1.0, 1.0, 1.0)
)
relative_pose: bool = False
cond_height: int = 512
cond_width: int = 512
cond_camera_distance: float = 1.6
cond_fovy_deg: float = 40.0
cond_elevation_deg: float = 0.0
cond_azimuth_deg: float = 0.0
num_workers: int = 16
eval_height: int = 512
eval_width: int = 512
eval_batch_size: int = 1
eval_elevation_deg: float = 0.0
eval_camera_distance: float = 1.6
eval_fovy_deg: float = 40.0
n_test_views: int = 120
num_views_output: int = 120
only_3dgs: bool = False
class CustomImageOrbitDataset(Dataset):
def __init__(self, reference_image, reference_mask, cfg: Any) -> None:
super().__init__()
self.cfg: CustomImageDataModuleConfig = parse_structured(CustomImageDataModuleConfig, cfg)
self.n_views = self.cfg.n_test_views
assert self.n_views % self.cfg.num_views_output == 0
self.all_images = reference_image
self.all_masks = reference_mask
azimuth_deg: Float[Tensor, "B"] = torch.linspace(0, 360.0, self.n_views + 1)[
: self.n_views
]
elevation_deg: Float[Tensor, "B"] = torch.full_like(
azimuth_deg, self.cfg.eval_elevation_deg
)
camera_distances: Float[Tensor, "B"] = torch.full_like(
elevation_deg, self.cfg.eval_camera_distance
)
elevation = elevation_deg * math.pi / 180
azimuth = azimuth_deg * math.pi / 180
# convert spherical coordinates to cartesian coordinates
# right hand coordinate system, x back, y right, z up
# elevation in (-90, 90), azimuth from +x to +y in (-180, 180)
camera_positions: Float[Tensor, "B 3"] = torch.stack(
[
camera_distances * torch.cos(elevation) * torch.cos(azimuth),
camera_distances * torch.cos(elevation) * torch.sin(azimuth),
camera_distances * torch.sin(elevation),
],
dim=-1,
)
# default scene center at origin
center: Float[Tensor, "B 3"] = torch.zeros_like(camera_positions)
# default camera up direction as +z
up: Float[Tensor, "B 3"] = torch.as_tensor([0, 0, 1], dtype=torch.float32)[
None, :
].repeat(self.n_views, 1)
fovy_deg: Float[Tensor, "B"] = torch.full_like(
elevation_deg, self.cfg.eval_fovy_deg
)
fovy = fovy_deg * math.pi / 180
lookat: Float[Tensor, "B 3"] = F.normalize(center - camera_positions, dim=-1)
right: Float[Tensor, "B 3"] = F.normalize(torch.cross(lookat, up), dim=-1)
up = F.normalize(torch.cross(right, lookat), dim=-1)
c2w3x4: Float[Tensor, "B 3 4"] = torch.cat(
[torch.stack([right, up, -lookat], dim=-1), camera_positions[:, :, None]],
dim=-1,
)
c2w: Float[Tensor, "B 4 4"] = torch.cat(
[c2w3x4, torch.zeros_like(c2w3x4[:, :1])], dim=1
)
c2w[:, 3, 3] = 1.0
# get directions by dividing directions_unit_focal by focal length
focal_length: Float[Tensor, "B"] = (
0.5 * self.cfg.eval_height / torch.tan(0.5 * fovy)
)
directions_unit_focal = get_ray_directions(
H=self.cfg.eval_height,
W=self.cfg.eval_width,
focal=1.0,
)
directions: Float[Tensor, "B H W 3"] = directions_unit_focal[
None, :, :, :
].repeat(self.n_views, 1, 1, 1)
directions[:, :, :, :2] = (
directions[:, :, :, :2] / focal_length[:, None, None, None]
)
# must use normalize=True to normalize directions here
rays_o, rays_d = get_rays(directions, c2w, keepdim=True)
intrinsic: Float[Tensor, "B 3 3"] = get_intrinsic_from_fov(
self.cfg.eval_fovy_deg * math.pi / 180,
H=self.cfg.eval_height,
W=self.cfg.eval_width,
bs=self.n_views,
)
intrinsic_normed: Float[Tensor, "B 3 3"] = intrinsic.clone()
intrinsic_normed[..., 0, 2] /= self.cfg.eval_width
intrinsic_normed[..., 1, 2] /= self.cfg.eval_height
intrinsic_normed[..., 0, 0] /= self.cfg.eval_width
intrinsic_normed[..., 1, 1] /= self.cfg.eval_height
self.rays_o, self.rays_d = rays_o, rays_d
self.intrinsic = intrinsic
self.intrinsic_normed = intrinsic_normed
self.c2w = c2w
self.camera_positions = camera_positions
self.background_color = torch.as_tensor(self.cfg.background_color)
# condition
self.intrinsic_cond = get_intrinsic_from_fov(
np.deg2rad(self.cfg.cond_fovy_deg),
H=self.cfg.cond_height,
W=self.cfg.cond_width,
)
self.intrinsic_normed_cond = self.intrinsic_cond.clone()
self.intrinsic_normed_cond[..., 0, 2] /= self.cfg.cond_width
self.intrinsic_normed_cond[..., 1, 2] /= self.cfg.cond_height
self.intrinsic_normed_cond[..., 0, 0] /= self.cfg.cond_width
self.intrinsic_normed_cond[..., 1, 1] /= self.cfg.cond_height
if self.cfg.relative_pose:
self.c2w_cond = torch.as_tensor(
[
[0, 0, 1, self.cfg.cond_camera_distance],
[1, 0, 0, 0],
[0, 1, 0, 0],
[0, 0, 0, 1],
]
).float()
else:
cond_elevation = self.cfg.cond_elevation_deg * math.pi / 180
cond_azimuth = self.cfg.cond_azimuth_deg * math.pi / 180
cond_camera_position: Float[Tensor, "3"] = torch.as_tensor(
[
self.cfg.cond_camera_distance * np.cos(cond_elevation) * np.cos(cond_azimuth),
self.cfg.cond_camera_distance * np.cos(cond_elevation) * np.sin(cond_azimuth),
self.cfg.cond_camera_distance * np.sin(cond_elevation),
], dtype=torch.float32
)
cond_center: Float[Tensor, "3"] = torch.zeros_like(cond_camera_position)
cond_up: Float[Tensor, "3"] = torch.as_tensor([0, 0, 1], dtype=torch.float32)
cond_lookat: Float[Tensor, "3"] = F.normalize(cond_center - cond_camera_position, dim=-1)
cond_right: Float[Tensor, "3"] = F.normalize(torch.cross(cond_lookat, cond_up), dim=-1)
cond_up = F.normalize(torch.cross(cond_right, cond_lookat), dim=-1)
cond_c2w3x4: Float[Tensor, "3 4"] = torch.cat(
[torch.stack([cond_right, cond_up, -cond_lookat], dim=-1), cond_camera_position[:, None]],
dim=-1,
)
cond_c2w: Float[Tensor, "4 4"] = torch.cat(
[cond_c2w3x4, torch.zeros_like(cond_c2w3x4[:1])], dim=0
)
cond_c2w[3, 3] = 1.0
self.c2w_cond = cond_c2w
def __len__(self):
if self.cfg.only_3dgs:
return len(self.all_images)
else:
return len(self.all_images) * self.n_views // self.cfg.num_views_output
def __getitem__(self, index):
if self.cfg.only_3dgs:
scene_index = index
view_index = [0]
else:
scene_index = index * self.cfg.num_views_output // self.n_views
view_start = index % (self.n_views // self.cfg.num_views_output)
view_index = list(range(self.n_views))[view_start * self.cfg.num_views_output :
(view_start + 1) * self.cfg.num_views_output]
mask_cond: Float[Tensor, "Hc Wc 1"] = self.resize_image(self.all_masks[scene_index].unsqueeze(2))
rgb_cond: Float[Tensor, "Hc Wc 3"] = self.resize_image(self.all_images[scene_index]) * mask_cond + self.background_color[None, None, :] * (1 - mask_cond)
out = {
"rgb_cond": rgb_cond,
"c2w_cond": self.c2w_cond.unsqueeze(0),
"mask_cond": mask_cond,
"intrinsic_cond": self.intrinsic_cond.unsqueeze(0),
"intrinsic_normed_cond": self.intrinsic_normed_cond.unsqueeze(0),
"view_index": torch.as_tensor(view_index),
"rays_o": self.rays_o[view_index],
"rays_d": self.rays_d[view_index],
"intrinsic": self.intrinsic[view_index],
"intrinsic_normed": self.intrinsic_normed[view_index],
"c2w": self.c2w[view_index],
"camera_positions": self.camera_positions[view_index],
}
out["c2w"][..., :3, 1:3] *= -1
out["c2w_cond"][..., :3, 1:3] *= -1
out["index"] = torch.as_tensor(scene_index)
out["background_color"] = self.background_color
instance_id = str(scene_index)
out["instance_id"] = instance_id
return out
def resize_image(self, img):
img_new = img.permute(2, 0, 1).unsqueeze(0)
img_new = F.interpolate(img_new, (self.cfg.cond_width, self.cfg.cond_height), mode="bilinear", align_corners=False).contiguous()
return img_new.permute(0, 2, 3, 1)
def collate(self, batch):
batch = torch.utils.data.default_collate(batch)
batch.update({"height": self.cfg.eval_height, "width": self.cfg.eval_width})
return batch
tgs/models/__init__.py
View File
tgs/models/image_feature.py
+48
View File
@@ -0,0 +1,48 @@
from dataclasses import dataclass
import torch
import torch.nn.functional as F
from einops import rearrange
from ..utils.base import BaseModule
from ..utils.ops import compute_distance_transform
from ..utils.typing import *
class ImageFeature(BaseModule):
@dataclass
class Config(BaseModule.Config):
use_rgb: bool = True
use_feature: bool = True
use_mask: bool = True
feature_dim: int = 128
out_dim: int = 133
backbone: str = "default"
freeze_backbone_params: bool = True
cfg: Config
def forward(self, rgb, mask=None, feature=None):
B, Nv, H, W = rgb.shape[:4]
rgb = rearrange(rgb, "B Nv H W C -> (B Nv) C H W")
if mask is not None:
mask = rearrange(mask, "B Nv H W C -> (B Nv) C H W")
assert feature is not None
# reshape dino tokens to image-like size
feature = rearrange(feature, "B (Nv Nt) C -> (B Nv) Nt C", Nv=Nv)
feature = feature[:, 1:].reshape(B * Nv, H // 14, W // 14, -1).permute(0, 3, 1, 2).contiguous()
feature = F.interpolate(feature, size=(H, W), mode='bilinear', align_corners=False)
if mask is not None and mask.is_floating_point():
mask = mask > 0.5
image_features = []
if self.cfg.use_rgb:
image_features.append(rgb)
if self.cfg.use_feature:
image_features.append(feature)
if self.cfg.use_mask:
image_features += [mask, compute_distance_transform(mask)]
# detach features, occur error when with grad
image_features = torch.cat(image_features, dim=1)#.detach()
return rearrange(image_features, "(B Nv) C H W -> B Nv C H W", B=B, Nv=Nv).squeeze(1)
tgs/models/networks.py
+204
View File
@@ -0,0 +1,204 @@
from dataclasses import dataclass
import torch
import torch.nn as nn
from einops import rearrange
import numpy as np
from ..utils.base import BaseModule
from ..utils.ops import get_activation
from ..utils.typing import *
class PointOutLayer(BaseModule):
@dataclass
class Config(BaseModule.Config):
in_channels: int = 1024
out_channels: int = 3
cfg: Config
def configure(self) -> None:
super().configure()
self.point_layer = nn.Linear(self.cfg.in_channels, self.cfg.out_channels)
self.initialize_weights()
def initialize_weights(self):
nn.init.constant_(self.point_layer.weight, 0)
nn.init.constant_(self.point_layer.bias, 0)
def forward(self, x):
return self.point_layer(x)
class TriplaneUpsampleNetwork(BaseModule):
@dataclass
class Config(BaseModule.Config):
in_channels: int = 1024
out_channels: int = 80
cfg: Config
def configure(self) -> None:
super().configure()
self.upsample = nn.ConvTranspose2d(
self.cfg.in_channels, self.cfg.out_channels, kernel_size=2, stride=2
)
def forward(
self, triplanes: Float[Tensor, "B 3 Ci Hp Wp"]
) -> Float[Tensor, "B 3 Co Hp2 Wp2"]:
triplanes_up = rearrange(
self.upsample(
rearrange(triplanes, "B Np Ci Hp Wp -> (B Np) Ci Hp Wp", Np=3)
),
"(B Np) Co Hp Wp -> B Np Co Hp Wp",
Np=3,
)
return triplanes_up
class MLP(nn.Module):
def __init__(
self,
dim_in: int,
dim_out: int,
n_neurons: int,
n_hidden_layers: int,
activation: str = "relu",
output_activation: Optional[str] = None,
bias: bool = True,
):
super().__init__()
layers = [
self.make_linear(
dim_in, n_neurons, is_first=True, is_last=False, bias=bias
),
self.make_activation(activation),
]
for i in range(n_hidden_layers - 1):
layers += [
self.make_linear(
n_neurons, n_neurons, is_first=False, is_last=False, bias=bias
),
self.make_activation(activation),
]
layers += [
self.make_linear(
n_neurons, dim_out, is_first=False, is_last=True, bias=bias
)
]
self.layers = nn.Sequential(*layers)
self.output_activation = get_activation(output_activation)
def forward(self, x):
x = self.layers(x)
x = self.output_activation(x)
return x
def make_linear(self, dim_in, dim_out, is_first, is_last, bias=True):
layer = nn.Linear(dim_in, dim_out, bias=bias)
return layer
def make_activation(self, activation):
if activation == "relu":
return nn.ReLU(inplace=True)
elif activation == "silu":
return nn.SiLU(inplace=True)
else:
raise NotImplementedError
class GSProjection(nn.Module):
def __init__(self,
in_channels: int = 80,
sh_degree: int = 3,
init_scaling: float = -5.0,
init_density: float = 0.1) -> None:
super().__init__()
self.out_keys = GS_KEYS + ["shs"]
self.out_channels = GS_CHANNELS + [(sh_degree + 1) ** 2 * 3]
self.out_layers = nn.ModuleList()
for key, ch in zip(self.out_keys, self.out_channels):
layer = nn.Linear(in_channels, ch)
# initialize
nn.init.constant_(layer.weight, 0)
nn.init.constant_(layer.bias, 0)
if key == "scaling":
nn.init.constant_(layer.bias, init_scaling)
elif key == "rotation":
nn.init.constant_(layer.bias, 0)
nn.init.constant_(layer.bias[0], 1.0)
elif key == "opacity":
inverse_sigmoid = lambda x: np.log(x / (1 - x))
nn.init.constant_(layer.bias, inverse_sigmoid(init_density))
self.out_layers.append(layer)
def forward(self, x):
ret = []
for k, layer in zip(self.out_keys, self.out_layers):
v = layer(x)
if k == "rotation":
v = torch.nn.functional.normalize(v)
elif k == "scaling":
v = torch.exp(v)
# v = v.detach() # FIXME: for DEBUG
elif k == "opacity":
v = torch.sigmoid(v)
# elif k == "shs":
# v = torch.reshape(v, (v.shape[0], -1, 3))
ret.append(v)
ret = torch.cat(ret, dim=-1)
return ret
def get_encoding(n_input_dims: int, config) -> nn.Module:
raise NotImplementedError
def get_mlp(n_input_dims, n_output_dims, config) -> nn.Module:
raise NotImplementedError
# Resnet Blocks for pointnet
class ResnetBlockFC(nn.Module):
''' Fully connected ResNet Block class.
Args:
size_in (int): input dimension
size_out (int): output dimension
size_h (int): hidden dimension
'''
def __init__(self, size_in, size_out=None, size_h=None):
super().__init__()
# Attributes
if size_out is None:
size_out = size_in
if size_h is None:
size_h = min(size_in, size_out)
self.size_in = size_in
self.size_h = size_h
self.size_out = size_out
# Submodules
self.fc_0 = nn.Linear(size_in, size_h)
self.fc_1 = nn.Linear(size_h, size_out)
self.actvn = nn.ReLU()
if size_in == size_out:
self.shortcut = None
else:
self.shortcut = nn.Linear(size_in, size_out, bias=False)
# Initialization
nn.init.zeros_(self.fc_1.weight)
def forward(self, x):
net = self.fc_0(self.actvn(x))
dx = self.fc_1(self.actvn(net))
if self.shortcut is not None:
x_s = self.shortcut(x)
else:
x_s = x
return x_s + dx
tgs/models/pointclouds/LICENSE_POINTNET
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2020 Songyou Peng, Michael Niemeyer, Lars Mescheder, Marc Pollefeys, Andreas Geiger
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
tgs/models/pointclouds/pointnet.py
+121
View File
@@ -0,0 +1,121 @@
# modified from https://github.com/autonomousvision/convolutional_occupancy_networks/blob/master/src/encoder/pointnet.py
from dataclasses import dataclass
import torch
import torch.nn as nn
from torch_scatter import scatter_mean, scatter_max
from tgs.utils.base import BaseModule
from tgs.models.networks import ResnetBlockFC
from tgs.utils.ops import scale_tensor
class LocalPoolPointnet(BaseModule):
''' PointNet-based encoder network with ResNet blocks for each point.
Number of input points are fixed.
Args:
c_dim (int): dimension of latent code c
dim (int): input points dimension
hidden_dim (int): hidden dimension of the network
scatter_type (str): feature aggregation when doing local pooling
plane_resolution (int): defined resolution for plane feature
padding (float): conventional padding paramter of ONet for unit cube, so [-0.5, 0.5] -> [-0.55, 0.55]
n_blocks (int): number of blocks ResNetBlockFC layers
'''
@dataclass
class Config(BaseModule.Config):
input_channels: int = 3
c_dim: int = 128
hidden_dim: int = 128
scatter_type: str = "max"
plane_size: int = 32
n_blocks: int = 5
radius: float = 1.
cfg: Config
def configure(self) -> None:
super().configure()
self.fc_pos = nn.Linear(self.cfg.input_channels, 2 * self.cfg.hidden_dim)
self.blocks = nn.ModuleList([
ResnetBlockFC(2 * self.cfg.hidden_dim, self.cfg.hidden_dim) for i in range(self.cfg.n_blocks)
])
self.fc_c = nn.Linear(self.cfg.hidden_dim, self.cfg.c_dim)
self.actvn = nn.ReLU()
if self.cfg.scatter_type == 'max':
self.scatter = scatter_max
elif self.cfg.scatter_type == 'mean':
self.scatter = scatter_mean
else:
raise ValueError('incorrect scatter type')
def generate_plane_features(self, index, c):
# acquire indices of features in plane
# xy = normalize_coordinate(p.clone(), plane=plane, padding=self.padding) # normalize to the range of (0, 1)
# index = self.coordinate2index(x, self.cfg.plane_size)
# scatter plane features from points
fea_plane = c.new_zeros(index.shape[0], self.cfg.c_dim, self.cfg.plane_size ** 2)
c = c.permute(0, 2, 1) # B x 512 x T
fea_plane = scatter_mean(c, index, out=fea_plane) # B x 512 x reso^2
fea_plane = fea_plane.reshape(index.shape[0], self.cfg.c_dim, self.cfg.plane_size, self.cfg.plane_size) # sparce matrix (B x 512 x reso x reso)
return fea_plane
def pool_local(self, xy, index, c):
bs, fea_dim = c.shape[0], c.shape[2]
keys = xy.keys()
c_out = 0
for key in keys:
# scatter plane features from points
fea = self.scatter(c.permute(0, 2, 1), index[key], dim_size=self.cfg.plane_size ** 2)
if self.scatter == scatter_max:
fea = fea[0]
# gather feature back to points
fea = fea.gather(dim=2, index=index[key].expand(-1, fea_dim, -1))
c_out += fea
return c_out.permute(0, 2, 1)
def coordinate2index(self, x):
x = (x * self.cfg.plane_size).long()
index = x[..., 0] + self.cfg.plane_size * x[..., 1]
assert index.max() < self.cfg.plane_size ** 2
return index[:, None, :]
def forward(self, p):
batch_size, T, D = p.shape
# acquire the index for each point
coord = {}
index = {}
position = torch.clamp(p[..., :3], -self.cfg.radius + 1e-6, self.cfg.radius - 1e-6)
position_norm = scale_tensor(position, (-self.cfg.radius, self.cfg.radius), (0, 1))
coord["xy"] = position_norm[..., [0, 1]]
coord["xz"] = position_norm[..., [0, 2]]
coord["yz"] = position_norm[..., [1, 2]]
index["xy"] = self.coordinate2index(coord["xy"])
index["xz"] = self.coordinate2index(coord["xz"])
index["yz"] = self.coordinate2index(coord["yz"])
net = self.fc_pos(p)
net = self.blocks[0](net)
for block in self.blocks[1:]:
pooled = self.pool_local(coord, index, net)
net = torch.cat([net, pooled], dim=2)
net = block(net)
c = self.fc_c(net)
features = torch.stack([
self.generate_plane_features(index["xy"], c),
self.generate_plane_features(index["xz"], c),
self.generate_plane_features(index["yz"], c)
], dim=1)
return features
tgs/models/pointclouds/simplepoint.py
+110
View File
@@ -0,0 +1,110 @@
from dataclasses import dataclass, field
import torch
from einops import rearrange
import tgs
from tgs.utils.base import BaseModule
from tgs.utils.typing import *
class SimplePointGenerator(BaseModule):
@dataclass
class Config(BaseModule.Config):
camera_embedder_cls: str = ""
camera_embedder: dict = field(default_factory=dict)
image_tokenizer_cls: str = ""
image_tokenizer: dict = field(default_factory=dict)
tokenizer_cls: str = ""
tokenizer: dict = field(default_factory=dict)
backbone_cls: str = ""
backbone: dict = field(default_factory=dict)
post_processor_cls: str = ""
post_processor: dict = field(default_factory=dict)
pointcloud_upsampling_cls: str = ""
pointcloud_upsampling: dict = field(default_factory=dict)
flip_c2w_cond: bool = True
cfg: Config
def configure(self) -> None:
super().configure()
self.image_tokenizer = tgs.find(self.cfg.image_tokenizer_cls)(
self.cfg.image_tokenizer
)
assert self.cfg.camera_embedder_cls == 'tgs.models.networks.MLP'
weights = self.cfg.camera_embedder.pop("weights") if "weights" in self.cfg.camera_embedder else None
self.camera_embedder = tgs.find(self.cfg.camera_embedder_cls)(**self.cfg.camera_embedder)
if weights:
from tgs.utils.misc import load_module_weights
weights_path, module_name = weights.split(":")
state_dict = load_module_weights(
weights_path, module_name=module_name, map_location="cpu"
)
self.camera_embedder.load_state_dict(state_dict)
self.tokenizer = tgs.find(self.cfg.tokenizer_cls)(self.cfg.tokenizer)
self.backbone = tgs.find(self.cfg.backbone_cls)(self.cfg.backbone)
self.post_processor = tgs.find(self.cfg.post_processor_cls)(
self.cfg.post_processor
)
self.pointcloud_upsampling = tgs.find(self.cfg.pointcloud_upsampling_cls)(self.cfg.pointcloud_upsampling)
def forward(self, batch, encoder_hidden_states=None, **kwargs):
batch_size, n_input_views = batch["rgb_cond"].shape[:2]
if encoder_hidden_states is None:
# Camera modulation
c2w_cond = batch["c2w_cond"].clone()
if self.cfg.flip_c2w_cond:
c2w_cond[..., :3, 1:3] *= -1
camera_extri = c2w_cond.view(*c2w_cond.shape[:-2], -1)
camera_intri = batch["intrinsic_normed_cond"].view(
*batch["intrinsic_normed_cond"].shape[:-2], -1)
camera_feats = torch.cat([camera_intri, camera_extri], dim=-1)
# camera_feats = rearrange(camera_feats, 'B Nv C -> (B Nv) C')
camera_feats = self.camera_embedder(camera_feats)
encoder_hidden_states: Float[Tensor, "B Cit Nit"] = self.image_tokenizer(
rearrange(batch["rgb_cond"], 'B Nv H W C -> B Nv C H W'),
modulation_cond=camera_feats,
)
encoder_hidden_states = rearrange(
encoder_hidden_states, 'B Nv C Nt -> B (Nv Nt) C', Nv=n_input_views)
tokens: Float[Tensor, "B Ct Nt"] = self.tokenizer(batch_size)
tokens = self.backbone(
tokens,
encoder_hidden_states=encoder_hidden_states,
modulation_cond=None,
)
pointclouds = self.post_processor(self.tokenizer.detokenize(tokens))
upsampling_input = {
"input_image_tokens": encoder_hidden_states.permute(0, 2, 1),
"input_image_tokens_global": encoder_hidden_states[:, :1],
"c2w_cond": c2w_cond,
"rgb_cond": batch["rgb_cond"],
"intrinsic_cond": batch["intrinsic_cond"],
"intrinsic_normed_cond": batch["intrinsic_normed_cond"],
"points": pointclouds.float()
}
up_results = self.pointcloud_upsampling(upsampling_input)
up_results.insert(0, pointclouds)
pointclouds = up_results[-1]
out = {
"points": pointclouds,
"up_results": up_results
}
return out
tgs/models/renderer.py
+427
View File
@@ -0,0 +1,427 @@
from dataclasses import dataclass, field
from collections import defaultdict
from diff_gaussian_rasterization import GaussianRasterizationSettings, GaussianRasterizer
from plyfile import PlyData, PlyElement
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
import math
from ..utils.typing import *
from ..utils.base import BaseModule
from ..utils.ops import trunc_exp
from ..utils.ops import scale_tensor
from .networks import MLP
from einops import rearrange, reduce
inverse_sigmoid = lambda x: np.log(x / (1 - x))
def getWorld2View2(R, t, translate=np.array([.0, .0, .0]), scale=1.0):
Rt = np.zeros((4, 4))
Rt[:3, :3] = R.transpose()
Rt[:3, 3] = t
Rt[3, 3] = 1.0
C2W = np.linalg.inv(Rt)
cam_center = C2W[:3, 3]
cam_center = (cam_center + translate) * scale
C2W[:3, 3] = cam_center
Rt = np.linalg.inv(C2W)
return np.float32(Rt)
def getProjectionMatrix(znear, zfar, fovX, fovY):
tanHalfFovY = math.tan((fovY / 2))
tanHalfFovX = math.tan((fovX / 2))
top = tanHalfFovY * znear
bottom = -top
right = tanHalfFovX * znear
left = -right
P = torch.zeros(4, 4)
z_sign = 1.0
P[0, 0] = 2.0 * znear / (right - left)
P[1, 1] = 2.0 * znear / (top - bottom)
P[0, 2] = (right + left) / (right - left)
P[1, 2] = (top + bottom) / (top - bottom)
P[3, 2] = z_sign
P[2, 2] = z_sign * zfar / (zfar - znear)
P[2, 3] = -(zfar * znear) / (zfar - znear)
return P
def intrinsic_to_fov(intrinsic, w, h):
fx, fy = intrinsic[0, 0], intrinsic[1, 1]
fov_x = 2 * torch.arctan2(w, 2 * fx)
fov_y = 2 * torch.arctan2(h, 2 * fy)
return fov_x, fov_y
class Camera:
def __init__(self, w2c, intrinsic, FoVx, FoVy, height, width, trans=np.array([0.0, 0.0, 0.0]), scale=1.0) -> None:
self.FoVx = FoVx
self.FoVy = FoVy
self.height = height
self.width = width
self.world_view_transform = w2c.transpose(0, 1)
self.zfar = 100.0
self.znear = 0.01
self.trans = trans
self.scale = scale
self.projection_matrix = getProjectionMatrix(znear=self.znear, zfar=self.zfar, fovX=self.FoVx, fovY=self.FoVy).transpose(0,1).to(w2c.device)
self.full_proj_transform = (self.world_view_transform.unsqueeze(0).bmm(self.projection_matrix.unsqueeze(0))).squeeze(0)
self.camera_center = self.world_view_transform.inverse()[3, :3]
@staticmethod
def from_c2w(c2w, intrinsic, height, width):
w2c = torch.inverse(c2w)
FoVx, FoVy = intrinsic_to_fov(intrinsic, w=torch.tensor(width, device=w2c.device), h=torch.tensor(height, device=w2c.device))
return Camera(w2c=w2c, intrinsic=intrinsic, FoVx=FoVx, FoVy=FoVy, height=height, width=width)
class GaussianModel(NamedTuple):
xyz: Tensor
opacity: Tensor
rotation: Tensor
scaling: Tensor
shs: Tensor
def construct_list_of_attributes(self):
l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
features_dc = self.shs[:, :1]
features_rest = self.shs[:, 1:]
for i in range(features_dc.shape[1]*features_dc.shape[2]):
l.append('f_dc_{}'.format(i))
for i in range(features_rest.shape[1]*features_rest.shape[2]):
l.append('f_rest_{}'.format(i))
l.append('opacity')
for i in range(self.scaling.shape[1]):
l.append('scale_{}'.format(i))
for i in range(self.rotation.shape[1]):
l.append('rot_{}'.format(i))
return l
def to_ply(self):
xyz = self.xyz.detach().cpu().numpy()
normals = np.zeros_like(xyz)
features_dc = self.shs[:, :1]
features_rest = self.shs[:, 1:]
f_dc = features_dc.detach().flatten(start_dim=1).contiguous().cpu().numpy()
f_rest = features_rest.detach().flatten(start_dim=1).contiguous().cpu().numpy()
opacities = inverse_sigmoid(torch.clamp(self.opacity, 1e-3, 1 - 1e-3).detach().cpu().numpy())
scale = np.log(self.scaling.detach().cpu().numpy())
rotation = self.rotation.detach().cpu().numpy()
dtype_full = [(attribute, 'f4') for attribute in self.construct_list_of_attributes()]
elements = np.empty(xyz.shape[0], dtype=dtype_full)
attributes = np.concatenate((xyz, normals, f_dc, f_rest, opacities, scale, rotation), axis=1)
elements[:] = list(map(tuple, attributes))
el = PlyElement.describe(elements, 'vertex')
return PlyData([el])
class GSLayer(BaseModule):
@dataclass
class Config(BaseModule.Config):
in_channels: int = 128
feature_channels: dict = field(default_factory=dict)
xyz_offset: bool = True
restrict_offset: bool = False
use_rgb: bool = False
clip_scaling: Optional[float] = None
init_scaling: float = -5.0
init_density: float = 0.1
cfg: Config
def configure(self, *args, **kwargs) -> None:
self.out_layers = nn.ModuleList()
for key, out_ch in self.cfg.feature_channels.items():
if key == "shs" and self.cfg.use_rgb:
out_ch = 3
layer = nn.Linear(self.cfg.in_channels, out_ch)
# initialize
if not (key == "shs" and self.cfg.use_rgb):
nn.init.constant_(layer.weight, 0)
nn.init.constant_(layer.bias, 0)
if key == "scaling":
nn.init.constant_(layer.bias, self.cfg.init_scaling)
elif key == "rotation":
nn.init.constant_(layer.bias, 0)
nn.init.constant_(layer.bias[0], 1.0)
elif key == "opacity":
nn.init.constant_(layer.bias, inverse_sigmoid(self.cfg.init_density))
self.out_layers.append(layer)
def forward(self, x, pts):
ret = {}
for k, layer in zip(self.cfg.feature_channels.keys(), self.out_layers):
v = layer(x)
if k == "rotation":
v = torch.nn.functional.normalize(v)
elif k == "scaling":
v = trunc_exp(v)
if self.cfg.clip_scaling is not None:
v = torch.clamp(v, min=0, max=self.cfg.clip_scaling)
elif k == "opacity":
v = torch.sigmoid(v)
elif k == "shs":
if self.cfg.use_rgb:
v = torch.sigmoid(v)
v = torch.reshape(v, (v.shape[0], -1, 3))
elif k == "xyz":
if self.cfg.restrict_offset:
max_step = 1.2 / 32
v = (torch.sigmoid(v) - 0.5) * max_step
v = v + pts if self.cfg.xyz_offset else pts
ret[k] = v
return GaussianModel(**ret)
class GS3DRenderer(BaseModule):
@dataclass
class Config(BaseModule.Config):
mlp_network_config: Optional[dict] = None
gs_out: dict = field(default_factory=dict)
sh_degree: int = 3
scaling_modifier: float = 1.0
random_background: bool = False
radius: float = 1.0
feature_reduction: str = "concat"
projection_feature_dim: int = 773
background_color: Tuple[float, float, float] = field(
default_factory=lambda: (1.0, 1.0, 1.0)
)
cfg: Config
def configure(self, *args, **kwargs) -> None:
if self.cfg.feature_reduction == "mean":
mlp_in = 80
elif self.cfg.feature_reduction == "concat":
mlp_in = 80 * 3
else:
raise NotImplementedError
mlp_in = mlp_in + self.cfg.projection_feature_dim
if self.cfg.mlp_network_config is not None:
self.mlp_net = MLP(mlp_in, self.cfg.gs_out.in_channels, **self.cfg.mlp_network_config)
else:
self.cfg.gs_out.in_channels = mlp_in
self.gs_net = GSLayer(self.cfg.gs_out)
def forward_gs(self, x, p):
if self.cfg.mlp_network_config is not None:
x = self.mlp_net(x)
return self.gs_net(x, p)
def forward_single_view(self,
gs: GaussianModel,
viewpoint_camera: Camera,
background_color: Optional[Float[Tensor, "3"]],
ret_mask: bool = True,
):
# Create zero tensor. We will use it to make pytorch return gradients of the 2D (screen-space) means
screenspace_points = torch.zeros_like(gs.xyz, dtype=gs.xyz.dtype, requires_grad=True, device=self.device) + 0
try:
screenspace_points.retain_grad()
except:
pass
bg_color = background_color
# Set up rasterization configuration
tanfovx = math.tan(viewpoint_camera.FoVx * 0.5)
tanfovy = math.tan(viewpoint_camera.FoVy * 0.5)
raster_settings = GaussianRasterizationSettings(
image_height=int(viewpoint_camera.height),
image_width=int(viewpoint_camera.width),
tanfovx=tanfovx,
tanfovy=tanfovy,
bg=bg_color,
scale_modifier=self.cfg.scaling_modifier,
viewmatrix=viewpoint_camera.world_view_transform,
projmatrix=viewpoint_camera.full_proj_transform.float(),
sh_degree=self.cfg.sh_degree,
campos=viewpoint_camera.camera_center,
prefiltered=False,
debug=False
)
rasterizer = GaussianRasterizer(raster_settings=raster_settings)
means3D = gs.xyz
means2D = screenspace_points
opacity = gs.opacity
# If precomputed 3d covariance is provided, use it. If not, then it will be computed from
# scaling / rotation by the rasterizer.
scales = None
rotations = None
cov3D_precomp = None
scales = gs.scaling
rotations = gs.rotation
# If precomputed colors are provided, use them. Otherwise, if it is desired to precompute colors
# from SHs in Python, do it. If not, then SH -> RGB conversion will be done by rasterizer.
shs = None
colors_precomp = None
if self.gs_net.cfg.use_rgb:
colors_precomp = gs.shs.squeeze(1)
else:
shs = gs.shs
# Rasterize visible Gaussians to image, obtain their radii (on screen).
with torch.autocast(device_type=self.device.type, dtype=torch.float32):
rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
means3D = means3D,
means2D = means2D,
shs = shs,
colors_precomp = colors_precomp,
opacities = opacity,
scales = scales,
rotations = rotations,
cov3D_precomp = cov3D_precomp)
ret = {
"comp_rgb": rendered_image.permute(1, 2, 0),
"comp_rgb_bg": bg_color
}
if ret_mask:
mask_bg_color = torch.zeros(3, dtype=torch.float32, device=self.device)
raster_settings = GaussianRasterizationSettings(
image_height=int(viewpoint_camera.height),
image_width=int(viewpoint_camera.width),
tanfovx=tanfovx,
tanfovy=tanfovy,
bg=mask_bg_color,
scale_modifier=self.cfg.scaling_modifier,
viewmatrix=viewpoint_camera.world_view_transform,
projmatrix=viewpoint_camera.full_proj_transform.float(),
sh_degree=0,
campos=viewpoint_camera.camera_center,
prefiltered=False,
debug=False
)
rasterizer = GaussianRasterizer(raster_settings=raster_settings)
with torch.autocast(device_type=self.device.type, dtype=torch.float32):
rendered_mask, radii, rendered_depth, rendered_alpha = rasterizer(
means3D = means3D,
means2D = means2D,
# shs = ,
colors_precomp = torch.ones_like(means3D),
opacities = opacity,
scales = scales,
rotations = rotations,
cov3D_precomp = cov3D_precomp)
ret["comp_mask"] = rendered_mask.permute(1, 2, 0)
return ret
def query_triplane(
self,
positions: Float[Tensor, "*B N 3"],
triplanes: Float[Tensor, "*B 3 Cp Hp Wp"],
) -> Dict[str, Tensor]:
batched = positions.ndim == 3
if not batched:
# no batch dimension
triplanes = triplanes[None, ...]
positions = positions[None, ...]
positions = scale_tensor(positions, (-self.cfg.radius, self.cfg.radius), (-1, 1))
indices2D: Float[Tensor, "B 3 N 2"] = torch.stack(
(positions[..., [0, 1]], positions[..., [0, 2]], positions[..., [1, 2]]),
dim=-3,
)
out: Float[Tensor, "B3 Cp 1 N"] = F.grid_sample(
rearrange(triplanes, "B Np Cp Hp Wp -> (B Np) Cp Hp Wp", Np=3),
rearrange(indices2D, "B Np N Nd -> (B Np) () N Nd", Np=3),
align_corners=False,
mode="bilinear",
)
if self.cfg.feature_reduction == "concat":
out = rearrange(out, "(B Np) Cp () N -> B N (Np Cp)", Np=3)
elif self.cfg.feature_reduction == "mean":
out = reduce(out, "(B Np) Cp () N -> B N Cp", Np=3, reduction="mean")
else:
raise NotImplementedError
if not batched:
out = out.squeeze(0)
return out
def forward_single_batch(
self,
gs_hidden_features: Float[Tensor, "Np Cp"],
query_points: Float[Tensor, "Np 3"],
c2ws: Float[Tensor, "Nv 4 4"],
intrinsics: Float[Tensor, "Nv 4 4"],
height: int,
width: int,
background_color: Optional[Float[Tensor, "3"]],
):
gs: GaussianModel = self.forward_gs(gs_hidden_features, query_points)
out_list = []
for c2w, intrinsic in zip(c2ws, intrinsics):
out_list.append(self.forward_single_view(
gs,
Camera.from_c2w(c2w, intrinsic, height, width),
background_color
))
out = defaultdict(list)
for out_ in out_list:
for k, v in out_.items():
out[k].append(v)
out = {k: torch.stack(v, dim=0) for k, v in out.items()}
out["3dgs"] = gs
return out
def forward(self,
gs_hidden_features: Float[Tensor, "B Np Cp"],
query_points: Float[Tensor, "B Np 3"],
c2w: Float[Tensor, "B Nv 4 4"],
intrinsic: Float[Tensor, "B Nv 4 4"],
height,
width,
additional_features: Optional[Float[Tensor, "B C H W"]] = None,
background_color: Optional[Float[Tensor, "B 3"]] = None,
**kwargs):
batch_size = gs_hidden_features.shape[0]
out_list = []
gs_hidden_features = self.query_triplane(query_points, gs_hidden_features)
if additional_features is not None:
gs_hidden_features = torch.cat([gs_hidden_features, additional_features], dim=-1)
for b in range(batch_size):
out_list.append(self.forward_single_batch(
gs_hidden_features[b],
query_points[b],
c2w[b],
intrinsic[b],
height, width,
background_color[b] if background_color is not None else None))
out = defaultdict(list)
for out_ in out_list:
for k, v in out_.items():
out[k].append(v)
for k, v in out.items():
if isinstance(v[0], torch.Tensor):
out[k] = torch.stack(v, dim=0)
else:
out[k] = v
return out
tgs/models/snowflake/LICENSE
+21
View File
@@ -0,0 +1,21 @@
MIT License
Copyright (c) 2021 AllenXiang
Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:
The above copyright notice and this permission notice shall be included in all
copies or substantial portions of the Software.
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
SOFTWARE.
tgs/models/snowflake/SPD.py
+68
View File
@@ -0,0 +1,68 @@
# -*- coding: utf-8 -*-
# @Author: Peng Xiang
import torch
import torch.nn as nn
from .utils import MLP_Res, MLP_CONV
from .skip_transformer import SkipTransformer
class SPD(nn.Module):
def __init__(self, dim_feat=512, up_factor=2, i=0, radius=1, bounding=True, global_feat=True):
"""Snowflake Point Deconvolution"""
super(SPD, self).__init__()
self.i = i
self.up_factor = up_factor
self.bounding = bounding
self.radius = radius
self.global_feat = global_feat
self.ps_dim = 32 if global_feat else 64
self.mlp_1 = MLP_CONV(in_channel=3, layer_dims=[64, 128])
self.mlp_2 = MLP_CONV(in_channel=128 * 2 + dim_feat if self.global_feat else 128, layer_dims=[256, 128])
self.skip_transformer = SkipTransformer(in_channel=128, dim=64)
self.mlp_ps = MLP_CONV(in_channel=128, layer_dims=[64, self.ps_dim])
self.ps = nn.ConvTranspose1d(self.ps_dim, 128, up_factor, up_factor, bias=False) # point-wise splitting
self.up_sampler = nn.Upsample(scale_factor=up_factor)
self.mlp_delta_feature = MLP_Res(in_dim=256, hidden_dim=128, out_dim=128)
self.mlp_delta = MLP_CONV(in_channel=128, layer_dims=[64, 3])
def forward(self, pcd_prev, feat_global=None, K_prev=None):
"""
Args:
pcd_prev: Tensor, (B, 3, N_prev)
feat_global: Tensor, (B, dim_feat, 1)
K_prev: Tensor, (B, 128, N_prev)
Returns:
pcd_child: Tensor, up sampled point cloud, (B, 3, N_prev * up_factor)
K_curr: Tensor, displacement feature of current step, (B, 128, N_prev * up_factor)
"""
b, _, n_prev = pcd_prev.shape
feat_1 = self.mlp_1(pcd_prev)
feat_1 = torch.cat([feat_1,
torch.max(feat_1, 2, keepdim=True)[0].repeat((1, 1, feat_1.size(2))),
feat_global.repeat(1, 1, feat_1.size(2))], 1) if self.global_feat else feat_1
Q = self.mlp_2(feat_1)
H = self.skip_transformer(pcd_prev, K_prev if K_prev is not None else Q, Q)
feat_child = self.mlp_ps(H)
feat_child = self.ps(feat_child) # (B, 128, N_prev * up_factor)
H_up = self.up_sampler(H)
K_curr = self.mlp_delta_feature(torch.cat([feat_child, H_up], 1))
delta = self.mlp_delta(torch.relu(K_curr))
if self.bounding:
delta = torch.tanh(delta) / self.radius**self.i # (B, 3, N_prev * up_factor)
pcd_child = self.up_sampler(pcd_prev)
pcd_child = pcd_child + delta
return pcd_child, K_curr
tgs/models/snowflake/SPD_crossattn.py
+81
View File
@@ -0,0 +1,81 @@
# -*- coding: utf-8 -*-
# @Author: Peng Xiang
import torch
import torch.nn as nn
from .utils import MLP_Res, MLP_CONV
from .skip_transformer import SkipTransformer
from .attention import ResidualTransformerBlock
class SPD_crossattn(nn.Module):
def __init__(self, dim_feat=512, up_factor=2, i=0, radius=1, bounding=True, global_feat=True):
"""Snowflake Point Deconvolution"""
super().__init__()
self.i = i
self.up_factor = up_factor
self.bounding = bounding
self.radius = radius
self.global_feat = global_feat
self.ps_dim = 32 if global_feat else 64
self.mlp_1 = MLP_CONV(in_channel=3, layer_dims=[64, 128])
self.pcd_image_attn = ResidualTransformerBlock(
device=torch.device('cuda'),
dtype=torch.float32,
n_data=128,
width=128,
heads=8,
init_scale=1.0,
)
self.mlp_2 = MLP_CONV(in_channel=128 * 2 + dim_feat if self.global_feat else 128, layer_dims=[256, 128])
self.skip_transformer = SkipTransformer(in_channel=128, dim=64)
self.mlp_ps = MLP_CONV(in_channel=128, layer_dims=[64, self.ps_dim])
self.ps = nn.ConvTranspose1d(self.ps_dim, 128, up_factor, up_factor, bias=False) # point-wise splitting
self.up_sampler = nn.Upsample(scale_factor=up_factor)
self.mlp_delta_feature = MLP_Res(in_dim=256, hidden_dim=128, out_dim=128)
self.mlp_delta = MLP_CONV(in_channel=128, layer_dims=[64, 3])
def forward(self, pcd_prev, feat_global=None, K_prev=None):
"""
Args:
pcd_prev: Tensor, (B, 3, N_prev)
feat_global: Tensor, (B, dim_feat, 1)
K_prev: Tensor, (B, 128, N_prev)
Returns:
pcd_child: Tensor, up sampled point cloud, (B, 3, N_prev * up_factor)
K_curr: Tensor, displacement feature of current step, (B, 128, N_prev * up_factor)
"""
b, _, n_prev = pcd_prev.shape
feat_1 = self.mlp_1(pcd_prev)
# feat_1 = torch.cat([feat_1,
# torch.max(feat_1, 2, keepdim=True)[0].repeat((1, 1, feat_1.size(2))),
# feat_global.repeat(1, 1, feat_1.size(2))], 1) if self.global_feat else feat_1
feat_1 = torch.permute(feat_1, (0, 2, 1))
feat_global = torch.permute(feat_global, (0, 2, 1))
feat_1 = self.pcd_image_attn(feat_1, feat_global)
Q = torch.permute(feat_1, (0, 2, 1))
# Q = self.mlp_2(feat_1)
H = self.skip_transformer(pcd_prev, K_prev if K_prev is not None else Q, Q)
feat_child = self.mlp_ps(H)
feat_child = self.ps(feat_child) # (B, 128, N_prev * up_factor)
H_up = self.up_sampler(H)
K_curr = self.mlp_delta_feature(torch.cat([feat_child, H_up], 1))
delta = self.mlp_delta(torch.relu(K_curr))
if self.bounding:
delta = torch.tanh(delta) / self.radius**self.i # (B, 3, N_prev * up_factor)
pcd_child = self.up_sampler(pcd_prev)
pcd_child = pcd_child + delta
return pcd_child, K_curr
tgs/models/snowflake/SPD_pp.py
+71
View File
@@ -0,0 +1,71 @@
import torch
import torch.nn as nn
from .utils import MLP_Res, MLP_CONV
from .skip_transformer import SkipTransformer
class SPD_pp(nn.Module):
def __init__(self, dim_feat=512, up_factor=2, i=0, radius=1, bounding=True, global_feat=True):
"""Snowflake Point Deconvolution"""
super(SPD_pp, self).__init__()
self.i = i
self.up_factor = up_factor
self.bounding = bounding
self.radius = radius
self.global_feat = global_feat
self.ps_dim = 32 if global_feat else 64
self.mlp_1 = MLP_CONV(in_channel=3, layer_dims=[64, 128])
self.mlp_2 = MLP_CONV(
in_channel=128 * 2 + dim_feat if self.global_feat else 128, layer_dims=[256, 128])
self.skip_transformer = SkipTransformer(in_channel=128, dim=64)
self.mlp_ps = MLP_CONV(in_channel=128, layer_dims=[64, self.ps_dim])
self.ps = nn.ConvTranspose1d(
self.ps_dim, 128, up_factor, up_factor, bias=False) # point-wise splitting
self.up_sampler = nn.Upsample(scale_factor=up_factor)
self.mlp_delta_feature = MLP_Res(
in_dim=256, hidden_dim=128, out_dim=128)
self.mlp_delta = MLP_CONV(in_channel=128, layer_dims=[64, 3])
def forward(self, pcd_prev, feat_cond=None, K_prev=None):
"""
Args:
pcd_prev: Tensor, (B, 3, N_prev)
feat_cond: Tensor, (B, dim_feat, N_prev)
K_prev: Tensor, (B, 128, N_prev)
Returns:
pcd_child: Tensor, up sampled point cloud, (B, 3, N_prev * up_factor)
K_curr: Tensor, displacement feature of current step, (B, 128, N_prev * up_factor)
"""
b, _, n_prev = pcd_prev.shape
feat_1 = self.mlp_1(pcd_prev)
feat_1 = torch.cat([feat_1,
torch.max(feat_1, 2, keepdim=True)[
0].repeat((1, 1, feat_1.size(2))),
feat_cond], 1) if self.global_feat else feat_1
Q = self.mlp_2(feat_1)
H = self.skip_transformer(
pcd_prev, K_prev if K_prev is not None else Q, Q)
feat_child = self.mlp_ps(H)
feat_child = self.ps(feat_child) # (B, 128, N_prev * up_factor)
H_up = self.up_sampler(H)
K_curr = self.mlp_delta_feature(torch.cat([feat_child, H_up], 1))
delta = self.mlp_delta(torch.relu(K_curr))
if self.bounding:
# (B, 3, N_prev * up_factor)
delta = torch.tanh(delta) / self.radius**self.i
pcd_child = self.up_sampler(pcd_prev)
pcd_child = pcd_child + delta
return pcd_child, K_curr
tgs/models/snowflake/attention.py
+239
View File
@@ -0,0 +1,239 @@
import torch
import torch.nn as nn
import math
import math
from typing import Optional
from typing import Callable, Iterable, Sequence, Union
import torch
def checkpoint(
func: Callable[..., Union[torch.Tensor, Sequence[torch.Tensor]]],
inputs: Sequence[torch.Tensor],
params: Iterable[torch.Tensor],
flag: bool,
):
"""
Evaluate a function without caching intermediate activations, allowing for
reduced memory at the expense of extra compute in the backward pass.
:param func: the function to evaluate.
:param inputs: the argument sequence to pass to `func`.
:param params: a sequence of parameters `func` depends on but does not
explicitly take as arguments.
:param flag: if False, disable gradient checkpointing.
"""
if flag:
args = tuple(inputs) + tuple(params)
return CheckpointFunction.apply(func, len(inputs), *args)
else:
return func(*inputs)
class CheckpointFunction(torch.autograd.Function):
@staticmethod
def forward(ctx, run_function, length, *args):
ctx.run_function = run_function
ctx.input_tensors = list(args[:length])
ctx.input_params = list(args[length:])
with torch.no_grad():
output_tensors = ctx.run_function(*ctx.input_tensors)
return output_tensors
@staticmethod
def backward(ctx, *output_grads):
ctx.input_tensors = [x.detach().requires_grad_(True)
for x in ctx.input_tensors]
with torch.enable_grad():
# Fixes a bug where the first op in run_function modifies the
# Tensor storage in place, which is not allowed for detach()'d
# Tensors.
shallow_copies = [x.view_as(x) for x in ctx.input_tensors]
output_tensors = ctx.run_function(*shallow_copies)
input_grads = torch.autograd.grad(
output_tensors,
ctx.input_tensors + ctx.input_params,
output_grads,
allow_unused=True,
)
del ctx.input_tensors
del ctx.input_params
del output_tensors
return (None, None) + input_grads
def init_linear(l, stddev):
nn.init.normal_(l.weight, std=stddev)
if l.bias is not None:
nn.init.constant_(l.bias, 0.0)
class MLP(nn.Module):
def __init__(self, *, device: torch.device, dtype: torch.dtype, width: int, init_scale: float):
super().__init__()
self.width = width
self.c_fc = nn.Linear(width, width * 4, device=device, dtype=dtype)
self.c_proj = nn.Linear(width * 4, width, device=device, dtype=dtype)
self.gelu = nn.GELU()
init_linear(self.c_fc, init_scale)
init_linear(self.c_proj, init_scale)
def forward(self, x):
return self.c_proj(self.gelu(self.c_fc(x)))
class QKVMultiheadCrossAttention(nn.Module):
def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int, n_data: int):
super().__init__()
self.device = device
self.dtype = dtype
self.heads = heads
self.n_data = n_data
def forward(self, q, kv):
_, n_ctx, _ = q.shape
bs, n_data, width = kv.shape
attn_ch = width // self.heads // 2
scale = 1 / math.sqrt(math.sqrt(attn_ch))
q = q.view(bs, n_ctx, self.heads, -1)
kv = kv.view(bs, n_data, self.heads, -1)
k, v = torch.split(kv, attn_ch, dim=-1)
weight = torch.einsum(
"bthc,bshc->bhts", q * scale, k * scale
) # More stable with f16 than dividing afterwards
wdtype = weight.dtype
weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
return torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)
class QKVMultiheadAttention(nn.Module):
def __init__(self, *, device: torch.device, dtype: torch.dtype, heads: int, n_ctx: int):
super().__init__()
self.device = device
self.dtype = dtype
self.heads = heads
self.n_ctx = n_ctx
def forward(self, qkv):
bs, n_ctx, width = qkv.shape
attn_ch = width // self.heads // 3
scale = 1 / math.sqrt(math.sqrt(attn_ch))
qkv = qkv.view(bs, n_ctx, self.heads, -1)
q, k, v = torch.split(qkv, attn_ch, dim=-1)
weight = torch.einsum(
"bthc,bshc->bhts", q * scale, k * scale
) # More stable with f16 than dividing afterwards
wdtype = weight.dtype
weight = torch.softmax(weight.float(), dim=-1).type(wdtype)
return torch.einsum("bhts,bshc->bthc", weight, v).reshape(bs, n_ctx, -1)
class MultiheadCrossAttention(nn.Module):
def __init__(
self,
*,
device: torch.device,
dtype: torch.dtype,
n_data: int,
width: int,
heads: int,
init_scale: float,
data_width: Optional[int] = None,
):
super().__init__()
self.n_data = n_data
self.width = width
self.heads = heads
self.data_width = width if data_width is None else data_width
self.c_q = nn.Linear(width, width, device=device, dtype=dtype)
self.c_kv = nn.Linear(self.data_width, width * 2,
device=device, dtype=dtype)
self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
self.attention = QKVMultiheadCrossAttention(
device=device, dtype=dtype, heads=heads, n_data=n_data
)
init_linear(self.c_q, init_scale)
init_linear(self.c_kv, init_scale)
init_linear(self.c_proj, init_scale)
def forward(self, x, data):
x = self.c_q(x)
data = self.c_kv(data)
x = checkpoint(self.attention, (x, data), (), True)
x = self.c_proj(x)
return x
class MultiheadAttention(nn.Module):
def __init__(
self,
*,
device: torch.device,
dtype: torch.dtype,
n_ctx: int,
width: int,
heads: int,
init_scale: float,
):
super().__init__()
self.n_ctx = n_ctx
self.width = width
self.heads = heads
self.c_qkv = nn.Linear(width, width * 3, device=device, dtype=dtype)
self.c_proj = nn.Linear(width, width, device=device, dtype=dtype)
self.attention = QKVMultiheadAttention(device=device, dtype=dtype, heads=heads, n_ctx=n_ctx)
init_linear(self.c_qkv, init_scale)
init_linear(self.c_proj, init_scale)
def forward(self, x):
x = self.c_qkv(x)
x = checkpoint(self.attention, (x,), (), True)
x = self.c_proj(x)
return x
class ResidualTransformerBlock(nn.Module):
def __init__(
self,
*,
device: torch.device,
dtype: torch.dtype,
n_data: int,
width: int,
heads: int,
data_width: Optional[int] = None,
init_scale: float = 1.0,
):
super().__init__()
if data_width is None:
data_width = width
self.attn_cross = MultiheadCrossAttention(
device=device,
dtype=dtype,
n_data=n_data,
width=width,
heads=heads,
data_width=data_width,
init_scale=init_scale,
)
self.attn_self = MultiheadAttention(
device=device,
dtype=dtype,
n_ctx=n_data,
width=width,
heads=heads,
init_scale=init_scale,
)
self.ln_1 = nn.LayerNorm(width, device=device, dtype=dtype)
self.ln_2 = nn.LayerNorm(data_width, device=device, dtype=dtype)
self.ln_3 = nn.LayerNorm(width, device=device, dtype=dtype)
self.mlp = MLP(device=device, dtype=dtype,
width=width, init_scale=init_scale)
self.ln_4 = nn.LayerNorm(width, device=device, dtype=dtype)
def forward(self, x: torch.Tensor, data: torch.Tensor):
x = x + self.attn_cross(self.ln_1(x), self.ln_2(data))
x = x + self.attn_self(self.ln_3(x))
x = x + self.mlp(self.ln_4(x))
return x
tgs/models/snowflake/model_spdpp.py
+188
View File
@@ -0,0 +1,188 @@
import torch
import torch.nn as nn
import torch.nn.functional as F
from tgs.utils.base import BaseModule
from tgs.utils.typing import *
from dataclasses import dataclass, field
from pytorch3d.renderer import (
AlphaCompositor,
NormWeightedCompositor,
PointsRasterizationSettings,
PointsRasterizer,
PointsRenderer)
from pytorch3d.renderer.cameras import CamerasBase
from pytorch3d.structures import Pointclouds
from pytorch3d.utils.camera_conversions import cameras_from_opencv_projection
from .utils import fps_subsample
from einops import rearrange
from .utils import MLP_CONV
from .SPD import SPD
from .SPD_crossattn import SPD_crossattn
from .SPD_pp import SPD_pp
SPD_BLOCK = {
'SPD': SPD,
'SPD_crossattn': SPD_crossattn,
'SPD_PP': SPD_pp,
}
def points_projection(points: Float[Tensor, "B Np 3"],
c2ws: Float[Tensor, "B 4 4"],
intrinsics: Float[Tensor, "B 3 3"],
local_features: Float[Tensor, "B C H W"],
raster_point_radius: float = 0.0075, # point size
raster_points_per_pixel: int = 1, # a single point per pixel, for now
bin_size: int = 0):
"""
points: (B, Np, 3)
"""
B, C, H, W = local_features.shape
device = local_features.device
raster_settings = PointsRasterizationSettings(
image_size=(H, W),
radius=raster_point_radius,
points_per_pixel=raster_points_per_pixel,
bin_size=bin_size,
)
Np = points.shape[1]
c2ws = c2ws.transpose(0, 1).flatten(0, 1)
intrinsics = intrinsics.transpose(0, 1).flatten(0, 1)
R = raster_settings.points_per_pixel
w2cs = torch.inverse(c2ws)
image_size = torch.as_tensor([H, W]).view(
1, 2).expand(w2cs.shape[0], -1).to(device)
cameras = cameras_from_opencv_projection(
w2cs[:, :3, :3], w2cs[:, :3, 3], intrinsics, image_size)
rasterize = PointsRasterizer(
cameras=cameras, raster_settings=raster_settings)
fragments = rasterize(Pointclouds(points))
fragments_idx: Tensor = fragments.idx.long()
visible_pixels = (fragments_idx > -1) # (B, H, W, R)
points_to_visible_pixels = fragments_idx[visible_pixels]
# Reshape local features to (B, H, W, R, C)
local_features = local_features.permute(
0, 2, 3, 1).unsqueeze(-2).expand(-1, -1, -1, R, -1) # (B, H, W, R, C)
# Get local features corresponding to visible points
local_features_proj = torch.zeros(B * Np, C, device=device)
local_features_proj[points_to_visible_pixels] = local_features[visible_pixels]
local_features_proj = local_features_proj.reshape(B, Np, C)
return local_features_proj
class Decoder(nn.Module):
def __init__(self, input_channels=1152, dim_feat=512, num_p0=512,
radius=1, bounding=True, up_factors=None,
SPD_type='SPD',
token_type='image_token'
):
super(Decoder, self).__init__()
# self.decoder_coarse = SeedGenerator(dim_feat=dim_feat, num_pc=num_p0)
if up_factors is None:
up_factors = [1]
else:
up_factors = up_factors
uppers = []
self.num_p0 = num_p0
self.mlp_feat_cond = MLP_CONV(in_channel=input_channels,
layer_dims=[dim_feat*2, dim_feat])
for i, factor in enumerate(up_factors):
uppers.append(
SPD_BLOCK[SPD_type](dim_feat=dim_feat, up_factor=factor,
i=i, bounding=bounding, radius=radius))
self.uppers = nn.ModuleList(uppers)
self.token_type = token_type
def calculate_pcl_token(self, pcl_token, up_factor):
up_token = F.interpolate(pcl_token, scale_factor=up_factor, mode='nearest')
return up_token
def calculate_image_token(self, pcd, input_image_tokens, batch):
"""
Args:
"""
batch_size, n_input_views = batch["rgb_cond"].shape[:2]
h_cond, w_cond = batch["rgb_cond"].shape[2:4]
input_image_tokens = rearrange(
input_image_tokens, '(B Nv) C Nt -> B (Nv Nt) C', Nv=n_input_views)
local_features = input_image_tokens[:, 1:].reshape(
batch_size, h_cond // 14, w_cond // 14, -1).permute(0, 3, 1, 2)
local_features = F.interpolate(local_features, size=(
h_cond, w_cond), mode='bilinear', align_corners=False)
batch['c2w_cond'][..., :3, 1:3] *= -1
local_features_proj = points_projection(
pcd,
batch['c2w_cond'],
batch['intrinsic_cond'],
local_features,
)
local_features_proj = local_features_proj.permute(0, 2, 1).contiguous()
return local_features_proj
def forward(self, x):
"""
Args:
points: Tensor, (b, num_p0, 3)
feat_cond: Tensor, (b, dim_feat) dinov2: 325x768
# partial_coarse: Tensor, (b, n_coarse, 3)
"""
points = x['points']
if self.token_type == 'pcl_token':
feat_cond = x['pcl_token']
elif self.token_type == 'image_token':
feat_cond = x['input_image_tokens']
feat_cond = self.mlp_feat_cond(feat_cond)
arr_pcd = []
feat_prev = None
pcd = torch.permute(points, (0, 2, 1)).contiguous()
pcl_up_scale = 1
for upper in self.uppers:
if self.token_type == 'pcl_token':
up_cond = self.calculate_pcl_token(
feat_cond, pcl_up_scale)
pcl_up_scale *= upper.up_factor
elif self.token_type == 'image_token':
up_cond = self.calculate_image_token(points, feat_cond, x)
pcd, feat_prev = upper(pcd, up_cond, feat_prev)
points = torch.permute(pcd, (0, 2, 1)).contiguous()
arr_pcd.append(points)
return arr_pcd
class SnowflakeModelSPDPP(BaseModule):
"""
apply PC^2 / PCL token to decoder
"""
@dataclass
class Config(BaseModule.Config):
input_channels: int = 1152
dim_feat: int = 128
num_p0: int = 512
radius: float = 1
bounding: bool = True
use_fps: bool = True
up_factors: List[int] = field(default_factory=lambda: [2, 2])
image_full_token_cond: bool = False
SPD_type: str = 'SPD_PP'
token_type: str = 'pcl_token'
cfg: Config
def configure(self) -> None:
super().configure()
self.decoder = Decoder(input_channels=self.cfg.input_channels,
dim_feat=self.cfg.dim_feat, num_p0=self.cfg.num_p0,
radius=self.cfg.radius, up_factors=self.cfg.up_factors, bounding=self.cfg.bounding,
SPD_type=self.cfg.SPD_type,
token_type=self.cfg.token_type
)
def forward(self, x):
results = self.decoder(x)
return results
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/__init__.py
+3
View File
@@ -0,0 +1,3 @@
import pointnet2_ops.pointnet2_modules
import pointnet2_ops.pointnet2_utils
from pointnet2_ops._version import __version__
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/ball_query.h
+5
View File
@@ -0,0 +1,5 @@
#pragma once
#include <torch/extension.h>
at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius,
const int nsample);
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/cuda_utils.h
+41
View File
@@ -0,0 +1,41 @@
#ifndef _CUDA_UTILS_H
#define _CUDA_UTILS_H
#include <ATen/ATen.h>
#include <ATen/cuda/CUDAContext.h>
#include <cmath>
#include <cuda.h>
#include <cuda_runtime.h>
#include <vector>
#define TOTAL_THREADS 512
inline int opt_n_threads(int work_size) {
const int pow_2 = std::log(static_cast<double>(work_size)) / std::log(2.0);
return max(min(1 << pow_2, TOTAL_THREADS), 1);
}
inline dim3 opt_block_config(int x, int y) {
const int x_threads = opt_n_threads(x);
const int y_threads =
max(min(opt_n_threads(y), TOTAL_THREADS / x_threads), 1);
dim3 block_config(x_threads, y_threads, 1);
return block_config;
}
#define CUDA_CHECK_ERRORS() \
do { \
cudaError_t err = cudaGetLastError(); \
if (cudaSuccess != err) { \
fprintf(stderr, "CUDA kernel failed : %s\n%s at L:%d in %s\n", \
cudaGetErrorString(err), __PRETTY_FUNCTION__, __LINE__, \
__FILE__); \
exit(-1); \
} \
} while (0)
#endif
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/group_points.h
+5
View File
@@ -0,0 +1,5 @@
#pragma once
#include <torch/extension.h>
at::Tensor group_points(at::Tensor points, at::Tensor idx);
at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n);
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/interpolate.h
+10
View File
@@ -0,0 +1,10 @@
#pragma once
#include <torch/extension.h>
#include <vector>
std::vector<at::Tensor> three_nn(at::Tensor unknowns, at::Tensor knows);
at::Tensor three_interpolate(at::Tensor points, at::Tensor idx,
at::Tensor weight);
at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx,
at::Tensor weight, const int m);
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/sampling.h
+6
View File
@@ -0,0 +1,6 @@
#pragma once
#include <torch/extension.h>
at::Tensor gather_points(at::Tensor points, at::Tensor idx);
at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx, const int n);
at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples);
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/utils.h
+25
View File
@@ -0,0 +1,25 @@
#pragma once
#include <ATen/cuda/CUDAContext.h>
#include <torch/extension.h>
#define CHECK_CUDA(x) \
do { \
AT_ASSERT(x.is_cuda(), #x " must be a CUDA tensor"); \
} while (0)
#define CHECK_CONTIGUOUS(x) \
do { \
AT_ASSERT(x.is_contiguous(), #x " must be a contiguous tensor"); \
} while (0)
#define CHECK_IS_INT(x) \
do { \
AT_ASSERT(x.scalar_type() == at::ScalarType::Int, \
#x " must be an int tensor"); \
} while (0)
#define CHECK_IS_FLOAT(x) \
do { \
AT_ASSERT(x.scalar_type() == at::ScalarType::Float, \
#x " must be a float tensor"); \
} while (0)
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query.cpp
+32
View File
@@ -0,0 +1,32 @@
#include "ball_query.h"
#include "utils.h"
void query_ball_point_kernel_wrapper(int b, int n, int m, float radius,
int nsample, const float *new_xyz,
const float *xyz, int *idx);
at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius,
const int nsample) {
CHECK_CONTIGUOUS(new_xyz);
CHECK_CONTIGUOUS(xyz);
CHECK_IS_FLOAT(new_xyz);
CHECK_IS_FLOAT(xyz);
if (new_xyz.is_cuda()) {
CHECK_CUDA(xyz);
}
at::Tensor idx =
torch::zeros({new_xyz.size(0), new_xyz.size(1), nsample},
at::device(new_xyz.device()).dtype(at::ScalarType::Int));
if (new_xyz.is_cuda()) {
query_ball_point_kernel_wrapper(xyz.size(0), xyz.size(1), new_xyz.size(1),
radius, nsample, new_xyz.data_ptr<float>(),
xyz.data_ptr<float>(), idx.data_ptr<int>());
} else {
AT_ASSERT(false, "CPU not supported");
}
return idx;
}
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query_gpu.cu
+54
View File
@@ -0,0 +1,54 @@
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include "cuda_utils.h"
// input: new_xyz(b, m, 3) xyz(b, n, 3)
// output: idx(b, m, nsample)
__global__ void query_ball_point_kernel(int b, int n, int m, float radius,
int nsample,
const float *__restrict__ new_xyz,
const float *__restrict__ xyz,
int *__restrict__ idx) {
int batch_index = blockIdx.x;
xyz += batch_index * n * 3;
new_xyz += batch_index * m * 3;
idx += m * nsample * batch_index;
int index = threadIdx.x;
int stride = blockDim.x;
float radius2 = radius * radius;
for (int j = index; j < m; j += stride) {
float new_x = new_xyz[j * 3 + 0];
float new_y = new_xyz[j * 3 + 1];
float new_z = new_xyz[j * 3 + 2];
for (int k = 0, cnt = 0; k < n && cnt < nsample; ++k) {
float x = xyz[k * 3 + 0];
float y = xyz[k * 3 + 1];
float z = xyz[k * 3 + 2];
float d2 = (new_x - x) * (new_x - x) + (new_y - y) * (new_y - y) +
(new_z - z) * (new_z - z);
if (d2 < radius2) {
if (cnt == 0) {
for (int l = 0; l < nsample; ++l) {
idx[j * nsample + l] = k;
}
}
idx[j * nsample + cnt] = k;
++cnt;
}
}
}
}
void query_ball_point_kernel_wrapper(int b, int n, int m, float radius,
int nsample, const float *new_xyz,
const float *xyz, int *idx) {
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
query_ball_point_kernel<<<b, opt_n_threads(m), 0, stream>>>(
b, n, m, radius, nsample, new_xyz, xyz, idx);
CUDA_CHECK_ERRORS();
}
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/bindings.cpp
+19
View File
@@ -0,0 +1,19 @@
#include "ball_query.h"
#include "group_points.h"
#include "interpolate.h"
#include "sampling.h"
PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
m.def("gather_points", &gather_points);
m.def("gather_points_grad", &gather_points_grad);
m.def("furthest_point_sampling", &furthest_point_sampling);
m.def("three_nn", &three_nn);
m.def("three_interpolate", &three_interpolate);
m.def("three_interpolate_grad", &three_interpolate_grad);
m.def("ball_query", &ball_query);
m.def("group_points", &group_points);
m.def("group_points_grad", &group_points_grad);
}
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points.cpp
+62
View File
@@ -0,0 +1,62 @@
#include "group_points.h"
#include "utils.h"
void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample,
const float *points, const int *idx,
float *out);
void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
int nsample, const float *grad_out,
const int *idx, float *grad_points);
at::Tensor group_points(at::Tensor points, at::Tensor idx) {
CHECK_CONTIGUOUS(points);
CHECK_CONTIGUOUS(idx);
CHECK_IS_FLOAT(points);
CHECK_IS_INT(idx);
if (points.is_cuda()) {
CHECK_CUDA(idx);
}
at::Tensor output =
torch::zeros({points.size(0), points.size(1), idx.size(1), idx.size(2)},
at::device(points.device()).dtype(at::ScalarType::Float));
if (points.is_cuda()) {
group_points_kernel_wrapper(points.size(0), points.size(1), points.size(2),
idx.size(1), idx.size(2),
points.data_ptr<float>(), idx.data_ptr<int>(),
output.data_ptr<float>());
} else {
AT_ASSERT(false, "CPU not supported");
}
return output;
}
at::Tensor group_points_grad(at::Tensor grad_out, at::Tensor idx, const int n) {
CHECK_CONTIGUOUS(grad_out);
CHECK_CONTIGUOUS(idx);
CHECK_IS_FLOAT(grad_out);
CHECK_IS_INT(idx);
if (grad_out.is_cuda()) {
CHECK_CUDA(idx);
}
at::Tensor output =
torch::zeros({grad_out.size(0), grad_out.size(1), n},
at::device(grad_out.device()).dtype(at::ScalarType::Float));
if (grad_out.is_cuda()) {
group_points_grad_kernel_wrapper(
grad_out.size(0), grad_out.size(1), n, idx.size(1), idx.size(2),
grad_out.data_ptr<float>(), idx.data_ptr<int>(),
output.data_ptr<float>());
} else {
AT_ASSERT(false, "CPU not supported");
}
return output;
}
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points_gpu.cu
+75
View File
@@ -0,0 +1,75 @@
#include <stdio.h>
#include <stdlib.h>
#include "cuda_utils.h"
// input: points(b, c, n) idx(b, npoints, nsample)
// output: out(b, c, npoints, nsample)
__global__ void group_points_kernel(int b, int c, int n, int npoints,
int nsample,
const float *__restrict__ points,
const int *__restrict__ idx,
float *__restrict__ out) {
int batch_index = blockIdx.x;
points += batch_index * n * c;
idx += batch_index * npoints * nsample;
out += batch_index * npoints * nsample * c;
const int index = threadIdx.y * blockDim.x + threadIdx.x;
const int stride = blockDim.y * blockDim.x;
for (int i = index; i < c * npoints; i += stride) {
const int l = i / npoints;
const int j = i % npoints;
for (int k = 0; k < nsample; ++k) {
int ii = idx[j * nsample + k];
out[(l * npoints + j) * nsample + k] = points[l * n + ii];
}
}
}
void group_points_kernel_wrapper(int b, int c, int n, int npoints, int nsample,
const float *points, const int *idx,
float *out) {
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
group_points_kernel<<<b, opt_block_config(npoints, c), 0, stream>>>(
b, c, n, npoints, nsample, points, idx, out);
CUDA_CHECK_ERRORS();
}
// input: grad_out(b, c, npoints, nsample), idx(b, npoints, nsample)
// output: grad_points(b, c, n)
__global__ void group_points_grad_kernel(int b, int c, int n, int npoints,
int nsample,
const float *__restrict__ grad_out,
const int *__restrict__ idx,
float *__restrict__ grad_points) {
int batch_index = blockIdx.x;
grad_out += batch_index * npoints * nsample * c;
idx += batch_index * npoints * nsample;
grad_points += batch_index * n * c;
const int index = threadIdx.y * blockDim.x + threadIdx.x;
const int stride = blockDim.y * blockDim.x;
for (int i = index; i < c * npoints; i += stride) {
const int l = i / npoints;
const int j = i % npoints;
for (int k = 0; k < nsample; ++k) {
int ii = idx[j * nsample + k];
atomicAdd(grad_points + l * n + ii,
grad_out[(l * npoints + j) * nsample + k]);
}
}
}
void group_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
int nsample, const float *grad_out,
const int *idx, float *grad_points) {
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
group_points_grad_kernel<<<b, opt_block_config(npoints, c), 0, stream>>>(
b, c, n, npoints, nsample, grad_out, idx, grad_points);
CUDA_CHECK_ERRORS();
}
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate.cpp
+99
View File
@@ -0,0 +1,99 @@
#include "interpolate.h"
#include "utils.h"
void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown,
const float *known, float *dist2, int *idx);
void three_interpolate_kernel_wrapper(int b, int c, int m, int n,
const float *points, const int *idx,
const float *weight, float *out);
void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m,
const float *grad_out,
const int *idx, const float *weight,
float *grad_points);
std::vector<at::Tensor> three_nn(at::Tensor unknowns, at::Tensor knows) {
CHECK_CONTIGUOUS(unknowns);
CHECK_CONTIGUOUS(knows);
CHECK_IS_FLOAT(unknowns);
CHECK_IS_FLOAT(knows);
if (unknowns.is_cuda()) {
CHECK_CUDA(knows);
}
at::Tensor idx =
torch::zeros({unknowns.size(0), unknowns.size(1), 3},
at::device(unknowns.device()).dtype(at::ScalarType::Int));
at::Tensor dist2 =
torch::zeros({unknowns.size(0), unknowns.size(1), 3},
at::device(unknowns.device()).dtype(at::ScalarType::Float));
if (unknowns.is_cuda()) {
three_nn_kernel_wrapper(unknowns.size(0), unknowns.size(1), knows.size(1),
unknowns.data_ptr<float>(), knows.data_ptr<float>(),
dist2.data_ptr<float>(), idx.data_ptr<int>());
} else {
AT_ASSERT(false, "CPU not supported");
}
return {dist2, idx};
}
at::Tensor three_interpolate(at::Tensor points, at::Tensor idx,
at::Tensor weight) {
CHECK_CONTIGUOUS(points);
CHECK_CONTIGUOUS(idx);
CHECK_CONTIGUOUS(weight);
CHECK_IS_FLOAT(points);
CHECK_IS_INT(idx);
CHECK_IS_FLOAT(weight);
if (points.is_cuda()) {
CHECK_CUDA(idx);
CHECK_CUDA(weight);
}
at::Tensor output =
torch::zeros({points.size(0), points.size(1), idx.size(1)},
at::device(points.device()).dtype(at::ScalarType::Float));
if (points.is_cuda()) {
three_interpolate_kernel_wrapper(
points.size(0), points.size(1), points.size(2), idx.size(1),
points.data_ptr<float>(), idx.data_ptr<int>(), weight.data_ptr<float>(),
output.data_ptr<float>());
} else {
AT_ASSERT(false, "CPU not supported");
}
return output;
}
at::Tensor three_interpolate_grad(at::Tensor grad_out, at::Tensor idx,
at::Tensor weight, const int m) {
CHECK_CONTIGUOUS(grad_out);
CHECK_CONTIGUOUS(idx);
CHECK_CONTIGUOUS(weight);
CHECK_IS_FLOAT(grad_out);
CHECK_IS_INT(idx);
CHECK_IS_FLOAT(weight);
if (grad_out.is_cuda()) {
CHECK_CUDA(idx);
CHECK_CUDA(weight);
}
at::Tensor output =
torch::zeros({grad_out.size(0), grad_out.size(1), m},
at::device(grad_out.device()).dtype(at::ScalarType::Float));
if (grad_out.is_cuda()) {
three_interpolate_grad_kernel_wrapper(
grad_out.size(0), grad_out.size(1), grad_out.size(2), m,
grad_out.data_ptr<float>(), idx.data_ptr<int>(),
weight.data_ptr<float>(), output.data_ptr<float>());
} else {
AT_ASSERT(false, "CPU not supported");
}
return output;
}
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate_gpu.cu
+154
View File
@@ -0,0 +1,154 @@
#include <math.h>
#include <stdio.h>
#include <stdlib.h>
#include "cuda_utils.h"
// input: unknown(b, n, 3) known(b, m, 3)
// output: dist2(b, n, 3), idx(b, n, 3)
__global__ void three_nn_kernel(int b, int n, int m,
const float *__restrict__ unknown,
const float *__restrict__ known,
float *__restrict__ dist2,
int *__restrict__ idx) {
int batch_index = blockIdx.x;
unknown += batch_index * n * 3;
known += batch_index * m * 3;
dist2 += batch_index * n * 3;
idx += batch_index * n * 3;
int index = threadIdx.x;
int stride = blockDim.x;
for (int j = index; j < n; j += stride) {
float ux = unknown[j * 3 + 0];
float uy = unknown[j * 3 + 1];
float uz = unknown[j * 3 + 2];
double best1 = 1e40, best2 = 1e40, best3 = 1e40;
int besti1 = 0, besti2 = 0, besti3 = 0;
for (int k = 0; k < m; ++k) {
float x = known[k * 3 + 0];
float y = known[k * 3 + 1];
float z = known[k * 3 + 2];
float d = (ux - x) * (ux - x) + (uy - y) * (uy - y) + (uz - z) * (uz - z);
if (d < best1) {
best3 = best2;
besti3 = besti2;
best2 = best1;
besti2 = besti1;
best1 = d;
besti1 = k;
} else if (d < best2) {
best3 = best2;
besti3 = besti2;
best2 = d;
besti2 = k;
} else if (d < best3) {
best3 = d;
besti3 = k;
}
}
dist2[j * 3 + 0] = best1;
dist2[j * 3 + 1] = best2;
dist2[j * 3 + 2] = best3;
idx[j * 3 + 0] = besti1;
idx[j * 3 + 1] = besti2;
idx[j * 3 + 2] = besti3;
}
}
void three_nn_kernel_wrapper(int b, int n, int m, const float *unknown,
const float *known, float *dist2, int *idx) {
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
three_nn_kernel<<<b, opt_n_threads(n), 0, stream>>>(b, n, m, unknown, known,
dist2, idx);
CUDA_CHECK_ERRORS();
}
// input: points(b, c, m), idx(b, n, 3), weight(b, n, 3)
// output: out(b, c, n)
__global__ void three_interpolate_kernel(int b, int c, int m, int n,
const float *__restrict__ points,
const int *__restrict__ idx,
const float *__restrict__ weight,
float *__restrict__ out) {
int batch_index = blockIdx.x;
points += batch_index * m * c;
idx += batch_index * n * 3;
weight += batch_index * n * 3;
out += batch_index * n * c;
const int index = threadIdx.y * blockDim.x + threadIdx.x;
const int stride = blockDim.y * blockDim.x;
for (int i = index; i < c * n; i += stride) {
const int l = i / n;
const int j = i % n;
float w1 = weight[j * 3 + 0];
float w2 = weight[j * 3 + 1];
float w3 = weight[j * 3 + 2];
int i1 = idx[j * 3 + 0];
int i2 = idx[j * 3 + 1];
int i3 = idx[j * 3 + 2];
out[i] = points[l * m + i1] * w1 + points[l * m + i2] * w2 +
points[l * m + i3] * w3;
}
}
void three_interpolate_kernel_wrapper(int b, int c, int m, int n,
const float *points, const int *idx,
const float *weight, float *out) {
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
three_interpolate_kernel<<<b, opt_block_config(n, c), 0, stream>>>(
b, c, m, n, points, idx, weight, out);
CUDA_CHECK_ERRORS();
}
// input: grad_out(b, c, n), idx(b, n, 3), weight(b, n, 3)
// output: grad_points(b, c, m)
__global__ void three_interpolate_grad_kernel(
int b, int c, int n, int m, const float *__restrict__ grad_out,
const int *__restrict__ idx, const float *__restrict__ weight,
float *__restrict__ grad_points) {
int batch_index = blockIdx.x;
grad_out += batch_index * n * c;
idx += batch_index * n * 3;
weight += batch_index * n * 3;
grad_points += batch_index * m * c;
const int index = threadIdx.y * blockDim.x + threadIdx.x;
const int stride = blockDim.y * blockDim.x;
for (int i = index; i < c * n; i += stride) {
const int l = i / n;
const int j = i % n;
float w1 = weight[j * 3 + 0];
float w2 = weight[j * 3 + 1];
float w3 = weight[j * 3 + 2];
int i1 = idx[j * 3 + 0];
int i2 = idx[j * 3 + 1];
int i3 = idx[j * 3 + 2];
atomicAdd(grad_points + l * m + i1, grad_out[i] * w1);
atomicAdd(grad_points + l * m + i2, grad_out[i] * w2);
atomicAdd(grad_points + l * m + i3, grad_out[i] * w3);
}
}
void three_interpolate_grad_kernel_wrapper(int b, int c, int n, int m,
const float *grad_out,
const int *idx, const float *weight,
float *grad_points) {
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
three_interpolate_grad_kernel<<<b, opt_block_config(n, c), 0, stream>>>(
b, c, n, m, grad_out, idx, weight, grad_points);
CUDA_CHECK_ERRORS();
}
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling.cpp
+87
View File
@@ -0,0 +1,87 @@
#include "sampling.h"
#include "utils.h"
void gather_points_kernel_wrapper(int b, int c, int n, int npoints,
const float *points, const int *idx,
float *out);
void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
const float *grad_out, const int *idx,
float *grad_points);
void furthest_point_sampling_kernel_wrapper(int b, int n, int m,
const float *dataset, float *temp,
int *idxs);
at::Tensor gather_points(at::Tensor points, at::Tensor idx) {
CHECK_CONTIGUOUS(points);
CHECK_CONTIGUOUS(idx);
CHECK_IS_FLOAT(points);
CHECK_IS_INT(idx);
if (points.is_cuda()) {
CHECK_CUDA(idx);
}
at::Tensor output =
torch::zeros({points.size(0), points.size(1), idx.size(1)},
at::device(points.device()).dtype(at::ScalarType::Float));
if (points.is_cuda()) {
gather_points_kernel_wrapper(points.size(0), points.size(1), points.size(2),
idx.size(1), points.data_ptr<float>(),
idx.data_ptr<int>(), output.data_ptr<float>());
} else {
AT_ASSERT(false, "CPU not supported");
}
return output;
}
at::Tensor gather_points_grad(at::Tensor grad_out, at::Tensor idx,
const int n) {
CHECK_CONTIGUOUS(grad_out);
CHECK_CONTIGUOUS(idx);
CHECK_IS_FLOAT(grad_out);
CHECK_IS_INT(idx);
if (grad_out.is_cuda()) {
CHECK_CUDA(idx);
}
at::Tensor output =
torch::zeros({grad_out.size(0), grad_out.size(1), n},
at::device(grad_out.device()).dtype(at::ScalarType::Float));
if (grad_out.is_cuda()) {
gather_points_grad_kernel_wrapper(grad_out.size(0), grad_out.size(1), n,
idx.size(1), grad_out.data_ptr<float>(),
idx.data_ptr<int>(),
output.data_ptr<float>());
} else {
AT_ASSERT(false, "CPU not supported");
}
return output;
}
at::Tensor furthest_point_sampling(at::Tensor points, const int nsamples) {
CHECK_CONTIGUOUS(points);
CHECK_IS_FLOAT(points);
at::Tensor output =
torch::zeros({points.size(0), nsamples},
at::device(points.device()).dtype(at::ScalarType::Int));
at::Tensor tmp =
torch::full({points.size(0), points.size(1)}, 1e10,
at::device(points.device()).dtype(at::ScalarType::Float));
if (points.is_cuda()) {
furthest_point_sampling_kernel_wrapper(
points.size(0), points.size(1), nsamples, points.data_ptr<float>(),
tmp.data_ptr<float>(), output.data_ptr<int>());
} else {
AT_ASSERT(false, "CPU not supported");
}
return output;
}
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling_gpu.cu
+229
View File
@@ -0,0 +1,229 @@
#include <stdio.h>
#include <stdlib.h>
#include "cuda_utils.h"
// input: points(b, c, n) idx(b, m)
// output: out(b, c, m)
__global__ void gather_points_kernel(int b, int c, int n, int m,
const float *__restrict__ points,
const int *__restrict__ idx,
float *__restrict__ out) {
for (int i = blockIdx.x; i < b; i += gridDim.x) {
for (int l = blockIdx.y; l < c; l += gridDim.y) {
for (int j = threadIdx.x; j < m; j += blockDim.x) {
int a = idx[i * m + j];
out[(i * c + l) * m + j] = points[(i * c + l) * n + a];
}
}
}
}
void gather_points_kernel_wrapper(int b, int c, int n, int npoints,
const float *points, const int *idx,
float *out) {
gather_points_kernel<<<dim3(b, c, 1), opt_n_threads(npoints), 0,
at::cuda::getCurrentCUDAStream()>>>(b, c, n, npoints,
points, idx, out);
CUDA_CHECK_ERRORS();
}
// input: grad_out(b, c, m) idx(b, m)
// output: grad_points(b, c, n)
__global__ void gather_points_grad_kernel(int b, int c, int n, int m,
const float *__restrict__ grad_out,
const int *__restrict__ idx,
float *__restrict__ grad_points) {
for (int i = blockIdx.x; i < b; i += gridDim.x) {
for (int l = blockIdx.y; l < c; l += gridDim.y) {
for (int j = threadIdx.x; j < m; j += blockDim.x) {
int a = idx[i * m + j];
atomicAdd(grad_points + (i * c + l) * n + a,
grad_out[(i * c + l) * m + j]);
}
}
}
}
void gather_points_grad_kernel_wrapper(int b, int c, int n, int npoints,
const float *grad_out, const int *idx,
float *grad_points) {
gather_points_grad_kernel<<<dim3(b, c, 1), opt_n_threads(npoints), 0,
at::cuda::getCurrentCUDAStream()>>>(
b, c, n, npoints, grad_out, idx, grad_points);
CUDA_CHECK_ERRORS();
}
__device__ void __update(float *__restrict__ dists, int *__restrict__ dists_i,
int idx1, int idx2) {
const float v1 = dists[idx1], v2 = dists[idx2];
const int i1 = dists_i[idx1], i2 = dists_i[idx2];
dists[idx1] = max(v1, v2);
dists_i[idx1] = v2 > v1 ? i2 : i1;
}
// Input dataset: (b, n, 3), tmp: (b, n)
// Ouput idxs (b, m)
template <unsigned int block_size>
__global__ void furthest_point_sampling_kernel(
int b, int n, int m, const float *__restrict__ dataset,
float *__restrict__ temp, int *__restrict__ idxs) {
if (m <= 0) return;
__shared__ float dists[block_size];
__shared__ int dists_i[block_size];
int batch_index = blockIdx.x;
dataset += batch_index * n * 3;
temp += batch_index * n;
idxs += batch_index * m;
int tid = threadIdx.x;
const int stride = block_size;
int old = 0;
if (threadIdx.x == 0) idxs[0] = old;
__syncthreads();
for (int j = 1; j < m; j++) {
int besti = 0;
float best = -1;
float x1 = dataset[old * 3 + 0];
float y1 = dataset[old * 3 + 1];
float z1 = dataset[old * 3 + 2];
for (int k = tid; k < n; k += stride) {
float x2, y2, z2;
x2 = dataset[k * 3 + 0];
y2 = dataset[k * 3 + 1];
z2 = dataset[k * 3 + 2];
float mag = (x2 * x2) + (y2 * y2) + (z2 * z2);
if (mag <= 1e-3) continue;
float d =
(x2 - x1) * (x2 - x1) + (y2 - y1) * (y2 - y1) + (z2 - z1) * (z2 - z1);
float d2 = min(d, temp[k]);
temp[k] = d2;
besti = d2 > best ? k : besti;
best = d2 > best ? d2 : best;
}
dists[tid] = best;
dists_i[tid] = besti;
__syncthreads();
if (block_size >= 512) {
if (tid < 256) {
__update(dists, dists_i, tid, tid + 256);
}
__syncthreads();
}
if (block_size >= 256) {
if (tid < 128) {
__update(dists, dists_i, tid, tid + 128);
}
__syncthreads();
}
if (block_size >= 128) {
if (tid < 64) {
__update(dists, dists_i, tid, tid + 64);
}
__syncthreads();
}
if (block_size >= 64) {
if (tid < 32) {
__update(dists, dists_i, tid, tid + 32);
}
__syncthreads();
}
if (block_size >= 32) {
if (tid < 16) {
__update(dists, dists_i, tid, tid + 16);
}
__syncthreads();
}
if (block_size >= 16) {
if (tid < 8) {
__update(dists, dists_i, tid, tid + 8);
}
__syncthreads();
}
if (block_size >= 8) {
if (tid < 4) {
__update(dists, dists_i, tid, tid + 4);
}
__syncthreads();
}
if (block_size >= 4) {
if (tid < 2) {
__update(dists, dists_i, tid, tid + 2);
}
__syncthreads();
}
if (block_size >= 2) {
if (tid < 1) {
__update(dists, dists_i, tid, tid + 1);
}
__syncthreads();
}
old = dists_i[0];
if (tid == 0) idxs[j] = old;
}
}
void furthest_point_sampling_kernel_wrapper(int b, int n, int m,
const float *dataset, float *temp,
int *idxs) {
unsigned int n_threads = opt_n_threads(n);
cudaStream_t stream = at::cuda::getCurrentCUDAStream();
switch (n_threads) {
case 512:
furthest_point_sampling_kernel<512>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 256:
furthest_point_sampling_kernel<256>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 128:
furthest_point_sampling_kernel<128>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 64:
furthest_point_sampling_kernel<64>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 32:
furthest_point_sampling_kernel<32>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 16:
furthest_point_sampling_kernel<16>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 8:
furthest_point_sampling_kernel<8>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 4:
furthest_point_sampling_kernel<4>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 2:
furthest_point_sampling_kernel<2>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
case 1:
furthest_point_sampling_kernel<1>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
break;
default:
furthest_point_sampling_kernel<512>
<<<b, n_threads, 0, stream>>>(b, n, m, dataset, temp, idxs);
}
CUDA_CHECK_ERRORS();
}
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_version.py
+1
View File
@@ -0,0 +1 @@
__version__ = "3.0.0"
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/pointnet2_modules.py
+209
View File
@@ -0,0 +1,209 @@
from typing import List, Optional, Tuple
import torch
import torch.nn as nn
import torch.nn.functional as F
from pointnet2_ops import pointnet2_utils
def build_shared_mlp(mlp_spec: List[int], bn: bool = True):
layers = []
for i in range(1, len(mlp_spec)):
layers.append(
nn.Conv2d(mlp_spec[i - 1], mlp_spec[i], kernel_size=1, bias=not bn)
)
if bn:
layers.append(nn.BatchNorm2d(mlp_spec[i]))
layers.append(nn.ReLU(True))
return nn.Sequential(*layers)
class _PointnetSAModuleBase(nn.Module):
def __init__(self):
super(_PointnetSAModuleBase, self).__init__()
self.npoint = None
self.groupers = None
self.mlps = None
def forward(
self, xyz: torch.Tensor, features: Optional[torch.Tensor]
) -> Tuple[torch.Tensor, torch.Tensor]:
r"""
Parameters
----------
xyz : torch.Tensor
(B, N, 3) tensor of the xyz coordinates of the features
features : torch.Tensor
(B, C, N) tensor of the descriptors of the the features
Returns
-------
new_xyz : torch.Tensor
(B, npoint, 3) tensor of the new features' xyz
new_features : torch.Tensor
(B, \sum_k(mlps[k][-1]), npoint) tensor of the new_features descriptors
"""
new_features_list = []
xyz_flipped = xyz.transpose(1, 2).contiguous()
new_xyz = (
pointnet2_utils.gather_operation(
xyz_flipped, pointnet2_utils.furthest_point_sample(xyz, self.npoint)
)
.transpose(1, 2)
.contiguous()
if self.npoint is not None
else None
)
for i in range(len(self.groupers)):
new_features = self.groupers[i](
xyz, new_xyz, features
) # (B, C, npoint, nsample)
new_features = self.mlps[i](new_features) # (B, mlp[-1], npoint, nsample)
new_features = F.max_pool2d(
new_features, kernel_size=[1, new_features.size(3)]
) # (B, mlp[-1], npoint, 1)
new_features = new_features.squeeze(-1) # (B, mlp[-1], npoint)
new_features_list.append(new_features)
return new_xyz, torch.cat(new_features_list, dim=1)
class PointnetSAModuleMSG(_PointnetSAModuleBase):
r"""Pointnet set abstrction layer with multiscale grouping
Parameters
----------
npoint : int
Number of features
radii : list of float32
list of radii to group with
nsamples : list of int32
Number of samples in each ball query
mlps : list of list of int32
Spec of the pointnet before the global max_pool for each scale
bn : bool
Use batchnorm
"""
def __init__(self, npoint, radii, nsamples, mlps, bn=True, use_xyz=True):
# type: (PointnetSAModuleMSG, int, List[float], List[int], List[List[int]], bool, bool) -> None
super(PointnetSAModuleMSG, self).__init__()
assert len(radii) == len(nsamples) == len(mlps)
self.npoint = npoint
self.groupers = nn.ModuleList()
self.mlps = nn.ModuleList()
for i in range(len(radii)):
radius = radii[i]
nsample = nsamples[i]
self.groupers.append(
pointnet2_utils.QueryAndGroup(radius, nsample, use_xyz=use_xyz)
if npoint is not None
else pointnet2_utils.GroupAll(use_xyz)
)
mlp_spec = mlps[i]
if use_xyz:
mlp_spec[0] += 3
self.mlps.append(build_shared_mlp(mlp_spec, bn))
class PointnetSAModule(PointnetSAModuleMSG):
r"""Pointnet set abstrction layer
Parameters
----------
npoint : int
Number of features
radius : float
Radius of ball
nsample : int
Number of samples in the ball query
mlp : list
Spec of the pointnet before the global max_pool
bn : bool
Use batchnorm
"""
def __init__(
self, mlp, npoint=None, radius=None, nsample=None, bn=True, use_xyz=True
):
# type: (PointnetSAModule, List[int], int, float, int, bool, bool) -> None
super(PointnetSAModule, self).__init__(
mlps=[mlp],
npoint=npoint,
radii=[radius],
nsamples=[nsample],
bn=bn,
use_xyz=use_xyz,
)
class PointnetFPModule(nn.Module):
r"""Propigates the features of one set to another
Parameters
----------
mlp : list
Pointnet module parameters
bn : bool
Use batchnorm
"""
def __init__(self, mlp, bn=True):
# type: (PointnetFPModule, List[int], bool) -> None
super(PointnetFPModule, self).__init__()
self.mlp = build_shared_mlp(mlp, bn=bn)
def forward(self, unknown, known, unknow_feats, known_feats):
# type: (PointnetFPModule, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor) -> torch.Tensor
r"""
Parameters
----------
unknown : torch.Tensor
(B, n, 3) tensor of the xyz positions of the unknown features
known : torch.Tensor
(B, m, 3) tensor of the xyz positions of the known features
unknow_feats : torch.Tensor
(B, C1, n) tensor of the features to be propigated to
known_feats : torch.Tensor
(B, C2, m) tensor of features to be propigated
Returns
-------
new_features : torch.Tensor
(B, mlp[-1], n) tensor of the features of the unknown features
"""
if known is not None:
dist, idx = pointnet2_utils.three_nn(unknown, known)
dist_recip = 1.0 / (dist + 1e-8)
norm = torch.sum(dist_recip, dim=2, keepdim=True)
weight = dist_recip / norm
interpolated_feats = pointnet2_utils.three_interpolate(
known_feats, idx, weight
)
else:
interpolated_feats = known_feats.expand(
*(known_feats.size()[0:2] + [unknown.size(1)])
)
if unknow_feats is not None:
new_features = torch.cat(
[interpolated_feats, unknow_feats], dim=1
) # (B, C2 + C1, n)
else:
new_features = interpolated_feats
new_features = new_features.unsqueeze(-1)
new_features = self.mlp(new_features)
return new_features.squeeze(-1)
tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/pointnet2_utils.py
+379
View File
@@ -0,0 +1,379 @@
import torch
import torch.nn as nn
import warnings
from torch.autograd import Function
from typing import *
try:
import pointnet2_ops._ext as _ext
except ImportError:
from torch.utils.cpp_extension import load
import glob
import os.path as osp
import os
warnings.warn("Unable to load pointnet2_ops cpp extension. JIT Compiling.")
_ext_src_root = osp.join(osp.dirname(__file__), "_ext-src")
_ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob(
osp.join(_ext_src_root, "src", "*.cu")
)
_ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*"))
os.environ["TORCH_CUDA_ARCH_LIST"] = "3.7+PTX;5.0;6.0;6.1;6.2;7.0;7.5"
_ext = load(
"_ext",
sources=_ext_sources,
extra_include_paths=[osp.join(_ext_src_root, "include")],
extra_cflags=["-O3"],
extra_cuda_cflags=["-O3", "-Xfatbin", "-compress-all"],
with_cuda=True,
)
class FurthestPointSampling(Function):
@staticmethod
def forward(ctx, xyz, npoint):
# type: (Any, torch.Tensor, int) -> torch.Tensor
r"""
Uses iterative furthest point sampling to select a set of npoint features that have the largest
minimum distance
Parameters
----------
xyz : torch.Tensor
(B, N, 3) tensor where N > npoint
npoint : int32
number of features in the sampled set
Returns
-------
torch.Tensor
(B, npoint) tensor containing the set
"""
out = _ext.furthest_point_sampling(xyz, npoint)
ctx.mark_non_differentiable(out)
return out
@staticmethod
def backward(ctx, grad_out):
return ()
furthest_point_sample = FurthestPointSampling.apply
class GatherOperation(Function):
@staticmethod
def forward(ctx, features, idx):
# type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor
r"""
Parameters
----------
features : torch.Tensor
(B, C, N) tensor
idx : torch.Tensor
(B, npoint) tensor of the features to gather
Returns
-------
torch.Tensor
(B, C, npoint) tensor
"""
ctx.save_for_backward(idx, features)
return _ext.gather_points(features, idx)
@staticmethod
def backward(ctx, grad_out):
idx, features = ctx.saved_tensors
N = features.size(2)
grad_features = _ext.gather_points_grad(grad_out.contiguous(), idx, N)
return grad_features, None
gather_operation = GatherOperation.apply
class ThreeNN(Function):
@staticmethod
def forward(ctx, unknown, known):
# type: (Any, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]
r"""
Find the three nearest neighbors of unknown in known
Parameters
----------
unknown : torch.Tensor
(B, n, 3) tensor of known features
known : torch.Tensor
(B, m, 3) tensor of unknown features
Returns
-------
dist : torch.Tensor
(B, n, 3) l2 distance to the three nearest neighbors
idx : torch.Tensor
(B, n, 3) index of 3 nearest neighbors
"""
dist2, idx = _ext.three_nn(unknown, known)
dist = torch.sqrt(dist2)
ctx.mark_non_differentiable(dist, idx)
return dist, idx
@staticmethod
def backward(ctx, grad_dist, grad_idx):
return ()
three_nn = ThreeNN.apply
class ThreeInterpolate(Function):
@staticmethod
def forward(ctx, features, idx, weight):
# type(Any, torch.Tensor, torch.Tensor, torch.Tensor) -> Torch.Tensor
r"""
Performs weight linear interpolation on 3 features
Parameters
----------
features : torch.Tensor
(B, c, m) Features descriptors to be interpolated from
idx : torch.Tensor
(B, n, 3) three nearest neighbors of the target features in features
weight : torch.Tensor
(B, n, 3) weights
Returns
-------
torch.Tensor
(B, c, n) tensor of the interpolated features
"""
ctx.save_for_backward(idx, weight, features)
return _ext.three_interpolate(features, idx, weight)
@staticmethod
def backward(ctx, grad_out):
# type: (Any, torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]
r"""
Parameters
----------
grad_out : torch.Tensor
(B, c, n) tensor with gradients of ouputs
Returns
-------
grad_features : torch.Tensor
(B, c, m) tensor with gradients of features
None
None
"""
idx, weight, features = ctx.saved_tensors
m = features.size(2)
grad_features = _ext.three_interpolate_grad(
grad_out.contiguous(), idx, weight, m
)
return grad_features, torch.zeros_like(idx), torch.zeros_like(weight)
three_interpolate = ThreeInterpolate.apply
class GroupingOperation(Function):
@staticmethod
def forward(ctx, features, idx):
# type: (Any, torch.Tensor, torch.Tensor) -> torch.Tensor
r"""
Parameters
----------
features : torch.Tensor
(B, C, N) tensor of features to group
idx : torch.Tensor
(B, npoint, nsample) tensor containing the indicies of features to group with
Returns
-------
torch.Tensor
(B, C, npoint, nsample) tensor
"""
ctx.save_for_backward(idx, features)
return _ext.group_points(features, idx)
@staticmethod
def backward(ctx, grad_out):
# type: (Any, torch.tensor) -> Tuple[torch.Tensor, torch.Tensor]
r"""
Parameters
----------
grad_out : torch.Tensor
(B, C, npoint, nsample) tensor of the gradients of the output from forward
Returns
-------
torch.Tensor
(B, C, N) gradient of the features
None
"""
idx, features = ctx.saved_tensors
N = features.size(2)
grad_features = _ext.group_points_grad(grad_out.contiguous(), idx, N)
return grad_features, torch.zeros_like(idx)
grouping_operation = GroupingOperation.apply
class BallQuery(Function):
@staticmethod
def forward(ctx, radius, nsample, xyz, new_xyz):
# type: (Any, float, int, torch.Tensor, torch.Tensor) -> torch.Tensor
r"""
Parameters
----------
radius : float
radius of the balls
nsample : int
maximum number of features in the balls
xyz : torch.Tensor
(B, N, 3) xyz coordinates of the features
new_xyz : torch.Tensor
(B, npoint, 3) centers of the ball query
Returns
-------
torch.Tensor
(B, npoint, nsample) tensor with the indicies of the features that form the query balls
"""
output = _ext.ball_query(new_xyz, xyz, radius, nsample)
ctx.mark_non_differentiable(output)
return output
@staticmethod
def backward(ctx, grad_out):
return ()
ball_query = BallQuery.apply
class QueryAndGroup(nn.Module):
r"""
Groups with a ball query of radius
Parameters
---------
radius : float32
Radius of ball
nsample : int32
Maximum number of features to gather in the ball
"""
def __init__(self, radius, nsample, use_xyz=True):
# type: (QueryAndGroup, float, int, bool) -> None
super(QueryAndGroup, self).__init__()
self.radius, self.nsample, self.use_xyz = radius, nsample, use_xyz
def forward(self, xyz, new_xyz, features=None):
# type: (QueryAndGroup, torch.Tensor. torch.Tensor, torch.Tensor) -> Tuple[Torch.Tensor]
r"""
Parameters
----------
xyz : torch.Tensor
xyz coordinates of the features (B, N, 3)
new_xyz : torch.Tensor
centriods (B, npoint, 3)
features : torch.Tensor
Descriptors of the features (B, C, N)
Returns
-------
new_features : torch.Tensor
(B, 3 + C, npoint, nsample) tensor
"""
idx = ball_query(self.radius, self.nsample, xyz, new_xyz)
xyz_trans = xyz.transpose(1, 2).contiguous()
grouped_xyz = grouping_operation(xyz_trans, idx) # (B, 3, npoint, nsample)
grouped_xyz -= new_xyz.transpose(1, 2).unsqueeze(-1)
if features is not None:
grouped_features = grouping_operation(features, idx)
if self.use_xyz:
new_features = torch.cat(
[grouped_xyz, grouped_features], dim=1
) # (B, C + 3, npoint, nsample)
else:
new_features = grouped_features
else:
assert (
self.use_xyz
), "Cannot have not features and not use xyz as a feature!"
new_features = grouped_xyz
return new_features
class GroupAll(nn.Module):
r"""
Groups all features
Parameters
---------
"""
def __init__(self, use_xyz=True):
# type: (GroupAll, bool) -> None
super(GroupAll, self).__init__()
self.use_xyz = use_xyz
def forward(self, xyz, new_xyz, features=None):
# type: (GroupAll, torch.Tensor, torch.Tensor, torch.Tensor) -> Tuple[torch.Tensor]
r"""
Parameters
----------
xyz : torch.Tensor
xyz coordinates of the features (B, N, 3)
new_xyz : torch.Tensor
Ignored
features : torch.Tensor
Descriptors of the features (B, C, N)
Returns
-------
new_features : torch.Tensor
(B, C + 3, 1, N) tensor
"""
grouped_xyz = xyz.transpose(1, 2).unsqueeze(2)
if features is not None:
grouped_features = features.unsqueeze(2)
if self.use_xyz:
new_features = torch.cat(
[grouped_xyz, grouped_features], dim=1
) # (B, 3 + C, 1, N)
else:
new_features = grouped_features
else:
new_features = grouped_xyz
return new_features
tgs/models/snowflake/pointnet2_ops_lib/setup.py
+41
View File
@@ -0,0 +1,41 @@
import glob
import os
import os.path as osp
from setuptools import find_packages, setup
import torch
from torch.utils.cpp_extension import BuildExtension, CUDAExtension
this_dir = osp.dirname(osp.abspath(__file__))
_ext_src_root = osp.join("pointnet2_ops", "_ext-src")
_ext_sources = glob.glob(osp.join(_ext_src_root, "src", "*.cpp")) + glob.glob(
osp.join(_ext_src_root, "src", "*.cu")
)
_ext_headers = glob.glob(osp.join(_ext_src_root, "include", "*"))
requirements = ["torch>=1.4"]
exec(open(osp.join("pointnet2_ops", "_version.py")).read())
os.environ["TORCH_CUDA_ARCH_LIST"] = ".".join(map(str, torch.cuda.get_device_capability()))
# os.environ["TORCH_CUDA_ARCH_LIST"] = "5.0;6.0;6.1;6.2;7.0;7.5;8.0;8.6"
setup(
name="pointnet2_ops",
version=__version__,
author="Erik Wijmans",
packages=find_packages(),
install_requires=requirements,
ext_modules=[
CUDAExtension(
name="pointnet2_ops._ext",
sources=_ext_sources,
extra_compile_args={
"cxx": ["-O3"],
"nvcc": ["-O3", "-Xfatbin", "-compress-all"],
},
include_dirs=[osp.join(this_dir, _ext_src_root, "include")],
)
],
cmdclass={"build_ext": BuildExtension},
include_package_data=True,
)
tgs/models/snowflake/skip_transformer.py
+69
View File
@@ -0,0 +1,69 @@
# -*- coding: utf-8 -*-
# @Author: Peng Xiang
import torch
from torch import nn, einsum
from .utils import MLP_Res, grouping_operation, query_knn
class SkipTransformer(nn.Module):
def __init__(self, in_channel, dim=256, n_knn=16, pos_hidden_dim=64, attn_hidden_multiplier=4):
super(SkipTransformer, self).__init__()
self.mlp_v = MLP_Res(in_dim=in_channel*2, hidden_dim=in_channel, out_dim=in_channel)
self.n_knn = n_knn
self.conv_key = nn.Conv1d(in_channel, dim, 1)
self.conv_query = nn.Conv1d(in_channel, dim, 1)
self.conv_value = nn.Conv1d(in_channel, dim, 1)
self.pos_mlp = nn.Sequential(
nn.Conv2d(3, pos_hidden_dim, 1),
nn.BatchNorm2d(pos_hidden_dim),
nn.ReLU(),
nn.Conv2d(pos_hidden_dim, dim, 1)
)
self.attn_mlp = nn.Sequential(
nn.Conv2d(dim, dim * attn_hidden_multiplier, 1),
nn.BatchNorm2d(dim * attn_hidden_multiplier),
nn.ReLU(),
nn.Conv2d(dim * attn_hidden_multiplier, dim, 1)
)
self.conv_end = nn.Conv1d(dim, in_channel, 1)
def forward(self, pos, key, query, include_self=True):
"""
Args:
pos: (B, 3, N)
key: (B, in_channel, N)
query: (B, in_channel, N)
include_self: boolean
Returns:
Tensor: (B, in_channel, N), shape context feature
"""
value = self.mlp_v(torch.cat([key, query], 1))
identity = value
key = self.conv_key(key)
query = self.conv_query(query)
value = self.conv_value(value)
b, dim, n = value.shape
pos_flipped = pos.permute(0, 2, 1).contiguous()
idx_knn = query_knn(self.n_knn, pos_flipped, pos_flipped, include_self=include_self)
key = grouping_operation(key, idx_knn) # b, dim, n, n_knn
qk_rel = query.reshape((b, -1, n, 1)) - key
pos_rel = pos.reshape((b, -1, n, 1)) - grouping_operation(pos, idx_knn) # b, 3, n, n_knn
pos_embedding = self.pos_mlp(pos_rel)
attention = self.attn_mlp(qk_rel + pos_embedding) # b, dim, n, n_knn
attention = torch.softmax(attention, -1)
value = value.reshape((b, -1, n, 1)) + pos_embedding #
agg = einsum('b c i j, b c i j -> b c i', attention, value) # b, dim, n
y = self.conv_end(agg)
return y + identity
tgs/models/snowflake/utils.py
+741
View File
@@ -0,0 +1,741 @@
# -*- coding: utf-8 -*-
# @Author: Peng Xiang
import types
import torch
import torch.nn.functional as F
import numpy as np
from torch import nn, einsum
from pointnet2_ops.pointnet2_utils import furthest_point_sample, \
gather_operation, ball_query, three_nn, three_interpolate, grouping_operation
class Conv1d(nn.Module):
def __init__(self, in_channel, out_channel, kernel_size=1, stride=1, if_bn=True, activation_fn=torch.relu):
super(Conv1d, self).__init__()
self.conv = nn.Conv1d(in_channel, out_channel, kernel_size, stride=stride)
self.if_bn = if_bn
self.bn = nn.BatchNorm1d(out_channel)
self.activation_fn = activation_fn
def forward(self, input):
out = self.conv(input)
if self.if_bn:
out = self.bn(out)
if self.activation_fn is not None:
out = self.activation_fn(out)
return out
class Conv2d(nn.Module):
def __init__(self, in_channel, out_channel, kernel_size=(1, 1), stride=(1, 1), if_bn=True, activation_fn=torch.relu):
super(Conv2d, self).__init__()
self.conv = nn.Conv2d(in_channel, out_channel, kernel_size, stride=stride)
self.if_bn = if_bn
self.bn = nn.BatchNorm2d(out_channel)
self.activation_fn = activation_fn
def forward(self, input):
out = self.conv(input)
if self.if_bn:
out = self.bn(out)
if self.activation_fn is not None:
out = self.activation_fn(out)
return out
class MLP(nn.Module):
def __init__(self, in_channel, layer_dims, bn=None):
super(MLP, self).__init__()
layers = []
last_channel = in_channel
for out_channel in layer_dims[:-1]:
layers.append(nn.Linear(last_channel, out_channel))
if bn:
layers.append(nn.BatchNorm1d(out_channel))
layers.append(nn.ReLU())
last_channel = out_channel
layers.append(nn.Linear(last_channel, layer_dims[-1]))
self.mlp = nn.Sequential(*layers)
def forward(self, inputs):
return self.mlp(inputs)
class MLP_CONV(nn.Module):
def __init__(self, in_channel, layer_dims, bn=None):
super(MLP_CONV, self).__init__()
layers = []
last_channel = in_channel
for out_channel in layer_dims[:-1]:
layers.append(nn.Conv1d(last_channel, out_channel, 1))
if bn:
layers.append(nn.BatchNorm1d(out_channel))
layers.append(nn.ReLU())
last_channel = out_channel
layers.append(nn.Conv1d(last_channel, layer_dims[-1], 1))
self.mlp = nn.Sequential(*layers)
def forward(self, inputs):
return self.mlp(inputs)
class MLP_Res(nn.Module):
def __init__(self, in_dim=128, hidden_dim=None, out_dim=128):
super(MLP_Res, self).__init__()
if hidden_dim is None:
hidden_dim = in_dim
self.conv_1 = nn.Conv1d(in_dim, hidden_dim, 1)
self.conv_2 = nn.Conv1d(hidden_dim, out_dim, 1)
self.conv_shortcut = nn.Conv1d(in_dim, out_dim, 1)
def forward(self, x):
"""
Args:
x: (B, out_dim, n)
"""
shortcut = self.conv_shortcut(x)
out = self.conv_2(torch.relu(self.conv_1(x))) + shortcut
return out
def sample_and_group(xyz, points, npoint, nsample, radius, use_xyz=True):
"""
Args:
xyz: Tensor, (B, 3, N)
points: Tensor, (B, f, N)
npoint: int
nsample: int
radius: float
use_xyz: boolean
Returns:
new_xyz: Tensor, (B, 3, npoint)
new_points: Tensor, (B, 3 | f+3 | f, npoint, nsample)
idx_local: Tensor, (B, npoint, nsample)
grouped_xyz: Tensor, (B, 3, npoint, nsample)
"""
xyz_flipped = xyz.permute(0, 2, 1).contiguous() # (B, N, 3)
new_xyz = gather_operation(xyz, furthest_point_sample(xyz_flipped, npoint)) # (B, 3, npoint)
idx = ball_query(radius, nsample, xyz_flipped, new_xyz.permute(0, 2, 1).contiguous()) # (B, npoint, nsample)
grouped_xyz = grouping_operation(xyz, idx) # (B, 3, npoint, nsample)
grouped_xyz -= new_xyz.unsqueeze(3).repeat(1, 1, 1, nsample)
if points is not None:
grouped_points = grouping_operation(points, idx) # (B, f, npoint, nsample)
if use_xyz:
new_points = torch.cat([grouped_xyz, grouped_points], 1)
else:
new_points = grouped_points
else:
new_points = grouped_xyz
return new_xyz, new_points, idx, grouped_xyz
def sample_and_group_all(xyz, points, use_xyz=True):
"""
Args:
xyz: Tensor, (B, 3, nsample)
points: Tensor, (B, f, nsample)
use_xyz: boolean
Returns:
new_xyz: Tensor, (B, 3, 1)
new_points: Tensor, (B, f|f+3|3, 1, nsample)
idx: Tensor, (B, 1, nsample)
grouped_xyz: Tensor, (B, 3, 1, nsample)
"""
b, _, nsample = xyz.shape
device = xyz.device
new_xyz = torch.zeros((1, 3, 1), dtype=torch.float, device=device).repeat(b, 1, 1)
grouped_xyz = xyz.reshape((b, 3, 1, nsample))
idx = torch.arange(nsample, device=device).reshape(1, 1, nsample).repeat(b, 1, 1)
if points is not None:
if use_xyz:
new_points = torch.cat([xyz, points], 1)
else:
new_points = points
new_points = new_points.unsqueeze(2)
else:
new_points = grouped_xyz
return new_xyz, new_points, idx, grouped_xyz
class PointNet_SA_Module(nn.Module):
def __init__(self, npoint, nsample, radius, in_channel, mlp, if_bn=True, group_all=False, use_xyz=True):
"""
Args:
npoint: int, number of points to sample
nsample: int, number of points in each local region
radius: float
in_channel: int, input channel of features(points)
mlp: list of int,
"""
super(PointNet_SA_Module, self).__init__()
self.npoint = npoint
self.nsample = nsample
self.radius = radius
self.mlp = mlp
self.group_all = group_all
self.use_xyz = use_xyz
if use_xyz:
in_channel += 3
last_channel = in_channel
self.mlp_conv = []
for out_channel in mlp:
self.mlp_conv.append(Conv2d(last_channel, out_channel, if_bn=if_bn))
last_channel = out_channel
self.mlp_conv = nn.Sequential(*self.mlp_conv)
def forward(self, xyz, points):
"""
Args:
xyz: Tensor, (B, 3, N)
points: Tensor, (B, f, N)
Returns:
new_xyz: Tensor, (B, 3, npoint)
new_points: Tensor, (B, mlp[-1], npoint)
"""
if self.group_all:
new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, self.use_xyz)
else:
new_xyz, new_points, idx, grouped_xyz = sample_and_group(xyz, points, self.npoint, self.nsample, self.radius, self.use_xyz)
new_points = self.mlp_conv(new_points)
new_points = torch.max(new_points, 3)[0]
return new_xyz, new_points
class PointNet_FP_Module(nn.Module):
def __init__(self, in_channel, mlp, use_points1=False, in_channel_points1=None, if_bn=True):
"""
Args:
in_channel: int, input channel of points2
mlp: list of int
use_points1: boolean, if use points
in_channel_points1: int, input channel of points1
"""
super(PointNet_FP_Module, self).__init__()
self.use_points1 = use_points1
if use_points1:
in_channel += in_channel_points1
last_channel = in_channel
self.mlp_conv = []
for out_channel in mlp:
self.mlp_conv.append(Conv1d(last_channel, out_channel, if_bn=if_bn))
last_channel = out_channel
self.mlp_conv = nn.Sequential(*self.mlp_conv)
def forward(self, xyz1, xyz2, points1, points2):
"""
Args:
xyz1: Tensor, (B, 3, N)
xyz2: Tensor, (B, 3, M)
points1: Tensor, (B, in_channel, N)
points2: Tensor, (B, in_channel, M)
Returns:MLP_CONV
new_points: Tensor, (B, mlp[-1], N)
"""
dist, idx = three_nn(xyz1.permute(0, 2, 1).contiguous(), xyz2.permute(0, 2, 1).contiguous())
dist = torch.clamp_min(dist, 1e-10) # (B, N, 3)
recip_dist = 1.0/dist
norm = torch.sum(recip_dist, 2, keepdim=True).repeat((1, 1, 3))
weight = recip_dist / norm
interpolated_points = three_interpolate(points2, idx, weight) # B, in_channel, N
if self.use_points1:
new_points = torch.cat([interpolated_points, points1], 1)
else:
new_points = interpolated_points
new_points = self.mlp_conv(new_points)
return new_points
def square_distance(src, dst):
"""
Calculate Euclid distance between each two points.
src^T * dst = xn * xm + yn * ym + zn * zm；
sum(src^2, dim=-1) = xn*xn + yn*yn + zn*zn;
sum(dst^2, dim=-1) = xm*xm + ym*ym + zm*zm;
dist = (xn - xm)^2 + (yn - ym)^2 + (zn - zm)^2
= sum(src**2,dim=-1)+sum(dst**2,dim=-1)-2*src^T*dst
Input:
src: source points, [B, N, C]
dst: target points, [B, M, C]
Output:
dist: per-point square distance, [B, N, M]
"""
B, N, _ = src.shape
_, M, _ = dst.shape
dist = -2 * torch.matmul(src, dst.permute(0, 2, 1)) # B, N, M
dist += torch.sum(src ** 2, -1).view(B, N, 1)
dist += torch.sum(dst ** 2, -1).view(B, 1, M)
return dist
def query_knn(nsample, xyz, new_xyz, include_self=True):
"""Find k-NN of new_xyz in xyz"""
pad = 0 if include_self else 1
sqrdists = square_distance(new_xyz, xyz) # B, S, N
idx = torch.argsort(sqrdists, dim=-1, descending=False)[:, :, pad: nsample+pad]
return idx.int()
def sample_and_group_knn(xyz, points, npoint, k, use_xyz=True, idx=None):
"""
Args:
xyz: Tensor, (B, 3, N)
points: Tensor, (B, f, N)
npoint: int
nsample: int
radius: float
use_xyz: boolean
Returns:
new_xyz: Tensor, (B, 3, npoint)
new_points: Tensor, (B, 3 | f+3 | f, npoint, nsample)
idx_local: Tensor, (B, npoint, nsample)
grouped_xyz: Tensor, (B, 3, npoint, nsample)
"""
xyz_flipped = xyz.permute(0, 2, 1).contiguous() # (B, N, 3)
new_xyz = gather_operation(xyz, furthest_point_sample(xyz_flipped, npoint)) # (B, 3, npoint)
if idx is None:
idx = query_knn(k, xyz_flipped, new_xyz.permute(0, 2, 1).contiguous())
grouped_xyz = grouping_operation(xyz, idx) # (B, 3, npoint, nsample)
grouped_xyz -= new_xyz.unsqueeze(3).repeat(1, 1, 1, k)
if points is not None:
grouped_points = grouping_operation(points, idx) # (B, f, npoint, nsample)
if use_xyz:
new_points = torch.cat([grouped_xyz, grouped_points], 1)
else:
new_points = grouped_points
else:
new_points = grouped_xyz
return new_xyz, new_points, idx, grouped_xyz
class PointNet_SA_Module_KNN(nn.Module):
def __init__(self, npoint, nsample, in_channel, mlp, if_bn=True, group_all=False, use_xyz=True, if_idx=False):
"""
Args:
npoint: int, number of points to sample
nsample: int, number of points in each local region
radius: float
in_channel: int, input channel of features(points)
mlp: list of int,
"""
super(PointNet_SA_Module_KNN, self).__init__()
self.npoint = npoint
self.nsample = nsample
self.mlp = mlp
self.group_all = group_all
self.use_xyz = use_xyz
self.if_idx = if_idx
if use_xyz:
in_channel += 3
last_channel = in_channel
self.mlp_conv = []
for out_channel in mlp[:-1]:
self.mlp_conv.append(Conv2d(last_channel, out_channel, if_bn=if_bn))
last_channel = out_channel
self.mlp_conv.append(Conv2d(last_channel, mlp[-1], if_bn=False, activation_fn=None))
self.mlp_conv = nn.Sequential(*self.mlp_conv)
def forward(self, xyz, points, idx=None):
"""
Args:
xyz: Tensor, (B, 3, N)
points: Tensor, (B, f, N)
Returns:
new_xyz: Tensor, (B, 3, npoint)
new_points: Tensor, (B, mlp[-1], npoint)
"""
if self.group_all:
new_xyz, new_points, idx, grouped_xyz = sample_and_group_all(xyz, points, self.use_xyz)
else:
new_xyz, new_points, idx, grouped_xyz = sample_and_group_knn(xyz, points, self.npoint, self.nsample, self.use_xyz, idx=idx)
new_points = self.mlp_conv(new_points)
new_points = torch.max(new_points, 3)[0]
if self.if_idx:
return new_xyz, new_points, idx
else:
return new_xyz, new_points
def fps_subsample(pcd, n_points=2048):
"""
Args
pcd: (b, 16384, 3)
returns
new_pcd: (b, n_points, 3)
"""
new_pcd = gather_operation(pcd.permute(0, 2, 1).contiguous(), furthest_point_sample(pcd, n_points))
new_pcd = new_pcd.permute(0, 2, 1).contiguous()
return new_pcd
class Transformer(nn.Module):
def __init__(self, in_channel, dim=256, n_knn=16, pos_hidden_dim=64, attn_hidden_multiplier=4):
super(Transformer, self).__init__()
self.n_knn = n_knn
self.conv_key = nn.Conv1d(dim, dim, 1)
self.conv_query = nn.Conv1d(dim, dim, 1)
self.conv_value = nn.Conv1d(dim, dim, 1)
self.pos_mlp = nn.Sequential(
nn.Conv2d(3, pos_hidden_dim, 1),
nn.BatchNorm2d(pos_hidden_dim),
nn.ReLU(),
nn.Conv2d(pos_hidden_dim, dim, 1)
)
self.attn_mlp = nn.Sequential(
nn.Conv2d(dim, dim * attn_hidden_multiplier, 1),
nn.BatchNorm2d(dim * attn_hidden_multiplier),
nn.ReLU(),
nn.Conv2d(dim * attn_hidden_multiplier, dim, 1)
)
self.linear_start = nn.Conv1d(in_channel, dim, 1)
self.linear_end = nn.Conv1d(dim, in_channel, 1)
def forward(self, x, pos):
"""feed forward of transformer
Args:
x: Tensor of features, (B, in_channel, n)
pos: Tensor of positions, (B, 3, n)
Returns:
y: Tensor of features with attention, (B, in_channel, n)
"""
identity = x
x = self.linear_start(x)
b, dim, n = x.shape
pos_flipped = pos.permute(0, 2, 1).contiguous()
idx_knn = query_knn(self.n_knn, pos_flipped, pos_flipped)
key = self.conv_key(x)
value = self.conv_value(x)
query = self.conv_query(x)
key = grouping_operation(key, idx_knn) # b, dim, n, n_knn
qk_rel = query.reshape((b, -1, n, 1)) - key
pos_rel = pos.reshape((b, -1, n, 1)) - grouping_operation(pos, idx_knn) # b, 3, n, n_knn
pos_embedding = self.pos_mlp(pos_rel) # b, dim, n, n_knn
attention = self.attn_mlp(qk_rel + pos_embedding)
attention = torch.softmax(attention, -1)
value = value.reshape((b, -1, n, 1)) + pos_embedding
agg = einsum('b c i j, b c i j -> b c i', attention, value) # b, dim, n
y = self.linear_end(agg)
return y+identity
class CouplingLayer(nn.Module):
def __init__(self, d, intermediate_dim, swap=False):
nn.Module.__init__(self)
self.d = d - (d // 2)
self.swap = swap
self.net_s_t = nn.Sequential(
nn.Linear(self.d, intermediate_dim),
nn.ReLU(inplace=True),
nn.Linear(intermediate_dim, intermediate_dim),
nn.ReLU(inplace=True),
nn.Linear(intermediate_dim, (d - self.d) * 2),
)
def forward(self, x, logpx=None, reverse=False):
if self.swap:
x = torch.cat([x[:, self.d:], x[:, :self.d]], 1)
in_dim = self.d
out_dim = x.shape[1] - self.d
s_t = self.net_s_t(x[:, :in_dim])
scale = torch.sigmoid(s_t[:, :out_dim] + 2.)
shift = s_t[:, out_dim:]
logdetjac = torch.sum(torch.log(scale).view(scale.shape[0], -1), 1, keepdim=True)
if not reverse:
y1 = x[:, self.d:] * scale + shift
delta_logp = -logdetjac
else:
y1 = (x[:, self.d:] - shift) / scale
delta_logp = logdetjac
y = torch.cat([x[:, :self.d], y1], 1) if not self.swap else torch.cat([y1, x[:, :self.d]], 1)
if logpx is None:
return y
else:
return y, logpx + delta_logp
class SequentialFlow(nn.Module):
"""A generalized nn.Sequential container for normalizing flows.
"""
def __init__(self, layersList):
super(SequentialFlow, self).__init__()
self.chain = nn.ModuleList(layersList)
def forward(self, x, logpx=None, reverse=False, inds=None):
if inds is None:
if reverse:
inds = range(len(self.chain) - 1, -1, -1)
else:
inds = range(len(self.chain))
if logpx is None:
for i in inds:
x = self.chain[i](x, reverse=reverse)
return x
else:
for i in inds:
x, logpx = self.chain[i](x, logpx, reverse=reverse)
return x, logpx
def build_latent_flow(args):
chain = []
for i in range(args.latent_flow_depth):
chain.append(CouplingLayer(args.latent_dim, args.latent_flow_hidden_dim, swap=(i % 2 == 0)))
return SequentialFlow(chain)
##################
## SpectralNorm ##
##################
POWER_ITERATION_FN = "spectral_norm_power_iteration"
class SpectralNorm(object):
def __init__(self, name='weight', dim=0, eps=1e-12):
self.name = name
self.dim = dim
self.eps = eps
def compute_weight(self, module, n_power_iterations):
if n_power_iterations < 0:
raise ValueError(
'Expected n_power_iterations to be non-negative, but '
'got n_power_iterations={}'.format(n_power_iterations)
)
weight = getattr(module, self.name + '_orig')
u = getattr(module, self.name + '_u')
v = getattr(module, self.name + '_v')
weight_mat = weight
if self.dim != 0:
# permute dim to front
weight_mat = weight_mat.permute(self.dim, * [d for d in range(weight_mat.dim()) if d != self.dim])
height = weight_mat.size(0)
weight_mat = weight_mat.reshape(height, -1)
with torch.no_grad():
for _ in range(n_power_iterations):
# Spectral norm of weight equals to `u^T W v`, where `u` and `v`
# are the first left and right singular vectors.
# This power iteration produces approximations of `u` and `v`.
v = F.normalize(torch.matmul(weight_mat.t(), u), dim=0, eps=self.eps)
u = F.normalize(torch.matmul(weight_mat, v), dim=0, eps=self.eps)
setattr(module, self.name + '_u', u)
setattr(module, self.name + '_v', v)
sigma = torch.dot(u, torch.matmul(weight_mat, v))
weight = weight / sigma
setattr(module, self.name, weight)
def remove(self, module):
weight = getattr(module, self.name)
delattr(module, self.name)
delattr(module, self.name + '_u')
delattr(module, self.name + '_orig')
module.register_parameter(self.name, torch.nn.Parameter(weight))
def get_update_method(self, module):
def update_fn(module, n_power_iterations):
self.compute_weight(module, n_power_iterations)
return update_fn
def __call__(self, module, unused_inputs):
del unused_inputs
self.compute_weight(module, n_power_iterations=0)
# requires_grad might be either True or False during inference.
if not module.training:
r_g = getattr(module, self.name + '_orig').requires_grad
setattr(module, self.name, getattr(module, self.name).detach().requires_grad_(r_g))
@staticmethod
def apply(module, name, dim, eps):
fn = SpectralNorm(name, dim, eps)
weight = module._parameters[name]
height = weight.size(dim)
u = F.normalize(weight.new_empty(height).normal_(0, 1), dim=0, eps=fn.eps)
v = F.normalize(weight.new_empty(int(weight.numel() / height)).normal_(0, 1), dim=0, eps=fn.eps)
delattr(module, fn.name)
module.register_parameter(fn.name + "_orig", weight)
# We still need to assign weight back as fn.name because all sorts of
# things may assume that it exists, e.g., when initializing weights.
# However, we can't directly assign as it could be an nn.Parameter and
# gets added as a parameter. Instead, we register weight.data as a
# buffer, which will cause weight to be included in the state dict
# and also supports nn.init due to shared storage.
module.register_buffer(fn.name, weight.data)
module.register_buffer(fn.name + "_u", u)
module.register_buffer(fn.name + "_v", v)
setattr(module, POWER_ITERATION_FN, types.MethodType(fn.get_update_method(module), module))
module.register_forward_pre_hook(fn)
return fn
def inplace_spectral_norm(module, name='weight', dim=None, eps=1e-12):
r"""Applies spectral normalization to a parameter in the given module.
.. math::
\mathbf{W} = \dfrac{\mathbf{W}}{\sigma(\mathbf{W})} \\
\sigma(\mathbf{W}) = \max_{\mathbf{h}: \mathbf{h} \ne 0} \dfrac{\|\mathbf{W} \mathbf{h}\|_2}{\|\mathbf{h}\|_2}
Spectral normalization stabilizes the training of discriminators (critics)
in Generaive Adversarial Networks (GANs) by rescaling the weight tensor
with spectral norm :math:`\sigma` of the weight matrix calculated using
power iteration method. If the dimension of the weight tensor is greater
than 2, it is reshaped to 2D in power iteration method to get spectral
norm. This is implemented via a hook that calculates spectral norm and
rescales weight before every :meth:`~Module.forward` call.
See `Spectral Normalization for Generative Adversarial Networks`_ .
.. _`Spectral Normalization for Generative Adversarial Networks`: https://arxiv.org/abs/1802.05957
Args:
module (nn.Module): containing module
name (str, optional): name of weight parameter
n_power_iterations (int, optional): number of power iterations to
calculate spectal norm
dim (int, optional): dimension corresponding to number of outputs,
the default is 0, except for modules that are instances of
ConvTranspose1/2/3d, when it is 1
eps (float, optional): epsilon for numerical stability in
calculating norms
Returns:
The original module with the spectal norm hook
Example::
>>> m = spectral_norm(nn.Linear(20, 40))
Linear (20 -> 40)
>>> m.weight_u.size()
torch.Size([20])
"""
if dim is None:
if isinstance(module, (torch.nn.ConvTranspose1d, torch.nn.ConvTranspose2d, torch.nn.ConvTranspose3d)):
dim = 1
else:
dim = 0
SpectralNorm.apply(module, name, dim=dim, eps=eps)
return module
def remove_spectral_norm(module, name='weight'):
r"""Removes the spectral normalization reparameterization from a module.
Args:
module (nn.Module): containing module
name (str, optional): name of weight parameter
Example:
>>> m = spectral_norm(nn.Linear(40, 10))
>>> remove_spectral_norm(m)
"""
for k, hook in module._forward_pre_hooks.items():
if isinstance(hook, SpectralNorm) and hook.name == name:
hook.remove(module)
del module._forward_pre_hooks[k]
return module
raise ValueError("spectral_norm of '{}' not found in {}".format(name, module))
def add_spectral_norm(model, logger=None):
"""Applies spectral norm to all modules within the scope of a CNF."""
def apply_spectral_norm(module):
if 'weight' in module._parameters:
if logger: logger.info("Adding spectral norm to {}".format(module))
inplace_spectral_norm(module, 'weight')
def find_coupling_layer(module):
if isinstance(module, CouplingLayer):
module.apply(apply_spectral_norm)
else:
for child in module.children():
find_coupling_layer(child)
find_coupling_layer(model)
def spectral_norm_power_iteration(model, n_power_iterations=1):
def recursive_power_iteration(module):
if hasattr(module, POWER_ITERATION_FN):
getattr(module, POWER_ITERATION_FN)(n_power_iterations)
model.apply(recursive_power_iteration)
def reparameterize_gaussian(mean, logvar):
std = torch.exp(0.5 * logvar)
eps = torch.randn(std.size()).to(mean)
return mean + std * eps
def gaussian_entropy(logvar):
const = 0.5 * float(logvar.size(1)) * (1. + np.log(np.pi * 2))
ent = 0.5 * logvar.sum(dim=1, keepdim=False) + const
return ent
def standard_normal_logprob(z):
dim = z.size(-1)
log_z = -0.5 * dim * np.log(2 * np.pi)
return log_z - z.pow(2) / 2
def truncated_normal_(tensor, mean=0, std=1, trunc_std=2):
"""
Taken from https://discuss.pytorch.org/t/implementing-truncated-normal-initializer/4778/15
"""
size = tensor.shape
tmp = tensor.new_empty(size + (4,)).normal_()
valid = (tmp < trunc_std) & (tmp > -trunc_std)
ind = valid.max(-1, keepdim=True)[1]
tensor.data.copy_(tmp.gather(-1, ind).squeeze(-1))
tensor.data.mul_(std).add_(mean)
return tensor
tgs/models/tokenizers/dinov2.py
+1175
View File
File diff suppressed because it is too large Load Diff
tgs/models/tokenizers/image.py
+135
View File
@@ -0,0 +1,135 @@
from dataclasses import dataclass
import torch
import torch.nn as nn
from einops import rearrange
from tgs.utils.base import BaseModule
from tgs.models.tokenizers.dinov2 import Dinov2Model
from tgs.models.transformers import Modulation
from tgs.utils.typing import *
class DINOV2SingleImageTokenizer(BaseModule):
@dataclass
class Config(BaseModule.Config):
pretrained_model_name_or_path: str = "facebook/dinov2-base"
width: int = 224
height: int = 224
modulation: bool = False
modulation_zero_init: bool = False
modulation_single_layer: bool = False
modulation_cond_dim: int = 16
freeze_backbone_params: bool = True
enable_memory_efficient_attention: bool = False
enable_gradient_checkpointing: bool = False
use_patch_embeddings: bool = False
patch_embeddings_aggr_method: str = 'concat'
cfg: Config
def configure(self) -> None:
super().configure()
model: Dinov2Model
if self.cfg.freeze_backbone_params:
# freeze dino backbone parameters
self.register_non_module(
"model",
Dinov2Model.from_pretrained(self.cfg.pretrained_model_name_or_path).to(
self.device
),
)
model = self.non_module("model")
for p in model.parameters():
p.requires_grad_(False)
model.eval()
else:
self.model = Dinov2Model.from_pretrained(
self.cfg.pretrained_model_name_or_path
).to(self.device)
model = self.model
model.set_use_memory_efficient_attention_xformers(
self.cfg.enable_memory_efficient_attention
)
model.set_gradient_checkpointing(self.cfg.enable_gradient_checkpointing)
# add modulation
if self.cfg.modulation:
modulations = []
for layer in model.encoder.layer:
norm1_modulation = Modulation(
model.config.hidden_size,
self.cfg.modulation_cond_dim,
zero_init=self.cfg.modulation_zero_init,
single_layer=self.cfg.modulation_single_layer,
)
norm2_modulation = Modulation(
model.config.hidden_size,
self.cfg.modulation_cond_dim,
zero_init=self.cfg.modulation_zero_init,
single_layer=self.cfg.modulation_single_layer,
)
layer.register_ada_norm_modulation(norm1_modulation, norm2_modulation)
modulations += [norm1_modulation, norm2_modulation]
self.modulations = nn.ModuleList(modulations)
self.register_buffer(
"image_mean",
torch.as_tensor([0.485, 0.456, 0.406]).reshape(1, 1, 3, 1, 1),
persistent=False,
)
self.register_buffer(
"image_std",
torch.as_tensor([0.229, 0.224, 0.225]).reshape(1, 1, 3, 1, 1),
persistent=False,
)
def forward(
self,
images: Float[Tensor, "B *N C H W"],
modulation_cond: Optional[Float[Tensor, "B *N Cc"]],
) -> Float[Tensor, "B *N Ct Nt"]:
model: Dinov2Model
if self.cfg.freeze_backbone_params:
model = self.non_module("model")
else:
model = self.model
packed = False
if images.ndim == 4:
packed = True
images = images.unsqueeze(1)
if modulation_cond is not None:
assert modulation_cond.ndim == 2
modulation_cond = modulation_cond.unsqueeze(1)
batch_size, n_input_views = images.shape[:2]
images = (images - self.image_mean) / self.image_std
out = model(
rearrange(images, "B N C H W -> (B N) C H W"),
modulation_cond=rearrange(modulation_cond, "B N Cc -> (B N) Cc")
if modulation_cond is not None
else None,
)
local_features, global_features = out.last_hidden_state, out.pooler_output
if self.cfg.use_patch_embeddings:
patch_embeddings = out.patch_embeddings
if self.cfg.patch_embeddings_aggr_method == 'concat':
local_features = torch.cat([local_features, patch_embeddings], dim=1)
elif self.cfg.patch_embeddings_aggr_method == 'add':
local_features = local_features + patch_embeddings
else:
raise NotImplementedError
local_features = local_features.permute(0, 2, 1)
local_features = rearrange(
local_features, "(B N) Ct Nt -> B N Ct Nt", B=batch_size
)
if packed:
local_features = local_features.squeeze(1)
return local_features
def detokenize(self, *args, **kwargs):
raise NotImplementedError
tgs/models/tokenizers/point.py
+29
View File
@@ -0,0 +1,29 @@
from dataclasses import dataclass
import torch.nn as nn
from tgs.utils.base import BaseModule
from tgs.utils.typing import *
import torch
class PointLearnablePositionalEmbedding(BaseModule):
@dataclass
class Config(BaseModule.Config):
num_pcl: int = 2048
num_channels: int = 512
cfg: Config
def configure(self) -> None:
super().configure()
self.pcl_embeddings = nn.Embedding(
self.cfg.num_pcl , self.cfg.num_channels
)
def forward(self, batch_size: int) -> Float[Tensor, "B Ct Nt"]:
range_ = torch.arange(self.cfg.num_pcl, device=self.device)
embeddings = self.pcl_embeddings(range_).unsqueeze(0).repeat((batch_size,1,1))
return torch.permute(embeddings, (0,2,1))
def detokenize(
self, tokens: Float[Tensor, "B Ct Nt"]
) -> Float[Tensor, "B 3 Ct Hp Wp"]:
return torch.permute(tokens, (0,2,1))
tgs/models/tokenizers/triplane.py
+52
View File
@@ -0,0 +1,52 @@
from dataclasses import dataclass
import math
import torch
import torch.nn as nn
from einops import rearrange, repeat
from tgs.utils.base import BaseModule
from tgs.utils.typing import *
class TriplaneLearnablePositionalEmbedding(BaseModule):
@dataclass
class Config(BaseModule.Config):
plane_size: int = 32
num_channels: int = 1024
cfg: Config
def configure(self) -> None:
super().configure()
self.embeddings = nn.Parameter(
torch.randn(
(3, self.cfg.num_channels, self.cfg.plane_size, self.cfg.plane_size),
dtype=torch.float32,
)
* 1
/ math.sqrt(self.cfg.num_channels)
)
def forward(self, batch_size: int, cond_embeddings: Float[Tensor, "B Ct"] = None) -> Float[Tensor, "B Ct Nt"]:
embeddings = repeat(self.embeddings, "Np Ct Hp Wp -> B Np Ct Hp Wp", B=batch_size)
if cond_embeddings is not None:
embeddings = embeddings + cond_embeddings
return rearrange(
embeddings,
"B Np Ct Hp Wp -> B Ct (Np Hp Wp)",
)
def detokenize(
self, tokens: Float[Tensor, "B Ct Nt"]
) -> Float[Tensor, "B 3 Ct Hp Wp"]:
batch_size, Ct, Nt = tokens.shape
assert Nt == self.cfg.plane_size**2 * 3
assert Ct == self.cfg.num_channels
return rearrange(
tokens,
"B Ct (Np Hp Wp) -> B Np Ct Hp Wp",
Np=3,
Hp=self.cfg.plane_size,
Wp=self.cfg.plane_size,
)
tgs/models/transformers.py
+908
View File
@@ -0,0 +1,908 @@
# Copyright 2023 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import torch
import torch.nn.functional as F
from torch import nn
from diffusers.utils.torch_utils import maybe_allow_in_graph
from diffusers.models.activations import get_activation
from diffusers.models.attention_processor import Attention
from diffusers.models.embeddings import CombinedTimestepLabelEmbeddings
from dataclasses import dataclass
from ..utils.base import BaseModule
from ..utils.typing import *
class MemoryEfficientAttentionMixin:
def enable_xformers_memory_efficient_attention(
self, attention_op: Optional[Callable] = None
):
r"""
Enable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/). When this
option is enabled, you should observe lower GPU memory usage and a potential speed up during inference. Speed
up during training is not guaranteed.
<Tip warning={true}>
⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes
precedent.
</Tip>
Parameters:
attention_op (`Callable`, *optional*):
Override the default `None` operator for use as `op` argument to the
[`memory_efficient_attention()`](https://facebookresearch.github.io/xformers/components/ops.html#xformers.ops.memory_efficient_attention)
function of xFormers.
Examples:
```py
>>> import torch
>>> from diffusers import DiffusionPipeline
>>> from xformers.ops import MemoryEfficientAttentionFlashAttentionOp
>>> pipe = DiffusionPipeline.from_pretrained("stabilityai/stable-diffusion-2-1", torch_dtype=torch.float16)
>>> pipe = pipe.to("cuda")
>>> pipe.enable_xformers_memory_efficient_attention(attention_op=MemoryEfficientAttentionFlashAttentionOp)
>>> # Workaround for not accepting attention shape using VAE for Flash Attention
>>> pipe.vae.enable_xformers_memory_efficient_attention(attention_op=None)
```
"""
self.set_use_memory_efficient_attention_xformers(True, attention_op)
def disable_xformers_memory_efficient_attention(self):
r"""
Disable memory efficient attention from [xFormers](https://facebookresearch.github.io/xformers/).
"""
self.set_use_memory_efficient_attention_xformers(False)
def set_use_memory_efficient_attention_xformers(
self, valid: bool, attention_op: Optional[Callable] = None
) -> None:
# Recursively walk through all the children.
# Any children which exposes the set_use_memory_efficient_attention_xformers method
# gets the message
def fn_recursive_set_mem_eff(module: torch.nn.Module):
if hasattr(module, "set_use_memory_efficient_attention_xformers"):
module.set_use_memory_efficient_attention_xformers(valid, attention_op)
for child in module.children():
fn_recursive_set_mem_eff(child)
for module in self.children():
if isinstance(module, torch.nn.Module):
fn_recursive_set_mem_eff(module)
@maybe_allow_in_graph
class GatedSelfAttentionDense(nn.Module):
r"""
A gated self-attention dense layer that combines visual features and object features.
Parameters:
query_dim (`int`): The number of channels in the query.
context_dim (`int`): The number of channels in the context.
n_heads (`int`): The number of heads to use for attention.
d_head (`int`): The number of channels in each head.
"""
def __init__(self, query_dim: int, context_dim: int, n_heads: int, d_head: int):
super().__init__()
# we need a linear projection since we need cat visual feature and obj feature
self.linear = nn.Linear(context_dim, query_dim)
self.attn = Attention(query_dim=query_dim, heads=n_heads, dim_head=d_head)
self.ff = FeedForward(query_dim, activation_fn="geglu")
self.norm1 = nn.LayerNorm(query_dim)
self.norm2 = nn.LayerNorm(query_dim)
self.register_parameter("alpha_attn", nn.Parameter(torch.tensor(0.0)))
self.register_parameter("alpha_dense", nn.Parameter(torch.tensor(0.0)))
self.enabled = True
def forward(self, x: torch.Tensor, objs: torch.Tensor) -> torch.Tensor:
if not self.enabled:
return x
n_visual = x.shape[1]
objs = self.linear(objs)
x = (
x
+ self.alpha_attn.tanh()
* self.attn(self.norm1(torch.cat([x, objs], dim=1)))[:, :n_visual, :]
)
x = x + self.alpha_dense.tanh() * self.ff(self.norm2(x))
return x
@maybe_allow_in_graph
class BasicTransformerBlock(nn.Module, MemoryEfficientAttentionMixin):
r"""
A basic Transformer block.
Parameters:
dim (`int`): The number of channels in the input and output.
num_attention_heads (`int`): The number of heads to use for multi-head attention.
attention_head_dim (`int`): The number of channels in each head.
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
cross_attention_dim (`int`, *optional*): The size of the encoder_hidden_states vector for cross attention.
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
num_embeds_ada_norm (:
obj: `int`, *optional*): The number of diffusion steps used during training. See `Transformer2DModel`.
attention_bias (:
obj: `bool`, *optional*, defaults to `False`): Configure if the attentions should contain a bias parameter.
only_cross_attention (`bool`, *optional*):
Whether to use only cross-attention layers. In this case two cross attention layers are used.
double_self_attention (`bool`, *optional*):
Whether to use two self-attention layers. In this case no cross attention layers are used.
upcast_attention (`bool`, *optional*):
Whether to upcast the attention computation to float32. This is useful for mixed precision training.
norm_elementwise_affine (`bool`, *optional*, defaults to `True`):
Whether to use learnable elementwise affine parameters for normalization.
norm_type (`str`, *optional*, defaults to `"layer_norm"`):
The normalization layer to use. Can be `"layer_norm"`, `"ada_norm"` or `"ada_norm_zero"`.
final_dropout (`bool` *optional*, defaults to False):
Whether to apply a final dropout after the last feed-forward layer.
attention_type (`str`, *optional*, defaults to `"default"`):
The type of attention to use. Can be `"default"` or `"gated"` or `"gated-text-image"`.
"""
def __init__(
self,
dim: int,
num_attention_heads: int,
attention_head_dim: int,
dropout=0.0,
cross_attention_dim: Optional[int] = None,
activation_fn: str = "geglu",
num_embeds_ada_norm: Optional[int] = None,
cond_dim_ada_norm_continuous: Optional[int] = None,
attention_bias: bool = False,
only_cross_attention: bool = False,
double_self_attention: bool = False,
upcast_attention: bool = False,
norm_elementwise_affine: bool = True,
norm_type: str = "layer_norm",
final_dropout: bool = False,
attention_type: str = "default",
):
super().__init__()
self.only_cross_attention = only_cross_attention
self.use_ada_layer_norm_zero = (
num_embeds_ada_norm is not None
) and norm_type == "ada_norm_zero"
self.use_ada_layer_norm = (
num_embeds_ada_norm is not None
) and norm_type == "ada_norm"
self.use_ada_layer_norm_continuous = (
cond_dim_ada_norm_continuous is not None
) and norm_type == "ada_norm_continuous"
assert (
int(self.use_ada_layer_norm)
+ int(self.use_ada_layer_norm_continuous)
+ int(self.use_ada_layer_norm_zero)
<= 1
)
if norm_type in ("ada_norm", "ada_norm_zero") and num_embeds_ada_norm is None:
raise ValueError(
f"`norm_type` is set to {norm_type}, but `num_embeds_ada_norm` is not defined. Please make sure to"
f" define `num_embeds_ada_norm` if setting `norm_type` to {norm_type}."
)
# Define 3 blocks. Each block has its own normalization layer.
# 1. Self-Attn
if self.use_ada_layer_norm:
self.norm1 = AdaLayerNorm(dim, num_embeds_ada_norm)
elif self.use_ada_layer_norm_continuous:
self.norm1 = AdaLayerNormContinuous(dim, cond_dim_ada_norm_continuous)
elif self.use_ada_layer_norm_zero:
self.norm1 = AdaLayerNormZero(dim, num_embeds_ada_norm)
else:
self.norm1 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
self.attn1 = Attention(
query_dim=dim,
heads=num_attention_heads,
dim_head=attention_head_dim,
dropout=dropout,
bias=attention_bias,
cross_attention_dim=cross_attention_dim if only_cross_attention else None,
upcast_attention=upcast_attention,
)
# 2. Cross-Attn
if cross_attention_dim is not None or double_self_attention:
# We currently only use AdaLayerNormZero for self attention where there will only be one attention block.
# I.e. the number of returned modulation chunks from AdaLayerZero would not make sense if returned during
# the second cross attention block.
if self.use_ada_layer_norm:
self.norm2 = AdaLayerNorm(dim, num_embeds_ada_norm)
elif self.use_ada_layer_norm_continuous:
self.norm2 = AdaLayerNormContinuous(dim, cond_dim_ada_norm_continuous)
else:
self.norm2 = nn.LayerNorm(
dim, elementwise_affine=norm_elementwise_affine
)
self.attn2 = Attention(
query_dim=dim,
cross_attention_dim=cross_attention_dim
if not double_self_attention
else None,
heads=num_attention_heads,
dim_head=attention_head_dim,
dropout=dropout,
bias=attention_bias,
upcast_attention=upcast_attention,
) # is self-attn if encoder_hidden_states is none
else:
self.norm2 = None
self.attn2 = None
# 3. Feed-forward
if self.use_ada_layer_norm_continuous:
self.norm3 = AdaLayerNormContinuous(dim, cond_dim_ada_norm_continuous)
else:
self.norm3 = nn.LayerNorm(dim, elementwise_affine=norm_elementwise_affine)
self.ff = FeedForward(
dim,
dropout=dropout,
activation_fn=activation_fn,
final_dropout=final_dropout,
)
# 4. Fuser
if attention_type == "gated" or attention_type == "gated-text-image":
self.fuser = GatedSelfAttentionDense(
dim, cross_attention_dim, num_attention_heads, attention_head_dim
)
# let chunk size default to None
self._chunk_size = None
self._chunk_dim = 0
def set_chunk_feed_forward(self, chunk_size: Optional[int], dim: int):
# Sets chunk feed-forward
self._chunk_size = chunk_size
self._chunk_dim = dim
def forward(
self,
hidden_states: torch.FloatTensor,
attention_mask: Optional[torch.FloatTensor] = None,
encoder_hidden_states: Optional[torch.FloatTensor] = None,
encoder_attention_mask: Optional[torch.FloatTensor] = None,
timestep: Optional[torch.LongTensor] = None,
modulation_cond: Optional[torch.FloatTensor] = None,
cross_attention_kwargs: Dict[str, Any] = None,
class_labels: Optional[torch.LongTensor] = None,
) -> torch.FloatTensor:
# Notice that normalization is always applied before the real computation in the following blocks.
# 0. Self-Attention
if self.use_ada_layer_norm:
norm_hidden_states = self.norm1(hidden_states, timestep)
elif self.use_ada_layer_norm_continuous:
norm_hidden_states = self.norm1(hidden_states, modulation_cond)
elif self.use_ada_layer_norm_zero:
norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.norm1(
hidden_states, timestep, class_labels, hidden_dtype=hidden_states.dtype
)
else:
norm_hidden_states = self.norm1(hidden_states)
# 1. Retrieve lora scale.
lora_scale = (
cross_attention_kwargs.get("scale", 1.0)
if cross_attention_kwargs is not None
else 1.0
)
# 2. Prepare GLIGEN inputs
cross_attention_kwargs = (
cross_attention_kwargs.copy() if cross_attention_kwargs is not None else {}
)
gligen_kwargs = cross_attention_kwargs.pop("gligen", None)
attn_output = self.attn1(
norm_hidden_states,
encoder_hidden_states=encoder_hidden_states
if self.only_cross_attention
else None,
attention_mask=attention_mask,
**cross_attention_kwargs,
)
if self.use_ada_layer_norm_zero:
attn_output = gate_msa.unsqueeze(1) * attn_output
hidden_states = attn_output + hidden_states
# 2.5 GLIGEN Control
if gligen_kwargs is not None:
hidden_states = self.fuser(hidden_states, gligen_kwargs["objs"])
# 2.5 ends
# 3. Cross-Attention
if self.attn2 is not None:
if self.use_ada_layer_norm:
norm_hidden_states = self.norm2(hidden_states, timestep)
elif self.use_ada_layer_norm_continuous:
norm_hidden_states = self.norm2(hidden_states, modulation_cond)
else:
norm_hidden_states = self.norm2(hidden_states)
attn_output = self.attn2(
norm_hidden_states,
encoder_hidden_states=encoder_hidden_states,
attention_mask=encoder_attention_mask,
**cross_attention_kwargs,
)
hidden_states = attn_output + hidden_states
# 4. Feed-forward
if self.use_ada_layer_norm_continuous:
norm_hidden_states = self.norm3(hidden_states, modulation_cond)
else:
norm_hidden_states = self.norm3(hidden_states)
if self.use_ada_layer_norm_zero:
norm_hidden_states = (
norm_hidden_states * (1 + scale_mlp[:, None]) + shift_mlp[:, None]
)
if self._chunk_size is not None:
# "feed_forward_chunk_size" can be used to save memory
if norm_hidden_states.shape[self._chunk_dim] % self._chunk_size != 0:
raise ValueError(
f"`hidden_states` dimension to be chunked: {norm_hidden_states.shape[self._chunk_dim]} has to be divisible by chunk size: {self._chunk_size}. Make sure to set an appropriate `chunk_size` when calling `unet.enable_forward_chunking`."
)
num_chunks = norm_hidden_states.shape[self._chunk_dim] // self._chunk_size
ff_output = torch.cat(
[
self.ff(hid_slice, scale=lora_scale)
for hid_slice in norm_hidden_states.chunk(
num_chunks, dim=self._chunk_dim
)
],
dim=self._chunk_dim,
)
else:
ff_output = self.ff(norm_hidden_states, scale=lora_scale)
if self.use_ada_layer_norm_zero:
ff_output = gate_mlp.unsqueeze(1) * ff_output
hidden_states = ff_output + hidden_states
return hidden_states
class FeedForward(nn.Module):
r"""
A feed-forward layer.
Parameters:
dim (`int`): The number of channels in the input.
dim_out (`int`, *optional*): The number of channels in the output. If not given, defaults to `dim`.
mult (`int`, *optional*, defaults to 4): The multiplier to use for the hidden dimension.
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to be used in feed-forward.
final_dropout (`bool` *optional*, defaults to False): Apply a final dropout.
"""
def __init__(
self,
dim: int,
dim_out: Optional[int] = None,
mult: int = 4,
dropout: float = 0.0,
activation_fn: str = "geglu",
final_dropout: bool = False,
):
super().__init__()
inner_dim = int(dim * mult)
dim_out = dim_out if dim_out is not None else dim
linear_cls = nn.Linear
if activation_fn == "gelu":
act_fn = GELU(dim, inner_dim)
if activation_fn == "gelu-approximate":
act_fn = GELU(dim, inner_dim, approximate="tanh")
elif activation_fn == "geglu":
act_fn = GEGLU(dim, inner_dim)
elif activation_fn == "geglu-approximate":
act_fn = ApproximateGELU(dim, inner_dim)
self.net = nn.ModuleList([])
# project in
self.net.append(act_fn)
# project dropout
self.net.append(nn.Dropout(dropout))
# project out
self.net.append(linear_cls(inner_dim, dim_out))
# FF as used in Vision Transformer, MLP-Mixer, etc. have a final dropout
if final_dropout:
self.net.append(nn.Dropout(dropout))
def forward(self, hidden_states: torch.Tensor, scale: float = 1.0) -> torch.Tensor:
for module in self.net:
hidden_states = module(hidden_states)
return hidden_states
class GELU(nn.Module):
r"""
GELU activation function with tanh approximation support with `approximate="tanh"`.
Parameters:
dim_in (`int`): The number of channels in the input.
dim_out (`int`): The number of channels in the output.
approximate (`str`, *optional*, defaults to `"none"`): If `"tanh"`, use tanh approximation.
"""
def __init__(self, dim_in: int, dim_out: int, approximate: str = "none"):
super().__init__()
self.proj = nn.Linear(dim_in, dim_out)
self.approximate = approximate
def gelu(self, gate: torch.Tensor) -> torch.Tensor:
if gate.device.type != "mps":
return F.gelu(gate, approximate=self.approximate)
# mps: gelu is not implemented for float16
return F.gelu(gate.to(dtype=torch.float32), approximate=self.approximate).to(
dtype=gate.dtype
)
def forward(self, hidden_states):
hidden_states = self.proj(hidden_states)
hidden_states = self.gelu(hidden_states)
return hidden_states
class GEGLU(nn.Module):
r"""
A variant of the gated linear unit activation function from https://arxiv.org/abs/2002.05202.
Parameters:
dim_in (`int`): The number of channels in the input.
dim_out (`int`): The number of channels in the output.
"""
def __init__(self, dim_in: int, dim_out: int):
super().__init__()
linear_cls = nn.Linear
self.proj = linear_cls(dim_in, dim_out * 2)
def gelu(self, gate: torch.Tensor) -> torch.Tensor:
if gate.device.type != "mps":
return F.gelu(gate)
# mps: gelu is not implemented for float16
return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
def forward(self, hidden_states, scale: float = 1.0):
args = ()
hidden_states, gate = self.proj(hidden_states, *args).chunk(2, dim=-1)
return hidden_states * self.gelu(gate)
class ApproximateGELU(nn.Module):
r"""
The approximate form of Gaussian Error Linear Unit (GELU). For more details, see section 2:
https://arxiv.org/abs/1606.08415.
Parameters:
dim_in (`int`): The number of channels in the input.
dim_out (`int`): The number of channels in the output.
"""
def __init__(self, dim_in: int, dim_out: int):
super().__init__()
self.proj = nn.Linear(dim_in, dim_out)
def forward(self, x: torch.Tensor) -> torch.Tensor:
x = self.proj(x)
return x * torch.sigmoid(1.702 * x)
class AdaLayerNorm(nn.Module):
r"""
Norm layer modified to incorporate timestep embeddings.
Parameters:
embedding_dim (`int`): The size of each embedding vector.
num_embeddings (`int`): The size of the dictionary of embeddings.
"""
def __init__(self, embedding_dim: int, num_embeddings: int):
super().__init__()
self.emb = nn.Embedding(num_embeddings, embedding_dim)
self.silu = nn.SiLU()
self.linear = nn.Linear(embedding_dim, embedding_dim * 2)
self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False)
def forward(self, x: torch.Tensor, timestep: torch.Tensor) -> torch.Tensor:
emb = self.linear(self.silu(self.emb(timestep)))
scale, shift = torch.chunk(emb, 2, dim=1)
x = self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
return x
class AdaLayerNormContinuous(nn.Module):
r"""
Norm layer modified to incorporate arbitrary continuous embeddings.
Parameters:
embedding_dim (`int`): The size of each embedding vector.
"""
def __init__(self, embedding_dim: int, condition_dim: int):
super().__init__()
self.silu = nn.SiLU()
self.linear1 = nn.Linear(condition_dim, condition_dim)
self.linear2 = nn.Linear(condition_dim, embedding_dim * 2)
self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False)
def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
emb = self.linear2(self.silu(self.linear1(condition)))
scale, shift = torch.chunk(emb, 2, dim=1)
x = self.norm(x) * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
return x
class Modulation(nn.Module):
def __init__(self, embedding_dim: int, condition_dim: int, zero_init: bool = False, single_layer: bool = False):
super().__init__()
self.silu = nn.SiLU()
if single_layer:
self.linear1 = nn.Identity()
else:
self.linear1 = nn.Linear(condition_dim, condition_dim)
self.linear2 = nn.Linear(condition_dim, embedding_dim * 2)
# Only zero init the last linear layer
if zero_init:
nn.init.zeros_(self.linear2.weight)
nn.init.zeros_(self.linear2.bias)
def forward(self, x: torch.Tensor, condition: torch.Tensor) -> torch.Tensor:
emb = self.linear2(self.silu(self.linear1(condition)))
scale, shift = torch.chunk(emb, 2, dim=1)
x = x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
return x
class AdaLayerNormZero(nn.Module):
r"""
Norm layer adaptive layer norm zero (adaLN-Zero).
Parameters:
embedding_dim (`int`): The size of each embedding vector.
num_embeddings (`int`): The size of the dictionary of embeddings.
"""
def __init__(self, embedding_dim: int, num_embeddings: int):
super().__init__()
self.emb = CombinedTimestepLabelEmbeddings(num_embeddings, embedding_dim)
self.silu = nn.SiLU()
self.linear = nn.Linear(embedding_dim, 6 * embedding_dim, bias=True)
self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6)
def forward(
self,
x: torch.Tensor,
timestep: torch.Tensor,
class_labels: torch.LongTensor,
hidden_dtype: Optional[torch.dtype] = None,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
emb = self.linear(
self.silu(self.emb(timestep, class_labels, hidden_dtype=hidden_dtype))
)
shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(
6, dim=1
)
x = self.norm(x) * (1 + scale_msa[:, None]) + shift_msa[:, None]
return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
class AdaGroupNorm(nn.Module):
r"""
GroupNorm layer modified to incorporate timestep embeddings.
Parameters:
embedding_dim (`int`): The size of each embedding vector.
num_embeddings (`int`): The size of the dictionary of embeddings.
num_groups (`int`): The number of groups to separate the channels into.
act_fn (`str`, *optional*, defaults to `None`): The activation function to use.
eps (`float`, *optional*, defaults to `1e-5`): The epsilon value to use for numerical stability.
"""
def __init__(
self,
embedding_dim: int,
out_dim: int,
num_groups: int,
act_fn: Optional[str] = None,
eps: float = 1e-5,
):
super().__init__()
self.num_groups = num_groups
self.eps = eps
if act_fn is None:
self.act = None
else:
self.act = get_activation(act_fn)
self.linear = nn.Linear(embedding_dim, out_dim * 2)
def forward(self, x: torch.Tensor, emb: torch.Tensor) -> torch.Tensor:
if self.act:
emb = self.act(emb)
emb = self.linear(emb)
emb = emb[:, :, None, None]
scale, shift = emb.chunk(2, dim=1)
x = F.group_norm(x, self.num_groups, eps=self.eps)
x = x * (1 + scale) + shift
return x
class Transformer1D(BaseModule, MemoryEfficientAttentionMixin):
"""
A 1D Transformer model for sequence data.
Parameters:
num_attention_heads (`int`, *optional*, defaults to 16): The number of heads to use for multi-head attention.
attention_head_dim (`int`, *optional*, defaults to 88): The number of channels in each head.
in_channels (`int`, *optional*):
The number of channels in the input and output (specify if the input is **continuous**).
num_layers (`int`, *optional*, defaults to 1): The number of layers of Transformer blocks to use.
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
cross_attention_dim (`int`, *optional*): The number of `encoder_hidden_states` dimensions to use.
activation_fn (`str`, *optional*, defaults to `"geglu"`): Activation function to use in feed-forward.
num_embeds_ada_norm ( `int`, *optional*):
The number of diffusion steps used during training. Pass if at least one of the norm_layers is
`AdaLayerNorm`. This is fixed during training since it is used to learn a number of embeddings that are
added to the hidden states.
During inference, you can denoise for up to but not more steps than `num_embeds_ada_norm`.
attention_bias (`bool`, *optional*):
Configure if the `TransformerBlocks` attention should contain a bias parameter.
"""
@dataclass
class Config(BaseModule.Config):
num_attention_heads: int = 16
attention_head_dim: int = 88
in_channels: Optional[int] = None
out_channels: Optional[int] = None
num_layers: int = 1
dropout: float = 0.0
norm_num_groups: int = 32
cross_attention_dim: Optional[int] = None
attention_bias: bool = False
activation_fn: str = "geglu"
num_embeds_ada_norm: Optional[int] = None
cond_dim_ada_norm_continuous: Optional[int] = None
only_cross_attention: bool = False
double_self_attention: bool = False
upcast_attention: bool = False
norm_type: str = "layer_norm"
norm_elementwise_affine: bool = True
attention_type: str = "default"
enable_memory_efficient_attention: bool = False
gradient_checkpointing: bool = False
cfg: Config
def configure(self) -> None:
super().configure()
self.num_attention_heads = self.cfg.num_attention_heads
self.attention_head_dim = self.cfg.attention_head_dim
inner_dim = self.num_attention_heads * self.attention_head_dim
linear_cls = nn.Linear
if self.cfg.norm_type == "layer_norm" and (
self.cfg.num_embeds_ada_norm is not None
or self.cfg.cond_dim_ada_norm_continuous is not None
):
raise ValueError("Incorrect norm_type.")
# 2. Define input layers
self.in_channels = self.cfg.in_channels
self.norm = torch.nn.GroupNorm(
num_groups=self.cfg.norm_num_groups,
num_channels=self.cfg.in_channels,
eps=1e-6,
affine=True,
)
self.proj_in = linear_cls(self.cfg.in_channels, inner_dim)
# 3. Define transformers blocks
self.transformer_blocks = nn.ModuleList(
[
BasicTransformerBlock(
inner_dim,
self.num_attention_heads,
self.attention_head_dim,
dropout=self.cfg.dropout,
cross_attention_dim=self.cfg.cross_attention_dim,
activation_fn=self.cfg.activation_fn,
num_embeds_ada_norm=self.cfg.num_embeds_ada_norm,
cond_dim_ada_norm_continuous=self.cfg.cond_dim_ada_norm_continuous,
attention_bias=self.cfg.attention_bias,
only_cross_attention=self.cfg.only_cross_attention,
double_self_attention=self.cfg.double_self_attention,
upcast_attention=self.cfg.upcast_attention,
norm_type=self.cfg.norm_type,
norm_elementwise_affine=self.cfg.norm_elementwise_affine,
attention_type=self.cfg.attention_type,
)
for d in range(self.cfg.num_layers)
]
)
# 4. Define output layers
self.out_channels = (
self.cfg.in_channels
if self.cfg.out_channels is None
else self.cfg.out_channels
)
self.proj_out = linear_cls(inner_dim, self.cfg.in_channels)
self.gradient_checkpointing = self.cfg.gradient_checkpointing
self.set_use_memory_efficient_attention_xformers(
self.cfg.enable_memory_efficient_attention
)
def forward(
self,
hidden_states: torch.Tensor,
encoder_hidden_states: Optional[torch.Tensor] = None,
timestep: Optional[torch.LongTensor] = None,
modulation_cond: Optional[torch.FloatTensor] = None,
class_labels: Optional[torch.LongTensor] = None,
cross_attention_kwargs: Dict[str, Any] = None,
attention_mask: Optional[torch.Tensor] = None,
encoder_attention_mask: Optional[torch.Tensor] = None,
):
"""
The [`Transformer1DModel`] forward method.
Args:
hidden_states (`torch.LongTensor` of shape `(batch size, num latent pixels)` if discrete, `torch.FloatTensor` of shape `(batch size, channel, height, width)` if continuous):
Input `hidden_states`.
encoder_hidden_states ( `torch.FloatTensor` of shape `(batch size, sequence len, embed dims)`, *optional*):
Conditional embeddings for cross attention layer. If not given, cross-attention defaults to
self-attention.
timestep ( `torch.LongTensor`, *optional*):
Used to indicate denoising step. Optional timestep to be applied as an embedding in `AdaLayerNorm`.
class_labels ( `torch.LongTensor` of shape `(batch size, num classes)`, *optional*):
Used to indicate class labels conditioning. Optional class labels to be applied as an embedding in
`AdaLayerZeroNorm`.
cross_attention_kwargs ( `Dict[str, Any]`, *optional*):
A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
`self.processor` in
[diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
attention_mask ( `torch.Tensor`, *optional*):
An attention mask of shape `(batch, key_tokens)` is applied to `encoder_hidden_states`. If `1` the mask
is kept, otherwise if `0` it is discarded. Mask will be converted into a bias, which adds large
negative values to the attention scores corresponding to "discard" tokens.
encoder_attention_mask ( `torch.Tensor`, *optional*):
Cross-attention mask applied to `encoder_hidden_states`. Two formats supported:
* Mask `(batch, sequence_length)` True = keep, False = discard.
* Bias `(batch, 1, sequence_length)` 0 = keep, -10000 = discard.
If `ndim == 2`: will be interpreted as a mask, then converted into a bias consistent with the format
above. This bias will be added to the cross-attention scores.
return_dict (`bool`, *optional*, defaults to `True`):
Whether or not to return a [`~models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain
tuple.
Returns:
If `return_dict` is True, an [`~models.transformer_2d.Transformer2DModelOutput`] is returned, otherwise a
`tuple` where the first element is the sample tensor.
"""
# ensure attention_mask is a bias, and give it a singleton query_tokens dimension.
# we may have done this conversion already, e.g. if we came here via UNet2DConditionModel#forward.
# we can tell by counting dims; if ndim == 2: it's a mask rather than a bias.
# expects mask of shape:
# [batch, key_tokens]
# adds singleton query_tokens dimension:
# [batch, 1, key_tokens]
# this helps to broadcast it as a bias over attention scores, which will be in one of the following shapes:
# [batch, heads, query_tokens, key_tokens] (e.g. torch sdp attn)
# [batch * heads, query_tokens, key_tokens] (e.g. xformers or classic attn)
if attention_mask is not None and attention_mask.ndim == 2:
# assume that mask is expressed as:
# (1 = keep, 0 = discard)
# convert mask into a bias that can be added to attention scores:
# (keep = +0, discard = -10000.0)
attention_mask = (1 - attention_mask.to(hidden_states.dtype)) * -10000.0
attention_mask = attention_mask.unsqueeze(1)
# convert encoder_attention_mask to a bias the same way we do for attention_mask
if encoder_attention_mask is not None and encoder_attention_mask.ndim == 2:
encoder_attention_mask = (
1 - encoder_attention_mask.to(hidden_states.dtype)
) * -10000.0
encoder_attention_mask = encoder_attention_mask.unsqueeze(1)
# 1. Input
batch, _, seq_len = hidden_states.shape
residual = hidden_states
hidden_states = self.norm(hidden_states)
inner_dim = hidden_states.shape[1]
hidden_states = hidden_states.permute(0, 2, 1).reshape(
batch, seq_len, inner_dim
)
hidden_states = self.proj_in(hidden_states)
# 2. Blocks
for block in self.transformer_blocks:
if self.training and self.gradient_checkpointing:
hidden_states = torch.utils.checkpoint.checkpoint(
block,
hidden_states,
attention_mask,
encoder_hidden_states,
encoder_attention_mask,
timestep,
modulation_cond,
cross_attention_kwargs,
class_labels,
use_reentrant=False,
)
else:
hidden_states = block(
hidden_states,
attention_mask=attention_mask,
encoder_hidden_states=encoder_hidden_states,
encoder_attention_mask=encoder_attention_mask,
timestep=timestep,
modulation_cond=modulation_cond,
cross_attention_kwargs=cross_attention_kwargs,
class_labels=class_labels,
)
# 3. Output
hidden_states = self.proj_out(hidden_states)
hidden_states = (
hidden_states.reshape(batch, seq_len, inner_dim)
.permute(0, 2, 1)
.contiguous()
)
output = hidden_states + residual
return output
tgs/utils/__init__.py
View File
tgs/utils/base.py
+129
View File
@@ -0,0 +1,129 @@
from dataclasses import dataclass
import torch.nn as nn
from .config import parse_structured
from .misc import get_device, load_module_weights
from .typing import *
class Configurable:
@dataclass
class Config:
pass
def __init__(self, cfg: Optional[dict] = None) -> None:
super().__init__()
self.cfg = parse_structured(self.Config, cfg)
class Updateable:
def do_update_step(
self, epoch: int, global_step: int, on_load_weights: bool = False
):
for attr in self.__dir__():
if attr.startswith("_"):
continue
try:
module = getattr(self, attr)
except:
continue # ignore attributes like property, which can't be retrived using getattr?
if isinstance(module, Updateable):
module.do_update_step(
epoch, global_step, on_load_weights=on_load_weights
)
self.update_step(epoch, global_step, on_load_weights=on_load_weights)
def do_update_step_end(self, epoch: int, global_step: int):
for attr in self.__dir__():
if attr.startswith("_"):
continue
try:
module = getattr(self, attr)
except:
continue # ignore attributes like property, which can't be retrived using getattr?
if isinstance(module, Updateable):
module.do_update_step_end(epoch, global_step)
self.update_step_end(epoch, global_step)
def update_step(self, epoch: int, global_step: int, on_load_weights: bool = False):
# override this method to implement custom update logic
# if on_load_weights is True, you should be careful doing things related to model evaluations,
# as the models and tensors are not guarenteed to be on the same device
pass
def update_step_end(self, epoch: int, global_step: int):
pass
def update_if_possible(module: Any, epoch: int, global_step: int) -> None:
if isinstance(module, Updateable):
module.do_update_step(epoch, global_step)
def update_end_if_possible(module: Any, epoch: int, global_step: int) -> None:
if isinstance(module, Updateable):
module.do_update_step_end(epoch, global_step)
class BaseObject(Updateable):
@dataclass
class Config:
pass
cfg: Config # add this to every subclass of BaseObject to enable static type checking
def __init__(
self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs
) -> None:
super().__init__()
self.cfg = parse_structured(self.Config, cfg)
self.device = get_device()
self.configure(*args, **kwargs)
def configure(self, *args, **kwargs) -> None:
pass
class BaseModule(nn.Module, Updateable):
@dataclass
class Config:
weights: Optional[str] = None
freeze: Optional[bool] = False
cfg: Config # add this to every subclass of BaseModule to enable static type checking
def __init__(
self, cfg: Optional[Union[dict, DictConfig]] = None, *args, **kwargs
) -> None:
super().__init__()
self.cfg = parse_structured(self.Config, cfg)
self.device = get_device()
self._non_modules = {}
self.configure(*args, **kwargs)
if self.cfg.weights is not None:
# format: path/to/weights:module_name
weights_path, module_name = self.cfg.weights.split(":")
state_dict = load_module_weights(
weights_path, module_name=module_name, map_location="cpu"
)
self.load_state_dict(state_dict, strict=False)
# self.do_update_step(
# epoch, global_step, on_load_weights=True
# ) # restore states
if self.cfg.freeze:
for params in self.parameters():
params.requires_grad = False
def configure(self, *args, **kwargs) -> None:
pass
def register_non_module(self, name: str, module: nn.Module) -> None:
# non-modules won't be treated as model parameters
if name in self._non_modules:
raise ValueError(f"Non-module {name} already exists!")
self._non_modules[name] = module
def non_module(self, name: str):
return self._non_modules.get(name, None)
tgs/utils/config.py
+74
View File
@@ -0,0 +1,74 @@
import os
from dataclasses import dataclass, field
from omegaconf import OmegaConf
from .typing import *
# ============ Register OmegaConf Recolvers ============= #
OmegaConf.register_new_resolver(
"calc_exp_lr_decay_rate", lambda factor, n: factor ** (1.0 / n), replace=True
)
OmegaConf.register_new_resolver("add", lambda a, b: a + b, replace=True)
OmegaConf.register_new_resolver("sub", lambda a, b: a - b, replace=True)
OmegaConf.register_new_resolver("mul", lambda a, b: a * b, replace=True)
OmegaConf.register_new_resolver("div", lambda a, b: a / b, replace=True)
OmegaConf.register_new_resolver("idiv", lambda a, b: a // b, replace=True)
OmegaConf.register_new_resolver("basename", lambda p: os.path.basename(p), replace=True)
OmegaConf.register_new_resolver("rmspace", lambda s, sub: s.replace(" ", sub), replace=True)
OmegaConf.register_new_resolver("tuple2", lambda s: [float(s), float(s)], replace=True)
OmegaConf.register_new_resolver("gt0", lambda s: s > 0, replace=True)
OmegaConf.register_new_resolver("not", lambda s: not s, replace=True)
OmegaConf.register_new_resolver("shsdim", lambda sh_degree: (sh_degree + 1) ** 2 * 3, replace=True)
# ======================================================= #
# ============== Automatic Name Resolvers =============== #
def get_naming_convention(cfg):
# TODO
name = f"tgs_{cfg.system.backbone.num_layers}"
return name
# ======================================================= #
@dataclass
class ExperimentConfig:
n_gpus: int = 1
data: dict = field(default_factory=dict)
system: dict = field(default_factory=dict)
def load_config(
*yamls: str, cli_args: list = [], from_string=False, makedirs=True, **kwargs
) -> Any:
if from_string:
parse_func = OmegaConf.create
else:
parse_func = OmegaConf.load
yaml_confs = []
for y in yamls:
conf = parse_func(y)
extends = conf.pop("extends", None)
if extends:
assert os.path.exists(extends), f"File {extends} does not exist."
yaml_confs.append(OmegaConf.load(extends))
yaml_confs.append(conf)
cli_conf = OmegaConf.from_cli(cli_args)
cfg = OmegaConf.merge(*yaml_confs, cli_conf, kwargs)
OmegaConf.resolve(cfg)
assert isinstance(cfg, DictConfig)
scfg: ExperimentConfig = parse_structured(ExperimentConfig, cfg)
return scfg
def config_to_primitive(config, resolve: bool = True) -> Any:
return OmegaConf.to_container(config, resolve=resolve)
def dump_config(path: str, config) -> None:
with open(path, "w") as fp:
OmegaConf.save(config=config, f=fp)
def parse_structured(fields: Any, cfg: Optional[Union[dict, DictConfig]] = None) -> Any:
scfg = OmegaConf.structured(fields(**cfg))
return scfg
tgs/utils/misc.py
+88
View File
@@ -0,0 +1,88 @@
import os
import re
import torch
from packaging import version
from .typing import *
def parse_version(ver: str):
return version.parse(ver)
def get_rank():
# SLURM_PROCID can be set even if SLURM is not managing the multiprocessing,
# therefore LOCAL_RANK needs to be checked first
rank_keys = ("RANK", "LOCAL_RANK", "SLURM_PROCID", "JSM_NAMESPACE_RANK")
for key in rank_keys:
rank = os.environ.get(key)
if rank is not None:
return int(rank)
return 0
def get_device():
return torch.device(f"cuda:{get_rank()}")
def load_module_weights(
path, module_name=None, ignore_modules=None, map_location=None
) -> Tuple[dict, int, int]:
if module_name is not None and ignore_modules is not None:
raise ValueError("module_name and ignore_modules cannot be both set")
if map_location is None:
map_location = get_device()
ckpt = torch.load(path, map_location=map_location)
state_dict = ckpt["state_dict"]
state_dict_to_load = state_dict
if ignore_modules is not None:
state_dict_to_load = {}
for k, v in state_dict.items():
ignore = any(
[k.startswith(ignore_module + ".") for ignore_module in ignore_modules]
)
if ignore:
continue
state_dict_to_load[k] = v
if module_name is not None:
state_dict_to_load = {}
for k, v in state_dict.items():
m = re.match(rf"^{module_name}\.(.*)$", k)
if m is None:
continue
state_dict_to_load[m.group(1)] = v
return state_dict_to_load
# convert a function into recursive style to handle nested dict/list/tuple variables
def make_recursive_func(func):
def wrapper(vars, *args, **kwargs):
if isinstance(vars, list):
return [wrapper(x, *args, **kwargs) for x in vars]
elif isinstance(vars, tuple):
return tuple([wrapper(x, *args, **kwargs) for x in vars])
elif isinstance(vars, dict):
return {k: wrapper(v, *args, **kwargs) for k, v in vars.items()}
else:
return func(vars, *args, **kwargs)
return wrapper
@make_recursive_func
def todevice(vars, device="cuda"):
if isinstance(vars, torch.Tensor):
return vars.to(device)
elif isinstance(vars, str):
return vars
elif isinstance(vars, bool):
return vars
elif isinstance(vars, float):
return vars
elif isinstance(vars, int):
return vars
else:
raise NotImplementedError("invalid input type {} for tensor2numpy".format(type(vars)))
tgs/utils/ops.py
+279
View File
@@ -0,0 +1,279 @@
import math
import numpy as np
import torch
import torch.nn.functional as F
from torch.autograd import Function
from torch.cuda.amp import custom_bwd, custom_fwd
from pytorch3d import io
from pytorch3d.renderer import (
PointsRasterizationSettings,
PointsRasterizer)
from pytorch3d.structures import Pointclouds
from pytorch3d.utils.camera_conversions import cameras_from_opencv_projection
import cv2
from .typing import *
ValidScale = Union[Tuple[float, float], Num[Tensor, "2 D"]]
def scale_tensor(
dat: Num[Tensor, "... D"], inp_scale: ValidScale, tgt_scale: ValidScale
):
if inp_scale is None:
inp_scale = (0, 1)
if tgt_scale is None:
tgt_scale = (0, 1)
if isinstance(tgt_scale, Tensor):
assert dat.shape[-1] == tgt_scale.shape[-1]
dat = (dat - inp_scale[0]) / (inp_scale[1] - inp_scale[0])
dat = dat * (tgt_scale[1] - tgt_scale[0]) + tgt_scale[0]
return dat
class _TruncExp(Function): # pylint: disable=abstract-method
# Implementation from torch-ngp:
# https://github.com/ashawkey/torch-ngp/blob/93b08a0d4ec1cc6e69d85df7f0acdfb99603b628/activation.py
@staticmethod
@custom_fwd(cast_inputs=torch.float32)
def forward(ctx, x): # pylint: disable=arguments-differ
ctx.save_for_backward(x)
return torch.exp(x)
@staticmethod
@custom_bwd
def backward(ctx, g): # pylint: disable=arguments-differ
x = ctx.saved_tensors[0]
return g * torch.exp(torch.clamp(x, max=15))
trunc_exp = _TruncExp.apply
def get_activation(name) -> Callable:
if name is None:
return lambda x: x
name = name.lower()
if name == "none":
return lambda x: x
elif name == "lin2srgb":
return lambda x: torch.where(
x > 0.0031308,
torch.pow(torch.clamp(x, min=0.0031308), 1.0 / 2.4) * 1.055 - 0.055,
12.92 * x,
).clamp(0.0, 1.0)
elif name == "exp":
return lambda x: torch.exp(x)
elif name == "shifted_exp":
return lambda x: torch.exp(x - 1.0)
elif name == "trunc_exp":
return trunc_exp
elif name == "shifted_trunc_exp":
return lambda x: trunc_exp(x - 1.0)
elif name == "sigmoid":
return lambda x: torch.sigmoid(x)
elif name == "tanh":
return lambda x: torch.tanh(x)
elif name == "shifted_softplus":
return lambda x: F.softplus(x - 1.0)
elif name == "scale_-11_01":
return lambda x: x * 0.5 + 0.5
else:
try:
return getattr(F, name)
except AttributeError:
raise ValueError(f"Unknown activation function: {name}")
def get_ray_directions(
H: int,
W: int,
focal: Union[float, Tuple[float, float]],
principal: Optional[Tuple[float, float]] = None,
use_pixel_centers: bool = True,
) -> Float[Tensor, "H W 3"]:
"""
Get ray directions for all pixels in camera coordinate.
Reference: https://www.scratchapixel.com/lessons/3d-basic-rendering/
ray-tracing-generating-camera-rays/standard-coordinate-systems
Inputs:
H, W, focal, principal, use_pixel_centers: image height, width, focal length, principal point and whether use pixel centers
Outputs:
directions: (H, W, 3), the direction of the rays in camera coordinate
"""
pixel_center = 0.5 if use_pixel_centers else 0
if isinstance(focal, float):
fx, fy = focal, focal
cx, cy = W / 2, H / 2
else:
fx, fy = focal
assert principal is not None
cx, cy = principal
i, j = torch.meshgrid(
torch.arange(W, dtype=torch.float32) + pixel_center,
torch.arange(H, dtype=torch.float32) + pixel_center,
indexing="xy",
)
directions: Float[Tensor, "H W 3"] = torch.stack(
[(i - cx) / fx, -(j - cy) / fy, -torch.ones_like(i)], -1
)
return directions
def get_rays(
directions: Float[Tensor, "... 3"],
c2w: Float[Tensor, "... 4 4"],
keepdim=False,
noise_scale=0.0,
) -> Tuple[Float[Tensor, "... 3"], Float[Tensor, "... 3"]]:
# Rotate ray directions from camera coordinate to the world coordinate
assert directions.shape[-1] == 3
if directions.ndim == 2: # (N_rays, 3)
if c2w.ndim == 2: # (4, 4)
c2w = c2w[None, :, :]
assert c2w.ndim == 3 # (N_rays, 4, 4) or (1, 4, 4)
rays_d = (directions[:, None, :] * c2w[:, :3, :3]).sum(-1) # (N_rays, 3)
rays_o = c2w[:, :3, 3].expand(rays_d.shape)
elif directions.ndim == 3: # (H, W, 3)
assert c2w.ndim in [2, 3]
if c2w.ndim == 2: # (4, 4)
rays_d = (directions[:, :, None, :] * c2w[None, None, :3, :3]).sum(
-1
) # (H, W, 3)
rays_o = c2w[None, None, :3, 3].expand(rays_d.shape)
elif c2w.ndim == 3: # (B, 4, 4)
rays_d = (directions[None, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
-1
) # (B, H, W, 3)
rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
elif directions.ndim == 4: # (B, H, W, 3)
assert c2w.ndim == 3 # (B, 4, 4)
rays_d = (directions[:, :, :, None, :] * c2w[:, None, None, :3, :3]).sum(
-1
) # (B, H, W, 3)
rays_o = c2w[:, None, None, :3, 3].expand(rays_d.shape)
# add camera noise to avoid grid-like artifect
# https://github.com/ashawkey/stable-dreamfusion/blob/49c3d4fa01d68a4f027755acf94e1ff6020458cc/nerf/utils.py#L373
if noise_scale > 0:
rays_o = rays_o + torch.randn(3, device=rays_o.device) * noise_scale
rays_d = rays_d + torch.randn(3, device=rays_d.device) * noise_scale
rays_d = F.normalize(rays_d, dim=-1)
if not keepdim:
rays_o, rays_d = rays_o.reshape(-1, 3), rays_d.reshape(-1, 3)
return rays_o, rays_d
def get_projection_matrix(
fovy: Union[float, Float[Tensor, "B"]], aspect_wh: float, near: float, far: float
) -> Float[Tensor, "*B 4 4"]:
if isinstance(fovy, float):
proj_mtx = torch.zeros(4, 4, dtype=torch.float32)
proj_mtx[0, 0] = 1.0 / (math.tan(fovy / 2.0) * aspect_wh)
proj_mtx[1, 1] = -1.0 / math.tan(
fovy / 2.0
) # add a negative sign here as the y axis is flipped in nvdiffrast output
proj_mtx[2, 2] = -(far + near) / (far - near)
proj_mtx[2, 3] = -2.0 * far * near / (far - near)
proj_mtx[3, 2] = -1.0
else:
batch_size = fovy.shape[0]
proj_mtx = torch.zeros(batch_size, 4, 4, dtype=torch.float32)
proj_mtx[:, 0, 0] = 1.0 / (torch.tan(fovy / 2.0) * aspect_wh)
proj_mtx[:, 1, 1] = -1.0 / torch.tan(
fovy / 2.0
) # add a negative sign here as the y axis is flipped in nvdiffrast output
proj_mtx[:, 2, 2] = -(far + near) / (far - near)
proj_mtx[:, 2, 3] = -2.0 * far * near / (far - near)
proj_mtx[:, 3, 2] = -1.0
return proj_mtx
def get_mvp_matrix(
c2w: Float[Tensor, "*B 4 4"], proj_mtx: Float[Tensor, "*B 4 4"]
) -> Float[Tensor, "*B 4 4"]:
# calculate w2c from c2w: R' = Rt, t' = -Rt * t
# mathematically equivalent to (c2w)^-1
if c2w.ndim == 2:
assert proj_mtx.ndim == 2
w2c: Float[Tensor, "4 4"] = torch.zeros(4, 4).to(c2w)
w2c[:3, :3] = c2w[:3, :3].permute(1, 0)
w2c[:3, 3:] = -c2w[:3, :3].permute(1, 0) @ c2w[:3, 3:]
w2c[3, 3] = 1.0
else:
w2c: Float[Tensor, "B 4 4"] = torch.zeros(c2w.shape[0], 4, 4).to(c2w)
w2c[:, :3, :3] = c2w[:, :3, :3].permute(0, 2, 1)
w2c[:, :3, 3:] = -c2w[:, :3, :3].permute(0, 2, 1) @ c2w[:, :3, 3:]
w2c[:, 3, 3] = 1.0
# calculate mvp matrix by proj_mtx @ w2c (mv_mtx)
mvp_mtx = proj_mtx @ w2c
return mvp_mtx
def get_intrinsic_from_fov(fov, H, W, bs=-1):
focal_length = 0.5 * H / np.tan(0.5 * fov)
intrinsic = np.identity(3, dtype=np.float32)
intrinsic[0, 0] = focal_length
intrinsic[1, 1] = focal_length
intrinsic[0, 2] = W / 2.0
intrinsic[1, 2] = H / 2.0
if bs > 0:
intrinsic = intrinsic[None].repeat(bs, axis=0)
return torch.from_numpy(intrinsic)
def points_projection(points: Float[Tensor, "B Np 3"],
c2ws: Float[Tensor, "B 4 4"],
intrinsics: Float[Tensor, "B 3 3"],
local_features: Float[Tensor, "B C H W"],
# Rasterization settings
raster_point_radius: float = 0.0075, # point size
raster_points_per_pixel: int = 1, # a single point per pixel, for now
bin_size: int = 0):
B, C, H, W = local_features.shape
device = local_features.device
raster_settings = PointsRasterizationSettings(
image_size=(H, W),
radius=raster_point_radius,
points_per_pixel=raster_points_per_pixel,
bin_size=bin_size,
)
Np = points.shape[1]
R = raster_settings.points_per_pixel
w2cs = torch.inverse(c2ws)
image_size = torch.as_tensor([H, W]).view(1, 2).expand(w2cs.shape[0], -1).to(device)
cameras = cameras_from_opencv_projection(w2cs[:, :3, :3], w2cs[:, :3, 3], intrinsics, image_size)
rasterize = PointsRasterizer(cameras=cameras, raster_settings=raster_settings)
fragments = rasterize(Pointclouds(points))
fragments_idx: Tensor = fragments.idx.long()
visible_pixels = (fragments_idx > -1) # (B, H, W, R)
points_to_visible_pixels = fragments_idx[visible_pixels]
# Reshape local features to (B, H, W, R, C)
local_features = local_features.permute(0, 2, 3, 1).unsqueeze(-2).expand(-1, -1, -1, R, -1) # (B, H, W, R, C)
# Get local features corresponding to visible points
local_features_proj = torch.zeros(B * Np, C, device=device)
local_features_proj[points_to_visible_pixels] = local_features[visible_pixels]
local_features_proj = local_features_proj.reshape(B, Np, C)
return local_features_proj
def compute_distance_transform(mask: torch.Tensor):
image_size = mask.shape[-1]
distance_transform = torch.stack([
torch.from_numpy(cv2.distanceTransform(
(1 - m), distanceType=cv2.DIST_L2, maskSize=cv2.DIST_MASK_3
) / (image_size / 2))
for m in mask.squeeze(1).detach().cpu().numpy().astype(np.uint8)
]).unsqueeze(1).clip(0, 1).to(mask.device)
return distance_transform
tgs/utils/saving.py
+314
View File
@@ -0,0 +1,314 @@
import os
import re
import shutil
import cv2
import imageio
import matplotlib.pyplot as plt
import numpy as np
import torch
from matplotlib import cm
from matplotlib.colors import LinearSegmentedColormap
from PIL import Image, ImageDraw
from .typing import *
class SaverMixin:
_save_dir: Optional[str] = None
def set_save_dir(self, save_dir: str):
self._save_dir = save_dir
def get_save_dir(self):
if self._save_dir is None:
raise ValueError("Save dir is not set")
return self._save_dir
def convert_data(self, data):
if data is None:
return None
elif isinstance(data, np.ndarray):
return data
elif isinstance(data, torch.Tensor):
return data.detach().cpu().numpy()
elif isinstance(data, list):
return [self.convert_data(d) for d in data]
elif isinstance(data, dict):
return {k: self.convert_data(v) for k, v in data.items()}
else:
raise TypeError(
"Data must be in type numpy.ndarray, torch.Tensor, list or dict, getting",
type(data),
)
def get_save_path(self, filename):
save_path = os.path.join(self.get_save_dir(), filename)
os.makedirs(os.path.dirname(save_path), exist_ok=True)
return save_path
DEFAULT_RGB_KWARGS = {"data_format": "HWC", "data_range": (0, 1)}
DEFAULT_UV_KWARGS = {
"data_format": "HWC",
"data_range": (0, 1),
"cmap": "checkerboard",
}
DEFAULT_GRAYSCALE_KWARGS = {"data_range": None, "cmap": "jet"}
DEFAULT_GRID_KWARGS = {"align": "max"}
def get_rgb_image_(self, img, data_format, data_range, rgba=False):
img = self.convert_data(img)
assert data_format in ["CHW", "HWC"]
if data_format == "CHW":
img = img.transpose(1, 2, 0)
if img.dtype != np.uint8:
img = img.clip(min=data_range[0], max=data_range[1])
img = (
(img - data_range[0]) / (data_range[1] - data_range[0]) * 255.0
).astype(np.uint8)
nc = 4 if rgba else 3
imgs = [img[..., start : start + nc] for start in range(0, img.shape[-1], nc)]
imgs = [
img_
if img_.shape[-1] == nc
else np.concatenate(
[
img_,
np.zeros(
(img_.shape[0], img_.shape[1], nc - img_.shape[2]),
dtype=img_.dtype,
),
],
axis=-1,
)
for img_ in imgs
]
img = np.concatenate(imgs, axis=1)
if rgba:
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
else:
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
return img
def _save_rgb_image(
self,
filename,
img,
data_format,
data_range
):
img = self.get_rgb_image_(img, data_format, data_range)
cv2.imwrite(filename, img)
def save_rgb_image(
self,
filename,
img,
data_format=DEFAULT_RGB_KWARGS["data_format"],
data_range=DEFAULT_RGB_KWARGS["data_range"],
) -> str:
save_path = self.get_save_path(filename)
self._save_rgb_image(save_path, img, data_format, data_range)
return save_path
def get_grayscale_image_(self, img, data_range, cmap):
img = self.convert_data(img)
img = np.nan_to_num(img)
if data_range is None:
img = (img - img.min()) / (img.max() - img.min())
else:
img = img.clip(data_range[0], data_range[1])
img = (img - data_range[0]) / (data_range[1] - data_range[0])
assert cmap in [None, "jet", "magma", "spectral"]
if cmap == None:
img = (img * 255.0).astype(np.uint8)
img = np.repeat(img[..., None], 3, axis=2)
elif cmap == "jet":
img = (img * 255.0).astype(np.uint8)
img = cv2.applyColorMap(img, cv2.COLORMAP_JET)
elif cmap == "magma":
img = 1.0 - img
base = cm.get_cmap("magma")
num_bins = 256
colormap = LinearSegmentedColormap.from_list(
f"{base.name}{num_bins}", base(np.linspace(0, 1, num_bins)), num_bins
)(np.linspace(0, 1, num_bins))[:, :3]
a = np.floor(img * 255.0)
b = (a + 1).clip(max=255.0)
f = img * 255.0 - a
a = a.astype(np.uint16).clip(0, 255)
b = b.astype(np.uint16).clip(0, 255)
img = colormap[a] + (colormap[b] - colormap[a]) * f[..., None]
img = (img * 255.0).astype(np.uint8)
elif cmap == "spectral":
colormap = plt.get_cmap("Spectral")
def blend_rgba(image):
image = image[..., :3] * image[..., -1:] + (
1.0 - image[..., -1:]
) # blend A to RGB
return image
img = colormap(img)
img = blend_rgba(img)
img = (img * 255).astype(np.uint8)
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
return img
def _save_grayscale_image(
self,
filename,
img,
data_range,
cmap,
):
img = self.get_grayscale_image_(img, data_range, cmap)
cv2.imwrite(filename, img)
def save_grayscale_image(
self,
filename,
img,
data_range=DEFAULT_GRAYSCALE_KWARGS["data_range"],
cmap=DEFAULT_GRAYSCALE_KWARGS["cmap"],
) -> str:
save_path = self.get_save_path(filename)
self._save_grayscale_image(save_path, img, data_range, cmap)
return save_path
def get_image_grid_(self, imgs, align):
if isinstance(imgs[0], list):
return np.concatenate(
[self.get_image_grid_(row, align) for row in imgs], axis=0
)
cols = []
for col in imgs:
assert col["type"] in ["rgb", "uv", "grayscale"]
if col["type"] == "rgb":
rgb_kwargs = self.DEFAULT_RGB_KWARGS.copy()
rgb_kwargs.update(col["kwargs"])
cols.append(self.get_rgb_image_(col["img"], **rgb_kwargs))
elif col["type"] == "uv":
uv_kwargs = self.DEFAULT_UV_KWARGS.copy()
uv_kwargs.update(col["kwargs"])
cols.append(self.get_uv_image_(col["img"], **uv_kwargs))
elif col["type"] == "grayscale":
grayscale_kwargs = self.DEFAULT_GRAYSCALE_KWARGS.copy()
grayscale_kwargs.update(col["kwargs"])
cols.append(self.get_grayscale_image_(col["img"], **grayscale_kwargs))
if align == "max":
h = max([col.shape[0] for col in cols])
w = max([col.shape[1] for col in cols])
elif align == "min":
h = min([col.shape[0] for col in cols])
w = min([col.shape[1] for col in cols])
elif isinstance(align, int):
h = align
w = align
elif (
isinstance(align, tuple)
and isinstance(align[0], int)
and isinstance(align[1], int)
):
h, w = align
else:
raise ValueError(
f"Unsupported image grid align: {align}, should be min, max, int or (int, int)"
)
for i in range(len(cols)):
if cols[i].shape[0] != h or cols[i].shape[1] != w:
cols[i] = cv2.resize(cols[i], (w, h), interpolation=cv2.INTER_LINEAR)
return np.concatenate(cols, axis=1)
def save_image_grid(
self,
filename,
imgs,
align=DEFAULT_GRID_KWARGS["align"],
texts: Optional[List[float]] = None,
):
save_path = self.get_save_path(filename)
img = self.get_image_grid_(imgs, align=align)
if texts is not None:
img = Image.fromarray(img)
draw = ImageDraw.Draw(img)
black, white = (0, 0, 0), (255, 255, 255)
for i, text in enumerate(texts):
draw.text((2, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
draw.text((0, (img.size[1] // len(texts)) * i + 1), f"{text}", white)
draw.text((2, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
draw.text((0, (img.size[1] // len(texts)) * i - 1), f"{text}", white)
draw.text((1, (img.size[1] // len(texts)) * i), f"{text}", black)
img = np.asarray(img)
cv2.imwrite(save_path, img)
return save_path
def save_image(self, filename, img) -> str:
save_path = self.get_save_path(filename)
img = self.convert_data(img)
assert img.dtype == np.uint8 or img.dtype == np.uint16
if img.ndim == 3 and img.shape[-1] == 3:
img = cv2.cvtColor(img, cv2.COLOR_RGB2BGR)
elif img.ndim == 3 and img.shape[-1] == 4:
img = cv2.cvtColor(img, cv2.COLOR_RGBA2BGRA)
cv2.imwrite(save_path, img)
return save_path
def save_img_sequence(
self,
filename,
img_dir,
matcher,
save_format="mp4",
fps=30,
) -> str:
assert save_format in ["gif", "mp4"]
if not filename.endswith(save_format):
filename += f".{save_format}"
save_path = self.get_save_path(filename)
matcher = re.compile(matcher)
img_dir = os.path.join(self.get_save_dir(), img_dir)
imgs = []
for f in os.listdir(img_dir):
if matcher.search(f):
imgs.append(f)
imgs = sorted(imgs, key=lambda f: int(matcher.search(f).groups()[0]))
imgs = [cv2.imread(os.path.join(img_dir, f)) for f in imgs]
if save_format == "gif":
imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
imageio.mimsave(save_path, imgs, fps=fps, palettesize=256)
elif save_format == "mp4":
imgs = [cv2.cvtColor(i, cv2.COLOR_BGR2RGB) for i in imgs]
imageio.mimsave(save_path, imgs, fps=fps)
return save_path
def save_img_sequences(
self,
seq_dir,
matcher,
save_format="mp4",
fps=30,
delete=True
):
seq_dir_ = os.path.join(self.get_save_dir(), seq_dir)
for f in os.listdir(seq_dir_):
img_dir_ = os.path.join(seq_dir_, f)
if not os.path.isdir(img_dir_):
continue
try:
self.save_img_sequence(
os.path.join(seq_dir, f),
os.path.join(seq_dir, f),
matcher,
save_format=save_format,
fps=fps
)
except:
raise ValueError(f"Video saving for directory {seq_dir_} failed!")
if delete:
shutil.rmtree(img_dir_)
tgs/utils/typing.py
+40
View File
@@ -0,0 +1,40 @@
"""
This module contains type annotations for the project, using
1. Python type hints (https://docs.python.org/3/library/typing.html) for Python objects
2. jaxtyping (https://github.com/google/jaxtyping/blob/main/API.md) for PyTorch tensors
Two types of typing checking can be used:
1. Static type checking with mypy (install with pip and enabled as the default linter in VSCode)
2. Runtime type checking with typeguard (install with pip and triggered at runtime, mainly for tensor dtype and shape checking)
"""
# Basic types
from typing import (
Any,
Callable,
Dict,
Iterable,
List,
Literal,
NamedTuple,
NewType,
Optional,
Sized,
Tuple,
Type,
TypeVar,
Union,
)
# Tensor dtype
# for jaxtyping usage, see https://github.com/google/jaxtyping/blob/main/API.md
from jaxtyping import Bool, Complex, Float, Inexact, Int, Integer, Num, Shaped, UInt
# Config type
from omegaconf import DictConfig
# PyTorch Tensor type
from torch import Tensor
# Runtime type checking decorator
from typeguard import typechecked as typechecker