diff --git a/.gitignore b/.gitignore index 7c0ae4b..0c507ed 100644 --- a/.gitignore +++ b/.gitignore @@ -155,4 +155,13 @@ cython_debug/ diff-gaussian-rasterization/ # Nvdiffrast -nvdiffrast/ \ No newline at end of file +nvdiffrast/ + +# Model files +*.ckpt +*.bin +*.pth +*.safetensors + +# Ignore test outputs +outputs/ \ No newline at end of file diff --git a/README.md b/README.md index 0cce225..3c422f2 100644 --- a/README.md +++ b/README.md @@ -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: + + + +- 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
+ *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: - - - -- 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
- *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: ``` diff --git a/_Example_Workflows/3DGS_Multi-View_Renders_to_ply.json b/_Example_Workflows/3DGS_Multi-View_Renders_to_ply.json index 400b719..9bc5965 100644 --- a/_Example_Workflows/3DGS_Multi-View_Renders_to_ply.json +++ b/_Example_Workflows/3DGS_Multi-View_Renders_to_ply.json @@ -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": [], diff --git a/_Example_Workflows/3DGS_Multi-View_Renders_to_ply_with_Stack_Orbit_Cam_Poses_node.json b/_Example_Workflows/3DGS_Multi-View_Renders_to_ply_with_Stack_Orbit_Cam_Poses_node.json index 96ae54d..d6470ff 100644 --- a/_Example_Workflows/3DGS_Multi-View_Renders_to_ply_with_Stack_Orbit_Cam_Poses_node.json +++ b/_Example_Workflows/3DGS_Multi-View_Renders_to_ply_with_Stack_Orbit_Cam_Poses_node.json @@ -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": [], diff --git a/_Example_Workflows/Images_to_3DGS_to_DMTetMesh.json b/_Example_Workflows/Images_to_3DGS_to_DMTetMesh.json index 09df809..2500522 100644 --- a/_Example_Workflows/Images_to_3DGS_to_DMTetMesh.json +++ b/_Example_Workflows/Images_to_3DGS_to_DMTetMesh.json @@ -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": [], diff --git a/_Example_Workflows/Load_3DGS_to_DMTetMesh.json b/_Example_Workflows/Load_3DGS_to_DMTetMesh.json index 06e2fab..9672e25 100644 --- a/_Example_Workflows/Load_3DGS_to_DMTetMesh.json +++ b/_Example_Workflows/Load_3DGS_to_DMTetMesh.json @@ -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": [], diff --git a/_Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS.json b/_Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS.json new file mode 100644 index 0000000..57ea542 --- /dev/null +++ b/_Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS.json @@ -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 +} \ No newline at end of file diff --git a/_Example_Workflows/_Example_Inputs_Files/a_pikachu_with_smily_face-removebg.png b/_Example_Workflows/_Example_Inputs_Files/a_pikachu_with_smily_face-removebg.png new file mode 100644 index 0000000..34ca7c4 Binary files /dev/null and b/_Example_Workflows/_Example_Inputs_Files/a_pikachu_with_smily_face-removebg.png differ diff --git a/algorithms/diff_texturing.py b/algorithms/diff_texturing.py index dab262a..5176a8f 100644 --- a/algorithms/diff_texturing.py +++ b/algorithms/diff_texturing.py @@ -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 diff --git a/algorithms/main_3DGS_renderer.py b/algorithms/main_3DGS_renderer.py index 65999c0..70270d1 100644 --- a/algorithms/main_3DGS_renderer.py +++ b/algorithms/main_3DGS_renderer.py @@ -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) diff --git a/algorithms/triplane_gaussian_transformers.py b/algorithms/triplane_gaussian_transformers.py new file mode 100644 index 0000000..61c6fde --- /dev/null +++ b/algorithms/triplane_gaussian_transformers.py @@ -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"}, + } + ] + ) \ No newline at end of file diff --git a/configs/tgs_config.yaml b/configs/tgs_config.yaml new file mode 100644 index 0000000..50eab4e --- /dev/null +++ b/configs/tgs_config.yaml @@ -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 \ No newline at end of file diff --git a/install.bat b/install.bat index df5a491..96ddaab 100644 --- a/install.bat +++ b/install.bat @@ -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%" ) diff --git a/mesh_processer/mesh.py b/mesh_processer/mesh.py index 43d6935..7fcc263 100644 --- a/mesh_processer/mesh.py +++ b/mesh_processer/mesh.py @@ -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 diff --git a/mesh_processer/mesh_utils.py b/mesh_processer/mesh_utils.py index f454dad..698b130 100644 --- a/mesh_processer/mesh_utils.py +++ b/mesh_processer/mesh_utils.py @@ -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 \ No newline at end of file diff --git a/nodes.py b/nodes.py index 86d0822..c165cd1 100644 --- a/nodes.py +++ b/nodes.py @@ -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, ) \ No newline at end of file + 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) + \ No newline at end of file diff --git a/requirements.txt b/requirements.txt index ce532f5..5cf47fc 100644 --- a/requirements.txt +++ b/requirements.txt @@ -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 diff --git a/test.bat b/test.bat index f8ae281..b12512a 100644 --- a/test.bat +++ b/test.bat @@ -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 \ No newline at end of file +cd tgs/models/snowflake/pointnet2_ops_lib && "%python_miniconda_exec%" -s setup.py install && cd ../../../../ \ No newline at end of file diff --git a/tgs/__init__.py b/tgs/__init__.py new file mode 100644 index 0000000..889caab --- /dev/null +++ b/tgs/__init__.py @@ -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 \ No newline at end of file diff --git a/tgs/data.py b/tgs/data.py new file mode 100644 index 0000000..153dd41 --- /dev/null +++ b/tgs/data.py @@ -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 \ No newline at end of file diff --git a/tgs/models/__init__.py b/tgs/models/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tgs/models/image_feature.py b/tgs/models/image_feature.py new file mode 100644 index 0000000..88d3424 --- /dev/null +++ b/tgs/models/image_feature.py @@ -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) \ No newline at end of file diff --git a/tgs/models/networks.py b/tgs/models/networks.py new file mode 100644 index 0000000..082671e --- /dev/null +++ b/tgs/models/networks.py @@ -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 \ No newline at end of file diff --git a/tgs/models/pointclouds/LICENSE_POINTNET b/tgs/models/pointclouds/LICENSE_POINTNET new file mode 100644 index 0000000..5358ea9 --- /dev/null +++ b/tgs/models/pointclouds/LICENSE_POINTNET @@ -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. \ No newline at end of file diff --git a/tgs/models/pointclouds/pointnet.py b/tgs/models/pointclouds/pointnet.py new file mode 100644 index 0000000..14b5cbe --- /dev/null +++ b/tgs/models/pointclouds/pointnet.py @@ -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 diff --git a/tgs/models/pointclouds/simplepoint.py b/tgs/models/pointclouds/simplepoint.py new file mode 100644 index 0000000..2fb028d --- /dev/null +++ b/tgs/models/pointclouds/simplepoint.py @@ -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 \ No newline at end of file diff --git a/tgs/models/renderer.py b/tgs/models/renderer.py new file mode 100644 index 0000000..b897303 --- /dev/null +++ b/tgs/models/renderer.py @@ -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 + \ No newline at end of file diff --git a/tgs/models/snowflake/LICENSE b/tgs/models/snowflake/LICENSE new file mode 100644 index 0000000..01c2b54 --- /dev/null +++ b/tgs/models/snowflake/LICENSE @@ -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. diff --git a/tgs/models/snowflake/SPD.py b/tgs/models/snowflake/SPD.py new file mode 100644 index 0000000..eae592c --- /dev/null +++ b/tgs/models/snowflake/SPD.py @@ -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 \ No newline at end of file diff --git a/tgs/models/snowflake/SPD_crossattn.py b/tgs/models/snowflake/SPD_crossattn.py new file mode 100644 index 0000000..90273e9 --- /dev/null +++ b/tgs/models/snowflake/SPD_crossattn.py @@ -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 \ No newline at end of file diff --git a/tgs/models/snowflake/SPD_pp.py b/tgs/models/snowflake/SPD_pp.py new file mode 100644 index 0000000..886444e --- /dev/null +++ b/tgs/models/snowflake/SPD_pp.py @@ -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 diff --git a/tgs/models/snowflake/attention.py b/tgs/models/snowflake/attention.py new file mode 100644 index 0000000..51d4d80 --- /dev/null +++ b/tgs/models/snowflake/attention.py @@ -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 \ No newline at end of file diff --git a/tgs/models/snowflake/model_spdpp.py b/tgs/models/snowflake/model_spdpp.py new file mode 100644 index 0000000..2cbb192 --- /dev/null +++ b/tgs/models/snowflake/model_spdpp.py @@ -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 \ No newline at end of file diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/__init__.py b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/__init__.py new file mode 100644 index 0000000..5fd361f --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/__init__.py @@ -0,0 +1,3 @@ +import pointnet2_ops.pointnet2_modules +import pointnet2_ops.pointnet2_utils +from pointnet2_ops._version import __version__ diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/ball_query.h b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/ball_query.h new file mode 100644 index 0000000..1bbc638 --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/ball_query.h @@ -0,0 +1,5 @@ +#pragma once +#include + +at::Tensor ball_query(at::Tensor new_xyz, at::Tensor xyz, const float radius, + const int nsample); diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/cuda_utils.h b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/cuda_utils.h new file mode 100644 index 0000000..0fd5b6e --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/cuda_utils.h @@ -0,0 +1,41 @@ +#ifndef _CUDA_UTILS_H +#define _CUDA_UTILS_H + +#include +#include +#include + +#include +#include + +#include + +#define TOTAL_THREADS 512 + +inline int opt_n_threads(int work_size) { + const int pow_2 = std::log(static_cast(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 diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/group_points.h b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/group_points.h new file mode 100644 index 0000000..ad20cda --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/group_points.h @@ -0,0 +1,5 @@ +#pragma once +#include + +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); diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/interpolate.h b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/interpolate.h new file mode 100644 index 0000000..26b3464 --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/interpolate.h @@ -0,0 +1,10 @@ +#pragma once + +#include +#include + +std::vector 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); diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/sampling.h b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/sampling.h new file mode 100644 index 0000000..d795271 --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/sampling.h @@ -0,0 +1,6 @@ +#pragma once +#include + +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); diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/utils.h b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/utils.h new file mode 100644 index 0000000..5f080ed --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/include/utils.h @@ -0,0 +1,25 @@ +#pragma once +#include +#include + +#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) diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query.cpp b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query.cpp new file mode 100644 index 0000000..b1797c1 --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query.cpp @@ -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(), + xyz.data_ptr(), idx.data_ptr()); + } else { + AT_ASSERT(false, "CPU not supported"); + } + + return idx; +} diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query_gpu.cu b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query_gpu.cu new file mode 100644 index 0000000..559aef9 --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/ball_query_gpu.cu @@ -0,0 +1,54 @@ +#include +#include +#include + +#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, n, m, radius, nsample, new_xyz, xyz, idx); + + CUDA_CHECK_ERRORS(); +} diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/bindings.cpp b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/bindings.cpp new file mode 100644 index 0000000..d1916ce --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/bindings.cpp @@ -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); +} diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points.cpp b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points.cpp new file mode 100644 index 0000000..285a4bd --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points.cpp @@ -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(), idx.data_ptr(), + output.data_ptr()); + } 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(), idx.data_ptr(), + output.data_ptr()); + } else { + AT_ASSERT(false, "CPU not supported"); + } + + return output; +} diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points_gpu.cu b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points_gpu.cu new file mode 100644 index 0000000..57c2b1b --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/group_points_gpu.cu @@ -0,0 +1,75 @@ +#include +#include + +#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, 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, c, n, npoints, nsample, grad_out, idx, grad_points); + + CUDA_CHECK_ERRORS(); +} diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate.cpp b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate.cpp new file mode 100644 index 0000000..cdee31c --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate.cpp @@ -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 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(), knows.data_ptr(), + dist2.data_ptr(), idx.data_ptr()); + } 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(), idx.data_ptr(), weight.data_ptr(), + output.data_ptr()); + } 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(), idx.data_ptr(), + weight.data_ptr(), output.data_ptr()); + } else { + AT_ASSERT(false, "CPU not supported"); + } + + return output; +} diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate_gpu.cu b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate_gpu.cu new file mode 100644 index 0000000..81c5548 --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/interpolate_gpu.cu @@ -0,0 +1,154 @@ +#include +#include +#include + +#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, 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, 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, c, n, m, grad_out, idx, weight, grad_points); + + CUDA_CHECK_ERRORS(); +} diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling.cpp b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling.cpp new file mode 100644 index 0000000..ddbdc11 --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling.cpp @@ -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(), + idx.data_ptr(), output.data_ptr()); + } 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(), + idx.data_ptr(), + output.data_ptr()); + } 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(), + tmp.data_ptr(), output.data_ptr()); + } else { + AT_ASSERT(false, "CPU not supported"); + } + + return output; +} diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling_gpu.cu b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling_gpu.cu new file mode 100644 index 0000000..fc573f0 --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_ext-src/src/sampling_gpu.cu @@ -0,0 +1,229 @@ +#include +#include + +#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<<>>(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<<>>( + 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 +__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, m, dataset, temp, idxs); + break; + case 256: + furthest_point_sampling_kernel<256> + <<>>(b, n, m, dataset, temp, idxs); + break; + case 128: + furthest_point_sampling_kernel<128> + <<>>(b, n, m, dataset, temp, idxs); + break; + case 64: + furthest_point_sampling_kernel<64> + <<>>(b, n, m, dataset, temp, idxs); + break; + case 32: + furthest_point_sampling_kernel<32> + <<>>(b, n, m, dataset, temp, idxs); + break; + case 16: + furthest_point_sampling_kernel<16> + <<>>(b, n, m, dataset, temp, idxs); + break; + case 8: + furthest_point_sampling_kernel<8> + <<>>(b, n, m, dataset, temp, idxs); + break; + case 4: + furthest_point_sampling_kernel<4> + <<>>(b, n, m, dataset, temp, idxs); + break; + case 2: + furthest_point_sampling_kernel<2> + <<>>(b, n, m, dataset, temp, idxs); + break; + case 1: + furthest_point_sampling_kernel<1> + <<>>(b, n, m, dataset, temp, idxs); + break; + default: + furthest_point_sampling_kernel<512> + <<>>(b, n, m, dataset, temp, idxs); + } + + CUDA_CHECK_ERRORS(); +} diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_version.py b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_version.py new file mode 100644 index 0000000..528787c --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/_version.py @@ -0,0 +1 @@ +__version__ = "3.0.0" diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/pointnet2_modules.py b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/pointnet2_modules.py new file mode 100644 index 0000000..a0ad4f6 --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/pointnet2_modules.py @@ -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) diff --git a/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/pointnet2_utils.py b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/pointnet2_utils.py new file mode 100644 index 0000000..150fccc --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/pointnet2_ops/pointnet2_utils.py @@ -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 diff --git a/tgs/models/snowflake/pointnet2_ops_lib/setup.py b/tgs/models/snowflake/pointnet2_ops_lib/setup.py new file mode 100644 index 0000000..5ebde53 --- /dev/null +++ b/tgs/models/snowflake/pointnet2_ops_lib/setup.py @@ -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, +) diff --git a/tgs/models/snowflake/skip_transformer.py b/tgs/models/snowflake/skip_transformer.py new file mode 100644 index 0000000..36462f7 --- /dev/null +++ b/tgs/models/snowflake/skip_transformer.py @@ -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 diff --git a/tgs/models/snowflake/utils.py b/tgs/models/snowflake/utils.py new file mode 100644 index 0000000..6661e38 --- /dev/null +++ b/tgs/models/snowflake/utils.py @@ -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 \ No newline at end of file diff --git a/tgs/models/tokenizers/dinov2.py b/tgs/models/tokenizers/dinov2.py new file mode 100644 index 0000000..853aa89 --- /dev/null +++ b/tgs/models/tokenizers/dinov2.py @@ -0,0 +1,1175 @@ +# coding=utf-8 +# Copyright 2023 Meta AI and The HuggingFace Inc. 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. +""" PyTorch DINOv2 model.""" + + +import collections.abc +import math +from typing import Dict, List, Optional, Set, Tuple, Union +from dataclasses import dataclass + +import torch +import torch.utils.checkpoint +from torch import nn +from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss + +from transformers.activations import ACT2FN +from transformers.modeling_outputs import ( + BackboneOutput, + BaseModelOutput, + BaseModelOutputWithPooling, + ImageClassifierOutput, +) +from transformers.modeling_utils import PreTrainedModel +from transformers.pytorch_utils import ( + find_pruneable_heads_and_indices, + prune_linear_layer, +) +from transformers.utils import ( + add_code_sample_docstrings, + add_start_docstrings, + add_start_docstrings_to_model_forward, + logging, + replace_return_docstrings, +) +from transformers.utils.backbone_utils import BackboneMixin +from transformers.models.dinov2.configuration_dinov2 import Dinov2Config + +from tgs.models.transformers import MemoryEfficientAttentionMixin +from tgs.utils.typing import * + + +logger = logging.get_logger(__name__) + +# General docstring +_CONFIG_FOR_DOC = "Dinov2Config" + +# Base docstring +_CHECKPOINT_FOR_DOC = "facebook/dinov2-base" +_EXPECTED_OUTPUT_SHAPE = [1, 257, 768] + +# Image classification docstring +_IMAGE_CLASS_CHECKPOINT = "facebook/dinov2-base" + + +DINOV2_PRETRAINED_MODEL_ARCHIVE_LIST = [ + "facebook/dinov2-base", + # See all DINOv2 models at https://huggingface.co/models?filter=dinov2 +] + + +class Dinov2Embeddings(nn.Module): + """ + Construct the CLS token, mask token, position and patch embeddings. + """ + + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + + self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size)) + # register as mask token as it's not used in optimization + # to avoid the use of find_unused_parameters_true + # self.mask_token = nn.Parameter(torch.zeros(1, config.hidden_size)) + self.register_buffer("mask_token", torch.zeros(1, config.hidden_size)) + self.patch_embeddings = Dinov2PatchEmbeddings(config) + num_patches = self.patch_embeddings.num_patches + self.position_embeddings = nn.Parameter( + torch.randn(1, num_patches + 1, config.hidden_size) + ) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + self.config = config + + def interpolate_pos_encoding( + self, embeddings: torch.Tensor, height: int, width: int + ) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher + resolution images. + + Source: + https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174 + """ + + num_patches = embeddings.shape[1] - 1 + num_positions = self.position_embeddings.shape[1] - 1 + if num_patches == num_positions and height == width: + return self.position_embeddings + class_pos_embed = self.position_embeddings[:, 0] + patch_pos_embed = self.position_embeddings[:, 1:] + dim = embeddings.shape[-1] + height = height // self.config.patch_size + width = width // self.config.patch_size + # we add a small number to avoid floating point error in the interpolation + # see discussion at https://github.com/facebookresearch/dino/issues/8 + height, width = height + 0.1, width + 0.1 + patch_pos_embed = patch_pos_embed.reshape( + 1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim + ) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed, + scale_factor=( + height / math.sqrt(num_positions), + width / math.sqrt(num_positions), + ), + mode="bicubic", + align_corners=False, + ) + if ( + int(height) != patch_pos_embed.shape[-2] + or int(width) != patch_pos_embed.shape[-1] + ): + raise ValueError( + "Width or height does not match with the interpolated position embeddings" + ) + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) + + def forward( + self, pixel_values: torch.Tensor, bool_masked_pos: Optional[torch.Tensor] = None, + ) -> torch.Tensor: + batch_size, _, height, width = pixel_values.shape + patch_embeddings = self.patch_embeddings(pixel_values) + embeddings = patch_embeddings + + if bool_masked_pos is not None: + embeddings = torch.where( + bool_masked_pos.unsqueeze(-1), + self.mask_token.to(embeddings.dtype).unsqueeze(0), + embeddings, + ) + + # add the [CLS] token to the embedded patch tokens + cls_tokens = self.cls_token.expand(batch_size, -1, -1) + embeddings = torch.cat((cls_tokens, embeddings), dim=1) + + # add positional encoding to each token + embeddings = embeddings + self.interpolate_pos_encoding( + embeddings, height, width + ) + + embeddings = self.dropout(embeddings) + + return embeddings + + +class Dinov2PatchEmbeddings(nn.Module): + """ + This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial + `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a + Transformer. + """ + + def __init__(self, config): + super().__init__() + image_size, patch_size = config.image_size, config.patch_size + num_channels, hidden_size = config.num_channels, config.hidden_size + + image_size = ( + image_size + if isinstance(image_size, collections.abc.Iterable) + else (image_size, image_size) + ) + patch_size = ( + patch_size + if isinstance(patch_size, collections.abc.Iterable) + else (patch_size, patch_size) + ) + num_patches = (image_size[1] // patch_size[1]) * ( + image_size[0] // patch_size[0] + ) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + + self.projection = nn.Conv2d( + num_channels, hidden_size, kernel_size=patch_size, stride=patch_size + ) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + num_channels = pixel_values.shape[1] + if num_channels != self.num_channels: + raise ValueError( + "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + f" Expected {self.num_channels} but got {num_channels}." + ) + embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2) + return embeddings + + +# Copied from transformers.models.vit.modeling_vit.ViTSelfAttention with ViT->Dinov2 +class Dinov2SelfAttention(nn.Module): + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + if config.hidden_size % config.num_attention_heads != 0 and not hasattr( + config, "embedding_size" + ): + raise ValueError( + f"The hidden size {config.hidden_size,} is not a multiple of the number of attention " + f"heads {config.num_attention_heads}." + ) + + self.num_attention_heads = config.num_attention_heads + self.attention_head_size = int(config.hidden_size / config.num_attention_heads) + self.all_head_size = self.num_attention_heads * self.attention_head_size + self.attention_probs_dropout_prob = config.attention_probs_dropout_prob + + self.query = nn.Linear( + config.hidden_size, self.all_head_size, bias=config.qkv_bias + ) + self.key = nn.Linear( + config.hidden_size, self.all_head_size, bias=config.qkv_bias + ) + self.value = nn.Linear( + config.hidden_size, self.all_head_size, bias=config.qkv_bias + ) + + self.dropout = nn.Dropout(config.attention_probs_dropout_prob) + + self.use_memory_efficient_attention_xformers: bool = False + + def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor: + new_x_shape = x.size()[:-1] + ( + self.num_attention_heads, + self.attention_head_size, + ) + x = x.view(new_x_shape) + return x.permute(0, 2, 1, 3) + + def forward( + self, + hidden_states, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + mixed_query_layer = self.query(hidden_states) + + if self.use_memory_efficient_attention_xformers: + import xformers + assert head_mask is None and not output_attentions + new_size = hidden_states.size()[:-1] + ( + self.num_attention_heads, + self.attention_head_size, + ) + key_layer = self.key(hidden_states).view(new_size) + value_layer = self.value(hidden_states).view(new_size) + query_layer = mixed_query_layer.view(new_size) + context_layer = xformers.ops.memory_efficient_attention( + query_layer, key_layer, value_layer, p=self.attention_probs_dropout_prob + ) + context_layer = context_layer.view(*hidden_states.size()[:-1], -1) + else: + key_layer = self.transpose_for_scores(self.key(hidden_states)) + value_layer = self.transpose_for_scores(self.value(hidden_states)) + query_layer = self.transpose_for_scores(mixed_query_layer) + + # Take the dot product between "query" and "key" to get the raw attention scores. + attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2)) + + attention_scores = attention_scores / math.sqrt(self.attention_head_size) + + # Normalize the attention scores to probabilities. + attention_probs = nn.functional.softmax(attention_scores, dim=-1) + + # This is actually dropping out entire tokens to attend to, which might + # seem a bit unusual, but is taken from the original Transformer paper. + attention_probs = self.dropout(attention_probs) + + # Mask heads if we want to + if head_mask is not None: + attention_probs = attention_probs * head_mask + + context_layer = torch.matmul(attention_probs, value_layer) + + context_layer = context_layer.permute(0, 2, 1, 3).contiguous() + new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,) + context_layer = context_layer.view(new_context_layer_shape) + + outputs = ( + (context_layer, attention_probs) if output_attentions else (context_layer,) + ) + + return outputs + + def set_use_memory_efficient_attention_xformers( + self, valid: bool, attention_op: Optional[Callable] = None + ): + self.use_memory_efficient_attention_xformers = valid + + +# Copied from transformers.models.vit.modeling_vit.ViTSelfOutput with ViT->Dinov2 +class Dinov2SelfOutput(nn.Module): + """ + The residual connection is defined in Dinov2Layer instead of here (as is the case with other models), due to the + layernorm applied before each block. + """ + + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + self.dense = nn.Linear(config.hidden_size, config.hidden_size) + self.dropout = nn.Dropout(config.hidden_dropout_prob) + + def forward( + self, hidden_states: torch.Tensor, input_tensor: torch.Tensor + ) -> torch.Tensor: + hidden_states = self.dense(hidden_states) + hidden_states = self.dropout(hidden_states) + + return hidden_states + + +# Copied from transformers.models.vit.modeling_vit.ViTAttention with ViT->Dinov2 +class Dinov2Attention(nn.Module): + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + self.attention = Dinov2SelfAttention(config) + self.output = Dinov2SelfOutput(config) + self.pruned_heads = set() + + def prune_heads(self, heads: Set[int]) -> None: + if len(heads) == 0: + return + heads, index = find_pruneable_heads_and_indices( + heads, + self.attention.num_attention_heads, + self.attention.attention_head_size, + self.pruned_heads, + ) + + # Prune linear layers + self.attention.query = prune_linear_layer(self.attention.query, index) + self.attention.key = prune_linear_layer(self.attention.key, index) + self.attention.value = prune_linear_layer(self.attention.value, index) + self.output.dense = prune_linear_layer(self.output.dense, index, dim=1) + + # Update hyper params and store pruned heads + self.attention.num_attention_heads = self.attention.num_attention_heads - len( + heads + ) + self.attention.all_head_size = ( + self.attention.attention_head_size * self.attention.num_attention_heads + ) + self.pruned_heads = self.pruned_heads.union(heads) + + def forward( + self, + hidden_states: torch.Tensor, + head_mask: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + self_outputs = self.attention(hidden_states, head_mask, output_attentions) + + attention_output = self.output(self_outputs[0], hidden_states) + + outputs = (attention_output,) + self_outputs[ + 1: + ] # add attentions if we output them + return outputs + + +class Dinov2LayerScale(nn.Module): + def __init__(self, config) -> None: + super().__init__() + self.lambda1 = nn.Parameter( + config.layerscale_value * torch.ones(config.hidden_size) + ) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + return hidden_state * self.lambda1 + + +# Copied from transformers.models.beit.modeling_beit.drop_path +def drop_path( + input: torch.Tensor, drop_prob: float = 0.0, training: bool = False +) -> torch.Tensor: + """ + Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks). + + Comment by Ross Wightman: This is the same as the DropConnect impl I created for EfficientNet, etc networks, + however, the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper... + See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for changing the + layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use 'survival rate' as the + argument. + """ + if drop_prob == 0.0 or not training: + return input + keep_prob = 1 - drop_prob + shape = (input.shape[0],) + (1,) * ( + input.ndim - 1 + ) # work with diff dim tensors, not just 2D ConvNets + random_tensor = keep_prob + torch.rand( + shape, dtype=input.dtype, device=input.device + ) + random_tensor.floor_() # binarize + output = input.div(keep_prob) * random_tensor + return output + + +# Copied from transformers.models.beit.modeling_beit.BeitDropPath +class Dinov2DropPath(nn.Module): + """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).""" + + def __init__(self, drop_prob: Optional[float] = None) -> None: + super().__init__() + self.drop_prob = drop_prob + + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + return drop_path(hidden_states, self.drop_prob, self.training) + + def extra_repr(self) -> str: + return "p={}".format(self.drop_prob) + + +class Dinov2MLP(nn.Module): + def __init__(self, config) -> None: + super().__init__() + in_features = out_features = config.hidden_size + hidden_features = int(config.hidden_size * config.mlp_ratio) + self.fc1 = nn.Linear(in_features, hidden_features, bias=True) + if isinstance(config.hidden_act, str): + self.activation = ACT2FN[config.hidden_act] + else: + self.activation = config.hidden_act + self.fc2 = nn.Linear(hidden_features, out_features, bias=True) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + hidden_state = self.fc1(hidden_state) + hidden_state = self.activation(hidden_state) + hidden_state = self.fc2(hidden_state) + return hidden_state + + +class Dinov2SwiGLUFFN(nn.Module): + def __init__(self, config) -> None: + super().__init__() + in_features = out_features = config.hidden_size + hidden_features = int(config.hidden_size * config.mlp_ratio) + hidden_features = (int(hidden_features * 2 / 3) + 7) // 8 * 8 + + self.weights_in = nn.Linear(in_features, 2 * hidden_features, bias=True) + self.weights_out = nn.Linear(hidden_features, out_features, bias=True) + + def forward(self, hidden_state: torch.Tensor) -> torch.Tensor: + hidden_state = self.weights_in(hidden_state) + x1, x2 = hidden_state.chunk(2, dim=-1) + hidden = nn.functional.silu(x1) * x2 + return self.weights_out(hidden) + + +class Dinov2Layer(nn.Module, MemoryEfficientAttentionMixin): + """This corresponds to the Block class in the original implementation.""" + + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + + self.norm1 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.norm1_modulation = None + self.attention = Dinov2Attention(config) + self.layer_scale1 = Dinov2LayerScale(config) + self.drop_path1 = ( + Dinov2DropPath(config.drop_path_rate) + if config.drop_path_rate > 0.0 + else nn.Identity() + ) + + self.norm2 = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + self.norm2_modulation = None + + if config.use_swiglu_ffn: + self.mlp = Dinov2SwiGLUFFN(config) + else: + self.mlp = Dinov2MLP(config) + self.layer_scale2 = Dinov2LayerScale(config) + self.drop_path2 = ( + Dinov2DropPath(config.drop_path_rate) + if config.drop_path_rate > 0.0 + else nn.Identity() + ) + + def forward( + self, + hidden_states: torch.Tensor, + head_mask: Optional[torch.Tensor] = None, + modulation_cond: Optional[torch.Tensor] = None, + output_attentions: bool = False, + ) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]: + hidden_states_norm = self.norm1(hidden_states) + if self.norm1_modulation is not None: + assert modulation_cond is not None + hidden_states_norm = self.norm1_modulation( + hidden_states_norm, modulation_cond + ) + self_attention_outputs = self.attention( + hidden_states_norm, # in Dinov2, layernorm is applied before self-attention + head_mask, + output_attentions=output_attentions, + ) + attention_output = self_attention_outputs[0] + + attention_output = self.layer_scale1(attention_output) + outputs = self_attention_outputs[ + 1: + ] # add self attentions if we output attention weights + + # first residual connection + hidden_states = attention_output + hidden_states + + # in Dinov2, layernorm is also applied after self-attention + layer_output = self.norm2(hidden_states) + if self.norm2_modulation is not None: + assert modulation_cond is not None + layer_output = self.norm2_modulation(layer_output, modulation_cond) + layer_output = self.mlp(layer_output) + layer_output = self.layer_scale2(layer_output) + + # second residual connection + layer_output = layer_output + hidden_states + + outputs = (layer_output,) + outputs + + return outputs + + def register_ada_norm_modulation(self, norm1_mod: nn.Module, norm2_mod: nn.Module): + self.norm1_modulation = norm1_mod + self.norm2_modulation = norm2_mod + + +# Copied from transformers.models.vit.modeling_vit.ViTEncoder with ViT->Dinov2 +class Dinov2Encoder(nn.Module, MemoryEfficientAttentionMixin): + def __init__(self, config: Dinov2Config) -> None: + super().__init__() + self.config = config + self.layer = nn.ModuleList( + [Dinov2Layer(config) for _ in range(config.num_hidden_layers)] + ) + self.gradient_checkpointing = False + + def forward( + self, + hidden_states: torch.Tensor, + head_mask: Optional[torch.Tensor] = None, + modulation_cond: Optional[torch.Tensor] = None, + output_attentions: bool = False, + output_hidden_states: bool = False, + return_dict: bool = True, + ) -> Union[tuple, BaseModelOutput]: + all_hidden_states = () if output_hidden_states else None + all_self_attentions = () if output_attentions else None + + for i, layer_module in enumerate(self.layer): + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + layer_head_mask = head_mask[i] if head_mask is not None else None + + if self.gradient_checkpointing and self.training: + + def create_custom_forward(module): + def custom_forward(*inputs): + return module(*inputs, output_attentions) + + return custom_forward + + layer_outputs = torch.utils.checkpoint.checkpoint( + create_custom_forward(layer_module), + hidden_states, + layer_head_mask, + modulation_cond, + ) + else: + layer_outputs = layer_module( + hidden_states, layer_head_mask, modulation_cond, output_attentions + ) + + hidden_states = layer_outputs[0] + + if output_attentions: + all_self_attentions = all_self_attentions + (layer_outputs[1],) + + if output_hidden_states: + all_hidden_states = all_hidden_states + (hidden_states,) + + if not return_dict: + return tuple( + v + for v in [hidden_states, all_hidden_states, all_self_attentions] + if v is not None + ) + return BaseModelOutput( + last_hidden_state=hidden_states, + hidden_states=all_hidden_states, + attentions=all_self_attentions, + ) + + +class Dinov2PreTrainedModel(PreTrainedModel): + """ + An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained + models. + """ + + config_class = Dinov2Config + base_model_prefix = "dinov2" + main_input_name = "pixel_values" + supports_gradient_checkpointing = True + + def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None: + """Initialize the weights""" + if isinstance(module, (nn.Linear, nn.Conv2d)): + # Upcast the input in `fp32` and cast it back to desired `dtype` to avoid + # `trunc_normal_cpu` not implemented in `half` issues + module.weight.data = nn.init.trunc_normal_( + module.weight.data.to(torch.float32), + mean=0.0, + std=self.config.initializer_range, + ).to(module.weight.dtype) + if module.bias is not None: + module.bias.data.zero_() + elif isinstance(module, nn.LayerNorm): + module.bias.data.zero_() + module.weight.data.fill_(1.0) + elif isinstance(module, Dinov2Embeddings): + module.position_embeddings.data = nn.init.trunc_normal_( + module.position_embeddings.data.to(torch.float32), + mean=0.0, + std=self.config.initializer_range, + ).to(module.position_embeddings.dtype) + + module.cls_token.data = nn.init.trunc_normal_( + module.cls_token.data.to(torch.float32), + mean=0.0, + std=self.config.initializer_range, + ).to(module.cls_token.dtype) + + def _set_gradient_checkpointing( + self, module: Dinov2Encoder, value: bool = False + ) -> None: + if isinstance(module, Dinov2Encoder): + module.gradient_checkpointing = value + + +DINOV2_START_DOCSTRING = r""" + This model is a PyTorch [torch.nn.Module](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) subclass. Use it + as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general usage and + behavior. + + Parameters: + config ([`Dinov2Config`]): Model configuration class with all the parameters of the model. + Initializing with a config file does not load the weights associated with the model, only the + configuration. Check out the [`~PreTrainedModel.from_pretrained`] method to load the model weights. +""" + +DINOV2_BASE_INPUTS_DOCSTRING = r""" + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See + [`BitImageProcessor.preprocess`] for details. + + bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, sequence_length)`): + Boolean masked positions. Indicates which patches are masked (1) and which aren't (0). Only relevant for + pre-training. + + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + +DINOV2_INPUTS_DOCSTRING = r""" + Args: + pixel_values (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`): + Pixel values. Pixel values can be obtained using [`AutoImageProcessor`]. See + [`BitImageProcessor.preprocess`] for details. + + head_mask (`torch.FloatTensor` of shape `(num_heads,)` or `(num_layers, num_heads)`, *optional*): + Mask to nullify selected heads of the self-attention modules. Mask values selected in `[0, 1]`: + + - 1 indicates the head is **not masked**, + - 0 indicates the head is **masked**. + + output_attentions (`bool`, *optional*): + Whether or not to return the attentions tensors of all attention layers. See `attentions` under returned + tensors for more detail. + output_hidden_states (`bool`, *optional*): + Whether or not to return the hidden states of all layers. See `hidden_states` under returned tensors for + more detail. + return_dict (`bool`, *optional*): + Whether or not to return a [`~utils.ModelOutput`] instead of a plain tuple. +""" + +@dataclass +class CustomBaseModelOutputWithPooling(BaseModelOutputWithPooling): + patch_embeddings: Optional[torch.FloatTensor] = None + + +@add_start_docstrings( + "The bare DINOv2 Model transformer outputting raw hidden-states without any specific head on top.", + DINOV2_START_DOCSTRING, +) +class Dinov2Model(Dinov2PreTrainedModel, MemoryEfficientAttentionMixin): + def __init__(self, config: Dinov2Config): + super().__init__(config) + self.config = config + + self.embeddings = Dinov2Embeddings(config) + self.encoder = Dinov2Encoder(config) + + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> Dinov2PatchEmbeddings: + return self.embeddings.patch_embeddings + + def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None: + """ + Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base + class PreTrainedModel + """ + for layer, heads in heads_to_prune.items(): + self.encoder.layer[layer].attention.prune_heads(heads) + + @add_start_docstrings_to_model_forward(DINOV2_BASE_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_CHECKPOINT_FOR_DOC, + output_type=BaseModelOutputWithPooling, + config_class=_CONFIG_FOR_DOC, + modality="vision", + expected_output=_EXPECTED_OUTPUT_SHAPE, + ) + def forward( + self, + pixel_values: Optional[torch.Tensor] = None, + bool_masked_pos: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + modulation_cond: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[Tuple, BaseModelOutputWithPooling]: + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + if pixel_values is None: + raise ValueError("You have to specify pixel_values") + + # Prepare head mask if needed + # 1.0 in head_mask indicate we keep the head + # attention_probs has shape bsz x n_heads x N x N + # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] + # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length] + head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers) + + embedding_output = self.embeddings( + pixel_values, bool_masked_pos=bool_masked_pos + ) + + encoder_outputs = self.encoder( + embedding_output, + head_mask=head_mask, + modulation_cond=modulation_cond, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + sequence_output = encoder_outputs[0] + sequence_output = self.layernorm(sequence_output) + pooled_output = sequence_output[:, 0, :] + + if not return_dict: + head_outputs = (sequence_output, pooled_output) + return head_outputs + encoder_outputs[1:] + + return CustomBaseModelOutputWithPooling( + last_hidden_state=sequence_output, + pooler_output=pooled_output, + hidden_states=encoder_outputs.hidden_states, + attentions=encoder_outputs.attentions, + patch_embeddings=embedding_output + ) + + def set_gradient_checkpointing(self, value: bool = False) -> None: + self._set_gradient_checkpointing(self.encoder, value) + + +@add_start_docstrings( + """ + Dinov2 Model transformer with an image classification head on top (a linear layer on top of the final hidden state + of the [CLS] token) e.g. for ImageNet. + """, + DINOV2_START_DOCSTRING, +) +class Dinov2ForImageClassification(Dinov2PreTrainedModel): + def __init__(self, config: Dinov2Config) -> None: + super().__init__(config) + + self.num_labels = config.num_labels + self.dinov2 = Dinov2Model(config) + + # Classifier head + self.classifier = ( + nn.Linear(config.hidden_size * 2, config.num_labels) + if config.num_labels > 0 + else nn.Identity() + ) + + # Initialize weights and apply final processing + self.post_init() + + @add_start_docstrings_to_model_forward(DINOV2_INPUTS_DOCSTRING) + @add_code_sample_docstrings( + checkpoint=_IMAGE_CLASS_CHECKPOINT, + output_type=ImageClassifierOutput, + config_class=_CONFIG_FOR_DOC, + ) + def forward( + self, + pixel_values: Optional[torch.Tensor] = None, + head_mask: Optional[torch.Tensor] = None, + labels: Optional[torch.Tensor] = None, + output_attentions: Optional[bool] = None, + output_hidden_states: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> Union[tuple, ImageClassifierOutput]: + r""" + labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*): + Labels for computing the image classification/regression loss. Indices should be in `[0, ..., + config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If + `config.num_labels > 1` a classification loss is computed (Cross-Entropy). + """ + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + + outputs = self.dinov2( + pixel_values, + head_mask=head_mask, + output_attentions=output_attentions, + output_hidden_states=output_hidden_states, + return_dict=return_dict, + ) + + sequence_output = outputs[0] # batch_size, sequence_length, hidden_size + + cls_token = sequence_output[:, 0] + patch_tokens = sequence_output[:, 1:] + + linear_input = torch.cat([cls_token, patch_tokens.mean(dim=1)], dim=1) + + logits = self.classifier(linear_input) + + loss = None + if labels is not None: + # move labels to correct device to enable model parallelism + labels = labels.to(logits.device) + if self.config.problem_type is None: + if self.num_labels == 1: + self.config.problem_type = "regression" + elif self.num_labels > 1 and ( + labels.dtype == torch.long or labels.dtype == torch.int + ): + self.config.problem_type = "single_label_classification" + else: + self.config.problem_type = "multi_label_classification" + + if self.config.problem_type == "regression": + loss_fct = MSELoss() + if self.num_labels == 1: + loss = loss_fct(logits.squeeze(), labels.squeeze()) + else: + loss = loss_fct(logits, labels) + elif self.config.problem_type == "single_label_classification": + loss_fct = CrossEntropyLoss() + loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1)) + elif self.config.problem_type == "multi_label_classification": + loss_fct = BCEWithLogitsLoss() + loss = loss_fct(logits, labels) + + if not return_dict: + output = (logits,) + outputs[2:] + return ((loss,) + output) if loss is not None else output + + return ImageClassifierOutput( + loss=loss, + logits=logits, + hidden_states=outputs.hidden_states, + attentions=outputs.attentions, + ) + + +@add_start_docstrings( + """ + Dinov2 backbone, to be used with frameworks like DETR and MaskFormer. + """, + DINOV2_START_DOCSTRING, +) +class Dinov2Backbone(Dinov2PreTrainedModel, BackboneMixin): + def __init__(self, config): + super().__init__(config) + super()._init_backbone(config) + + self.num_features = [ + config.hidden_size for _ in range(config.num_hidden_layers + 1) + ] + self.embeddings = Dinov2Embeddings(config) + self.encoder = Dinov2Encoder(config) + + self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps) + + # Initialize weights and apply final processing + self.post_init() + + def get_input_embeddings(self) -> Dinov2PatchEmbeddings: + return self.embeddings.patch_embeddings + + @add_start_docstrings_to_model_forward(DINOV2_INPUTS_DOCSTRING) + @replace_return_docstrings(output_type=BackboneOutput, config_class=_CONFIG_FOR_DOC) + def forward( + self, + pixel_values: torch.Tensor, + output_hidden_states: Optional[bool] = None, + output_attentions: Optional[bool] = None, + return_dict: Optional[bool] = None, + ) -> BackboneOutput: + """ + Returns: + + Examples: + + ```python + >>> from transformers import AutoImageProcessor, AutoBackbone + >>> import torch + >>> from PIL import Image + >>> import requests + + >>> url = "http://images.cocodataset.org/val2017/000000039769.jpg" + >>> image = Image.open(requests.get(url, stream=True).raw) + + >>> processor = AutoImageProcessor.from_pretrained("facebook/dinov2-base") + >>> model = AutoBackbone.from_pretrained( + ... "facebook/dinov2-base", out_features=["stage2", "stage5", "stage8", "stage11"] + ... ) + + >>> inputs = processor(image, return_tensors="pt") + + >>> outputs = model(**inputs) + >>> feature_maps = outputs.feature_maps + >>> list(feature_maps[-1].shape) + [1, 768, 16, 16] + ```""" + return_dict = ( + return_dict if return_dict is not None else self.config.use_return_dict + ) + output_hidden_states = ( + output_hidden_states + if output_hidden_states is not None + else self.config.output_hidden_states + ) + output_attentions = ( + output_attentions + if output_attentions is not None + else self.config.output_attentions + ) + + embedding_output = self.embeddings(pixel_values) + + outputs = self.encoder( + embedding_output, + output_hidden_states=True, + output_attentions=output_attentions, + return_dict=return_dict, + ) + + hidden_states = outputs.hidden_states if return_dict else outputs[1] + + feature_maps = () + for stage, hidden_state in zip(self.stage_names, hidden_states): + if stage in self.out_features: + if self.config.apply_layernorm: + hidden_state = self.layernorm(hidden_state) + if self.config.reshape_hidden_states: + batch_size, _, height, width = pixel_values.shape + patch_size = self.config.patch_size + hidden_state = hidden_state[:, 1:, :].reshape( + batch_size, width // patch_size, height // patch_size, -1 + ) + hidden_state = hidden_state.permute(0, 3, 1, 2).contiguous() + feature_maps += (hidden_state,) + + if not return_dict: + if output_hidden_states: + output = (feature_maps,) + outputs[1:] + else: + output = (feature_maps,) + outputs[2:] + return output + + return BackboneOutput( + feature_maps=feature_maps, + hidden_states=outputs.hidden_states if output_hidden_states else None, + attentions=outputs.attentions if output_attentions else None, + ) + + + +class CustomPatchEmbeddings(nn.Module): + """ + This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial + `hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a + Transformer. + """ + + def __init__(self, image_size: int, patch_size: int, num_channels: int, hidden_size: int): + super().__init__() + + image_size = ( + image_size + if isinstance(image_size, collections.abc.Iterable) + else (image_size, image_size) + ) + patch_size = ( + patch_size + if isinstance(patch_size, collections.abc.Iterable) + else (patch_size, patch_size) + ) + num_patches = (image_size[1] // patch_size[1]) * ( + image_size[0] // patch_size[0] + ) + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.num_patches = num_patches + + self.projection = nn.Conv2d( + num_channels, hidden_size, kernel_size=patch_size, stride=patch_size + ) + + def forward(self, pixel_values: torch.Tensor) -> torch.Tensor: + num_channels = pixel_values.shape[1] + if num_channels != self.num_channels: + raise ValueError( + "Make sure that the channel dimension of the pixel values match with the one set in the configuration." + f" Expected {self.num_channels} but got {num_channels}." + ) + embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2) + return embeddings + + +class CustomEmbeddings(nn.Module): + """ + Construct the CLS token, mask token, position and patch embeddings. + """ + + def __init__(self, image_size: int, patch_size: int, num_channels: int, hidden_size: int) -> None: + super().__init__() + + self.image_size = image_size + self.patch_size = patch_size + self.num_channels = num_channels + self.hidden_size = hidden_size + + self.cls_token = nn.Parameter(torch.randn(1, 1, self.hidden_size)) + + self.patch_embeddings = CustomPatchEmbeddings(image_size, patch_size, num_channels, hidden_size) + num_patches = self.patch_embeddings.num_patches + self.position_embeddings = nn.Parameter( + torch.randn(1, num_patches + 1, self.hidden_size) + ) + + def interpolate_pos_encoding( + self, embeddings: torch.Tensor, height: int, width: int + ) -> torch.Tensor: + """ + This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher + resolution images. + + Source: + https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174 + """ + + num_patches = embeddings.shape[1] - 1 + num_positions = self.position_embeddings.shape[1] - 1 + if num_patches == num_positions and height == width: + return self.position_embeddings + class_pos_embed = self.position_embeddings[:, 0] + patch_pos_embed = self.position_embeddings[:, 1:] + dim = embeddings.shape[-1] + height = height // self.patch_size + width = width // self.patch_size + # we add a small number to avoid floating point error in the interpolation + # see discussion at https://github.com/facebookresearch/dino/issues/8 + height, width = height + 0.1, width + 0.1 + patch_pos_embed = patch_pos_embed.reshape( + 1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim + ) + patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2) + patch_pos_embed = nn.functional.interpolate( + patch_pos_embed, + scale_factor=( + height / math.sqrt(num_positions), + width / math.sqrt(num_positions), + ), + mode="bicubic", + align_corners=False, + ) + if ( + int(height) != patch_pos_embed.shape[-2] + or int(width) != patch_pos_embed.shape[-1] + ): + raise ValueError( + "Width or height does not match with the interpolated position embeddings" + ) + patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim) + return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1) + + def forward( + self, pixel_values: torch.Tensor, + ) -> torch.Tensor: + batch_size, _, height, width = pixel_values.shape + patch_embeddings = self.patch_embeddings(pixel_values) + embeddings = patch_embeddings + + # add the [CLS] token to the embedded patch tokens + cls_tokens = self.cls_token.expand(batch_size, -1, -1) + embeddings = torch.cat((cls_tokens, embeddings), dim=1) + + # add positional encoding to each token + embeddings = embeddings + self.interpolate_pos_encoding( + embeddings, height, width + ) + + return embeddings \ No newline at end of file diff --git a/tgs/models/tokenizers/image.py b/tgs/models/tokenizers/image.py new file mode 100644 index 0000000..eb8de83 --- /dev/null +++ b/tgs/models/tokenizers/image.py @@ -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 \ No newline at end of file diff --git a/tgs/models/tokenizers/point.py b/tgs/models/tokenizers/point.py new file mode 100644 index 0000000..99afb6c --- /dev/null +++ b/tgs/models/tokenizers/point.py @@ -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)) \ No newline at end of file diff --git a/tgs/models/tokenizers/triplane.py b/tgs/models/tokenizers/triplane.py new file mode 100644 index 0000000..8bd6774 --- /dev/null +++ b/tgs/models/tokenizers/triplane.py @@ -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, + ) diff --git a/tgs/models/transformers.py b/tgs/models/transformers.py new file mode 100644 index 0000000..82ccb14 --- /dev/null +++ b/tgs/models/transformers.py @@ -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. + + + + ⚠️ When memory efficient attention and sliced attention are both enabled, memory efficient attention takes + precedent. + + + + 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 diff --git a/tgs/utils/__init__.py b/tgs/utils/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/tgs/utils/base.py b/tgs/utils/base.py new file mode 100644 index 0000000..e98f63c --- /dev/null +++ b/tgs/utils/base.py @@ -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) diff --git a/tgs/utils/config.py b/tgs/utils/config.py new file mode 100644 index 0000000..abd8520 --- /dev/null +++ b/tgs/utils/config.py @@ -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 diff --git a/tgs/utils/misc.py b/tgs/utils/misc.py new file mode 100644 index 0000000..7daf634 --- /dev/null +++ b/tgs/utils/misc.py @@ -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))) diff --git a/tgs/utils/ops.py b/tgs/utils/ops.py new file mode 100644 index 0000000..29f0fd4 --- /dev/null +++ b/tgs/utils/ops.py @@ -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 \ No newline at end of file diff --git a/tgs/utils/saving.py b/tgs/utils/saving.py new file mode 100644 index 0000000..f8be1bc --- /dev/null +++ b/tgs/utils/saving.py @@ -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_) diff --git a/tgs/utils/typing.py b/tgs/utils/typing.py new file mode 100644 index 0000000..dee9f96 --- /dev/null +++ b/tgs/utils/typing.py @@ -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