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

Add Instant_NGP

Browse Source
This commit is contained in:
MrForExample
2024-04-01 15:50:58 +02:00
parent 6cd7b8f390
commit 086567f3e9
12 changed files with 2117 additions and 424 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+6 -2
View File
@@ -13,7 +13,7 @@
### Note: this project is still a WIP
## Currently support:
- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 12/03/2024**]
- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 01/04/2024**]
- **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
- [tripoSR-layered-diffusion workflow](https://github.com/C0nsumption/Consume-ComfyUI-Workflows/tree/main/assets/tripo_sr/00) by [@Consumption](https://twitter.com/c0nsumption_)
@@ -63,6 +63,10 @@
- Azimuth: In top view, from angle 0 rotate 360 degree with step -90 you get (0, -90, -180/180, 90, 0), in this case camera rotates clock-wise, vice versa.
- Elevation: 0 when camera points horizontally forward, pointing down to the ground is negitive angle, vice versa.
- **Instant NGP**: [nerfacc](https://github.com/nerfstudio-project/nerfacc)
- Multi-View images as inputs
- Export to 3D Mesh using marching cubes
- **3D Gaussian Splatting**
- [Improved Differential Gaussian Rasterization](https://github.com/ashawkey/diff-gaussian-rasterization)
- Better Compactness-based Densification method from [Gsgen](https://gsgen3d.github.io/),
@@ -74,7 +78,7 @@
- **Gaussian Splatting Orbit Renderer**
- Render 3DGS to images sequences or video, given a 3DGS file and camera poses generated by **Stack Orbit Camera Poses** node
- **Multi-View images to UVTexture**
- **Fitting_Mesh_With_Multiview_Images**
- Bake Multi-View images into UVTexture of given 3D mesh using [Nvdiffrast](https://github.com/NVlabs/nvdiffrast), supports:
- Export to .obj, .ply, .glb
_Example_Workflows/Multi-View-Images_to_3DGS_to_3DMesh(DMTet and DiffRast).json
+351 -396
View File
@@ -1,6 +1,6 @@
{
"last_node_id": 47,
"last_link_id": 76,
"last_node_id": 48,
"last_link_id": 81,
"nodes": [
{
"id": 9,
@@ -187,38 +187,6 @@
"Node name for S&R": "MaskToImage"
}
},
{
"id": 37,
"type": "[Comfy3D] Preview 3DMesh",
"pos": [
3740,
110
],
"size": [
600,
500
],
"flags": {},
"order": 29,
"mode": 0,
"inputs": [
{
"name": "mesh_file_path",
"type": "STRING",
"link": 55,
"widget": {
"name": "mesh_file_path"
}
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Preview 3DMesh"
},
"widgets_values": [
"",
null
]
},
{
"id": 39,
"type": "Reroute",
@@ -328,42 +296,6 @@
"horizontal": false
}
},
{
"id": 43,
"type": "Reroute",
"pos": [
2195,
-25
],
"size": [
75,
26
],
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 64
}
],
"outputs": [
{
"name": "",
"type": "ORBIT_CAMPOSES",
"links": [
65
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 35,
"type": "[Comfy3D] Convert 3DGS To Pointcloud",
@@ -414,7 +346,7 @@
"flags": {
"collapsed": false
},
"order": 26,
"order": 25,
"mode": 0,
"inputs": [
{
@@ -466,7 +398,7 @@
500
],
"flags": {},
"order": 27,
"order": 26,
"mode": 0,
"inputs": [
{
@@ -560,42 +492,6 @@
false
]
},
{
"id": 40,
"type": "Reroute",
"pos": [
2288,
-552
],
"size": [
75,
26
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 59
}
],
"outputs": [
{
"name": "",
"type": "IMAGE",
"links": [
69
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 32,
"type": "[Comfy3D] Gaussian Splatting Orbit Renderer",
@@ -655,246 +551,6 @@
0
]
},
{
"id": 42,
"type": "Reroute",
"pos": [
2241,
-484
],
"size": [
75,
26
],
"flags": {},
"order": 15,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 62
}
],
"outputs": [
{
"name": "",
"type": "MASK",
"links": [
63
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 36,
"type": "[Comfy3D] Save 3D Mesh",
"pos": [
3399,
110
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 28,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESH",
"link": 54
}
],
"outputs": [
{
"name": "save_path",
"type": "STRING",
"links": [
55
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Save 3D Mesh"
},
"widgets_values": [
"MeshTest\\3DMesh_%Y-%m-%d-%M-%S-%f.obj"
]
},
{
"id": 34,
"type": "[Comfy3D] Deep Marching Tetrahedrons",
"pos": [
2332,
249
],
"size": {
"0": 443.4000244140625,
"1": 314
},
"flags": {},
"order": 19,
"mode": 0,
"inputs": [
{
"name": "reference_points_cloud",
"type": "POINTCLOUD",
"link": 53
},
{
"name": "reference_images",
"type": "IMAGE",
"link": null
},
{
"name": "reference_masks",
"type": "MASK",
"link": null
}
],
"outputs": [
{
"name": "trained_mesh",
"type": "MESH",
"links": [
56,
68
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Deep Marching Tetrahedrons"
},
"widgets_values": [
3000,
1,
0.1,
1,
1,
128,
0.0001,
2,
128,
5
]
},
{
"id": 46,
"type": "[Comfy3D] Save 3D Mesh",
"pos": [
2843,
495
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 25,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESH",
"link": 68
}
],
"outputs": [
{
"name": "save_path",
"type": "STRING",
"links": [],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Save 3D Mesh"
},
"widgets_values": [
"MeshTest\\3DMesh_%Y-%m-%d-%M-%S-%f.obj"
]
},
{
"id": 38,
"type": "[Comfy3D] Bake Texture To Mesh",
"pos": [
2840,
110
],
"size": {
"0": 531.5999755859375,
"1": 334
},
"flags": {},
"order": 24,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 69
},
{
"name": "reference_masks",
"type": "MASK",
"link": 63
},
{
"name": "reference_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 65
},
{
"name": "mesh",
"type": "MESH",
"link": 56
}
],
"outputs": [
{
"name": "trained_mesh",
"type": "MESH",
"links": [
54
],
"shape": 3,
"slot_index": 0
},
{
"name": "baked_texture",
"type": "IMAGE",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Bake Texture To Mesh"
},
"widgets_values": [
49.1,
1024,
1024,
1000,
1,
0.1,
false,
0.01,
0.5,
false
]
},
{
"id": 7,
"type": "VHS_LoadImages",
@@ -1276,6 +932,313 @@
],
"color": "#432",
"bgcolor": "#653"
},
{
"id": 40,
"type": "Reroute",
"pos": [
2288,
-552
],
"size": [
75,
26
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 59
}
],
"outputs": [
{
"name": "",
"type": "IMAGE",
"links": [
77
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 42,
"type": "Reroute",
"pos": [
2241,
-484
],
"size": [
75,
26
],
"flags": {},
"order": 15,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 62
}
],
"outputs": [
{
"name": "",
"type": "MASK",
"links": [
78
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 34,
"type": "[Comfy3D] Deep Marching Tetrahedrons",
"pos": [
2332,
249
],
"size": {
"0": 443.4000244140625,
"1": 314
},
"flags": {},
"order": 19,
"mode": 0,
"inputs": [
{
"name": "reference_points_cloud",
"type": "POINTCLOUD",
"link": 53
},
{
"name": "reference_images",
"type": "IMAGE",
"link": null
},
{
"name": "reference_masks",
"type": "MASK",
"link": null
}
],
"outputs": [
{
"name": "trained_mesh",
"type": "MESH",
"links": [
79
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Deep Marching Tetrahedrons"
},
"widgets_values": [
3000,
1,
0.1,
1,
1,
128,
0.0001,
2,
128,
5
]
},
{
"id": 43,
"type": "Reroute",
"pos": [
2195,
-25
],
"size": [
75,
26
],
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 64
}
],
"outputs": [
{
"name": "",
"type": "ORBIT_CAMPOSES",
"links": [
80
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 48,
"type": "[Comfy3D] Fitting Mesh With Multiview Images",
"pos": [
2831,
90
],
"size": {
"0": 531.5999755859375,
"1": 358
},
"flags": {},
"order": 24,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 77
},
{
"name": "reference_masks",
"type": "MASK",
"link": 78
},
{
"name": "reference_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 80
},
{
"name": "mesh",
"type": "MESH",
"link": 79
}
],
"outputs": [
{
"name": "trained_mesh",
"type": "MESH",
"links": [
81
],
"shape": 3,
"slot_index": 0
},
{
"name": "baked_texture",
"type": "IMAGE",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Fitting Mesh With Multiview Images"
},
"widgets_values": [
49.1,
1024,
1024,
1024,
3,
0.001,
false,
0.0001,
0.5,
512,
false
]
},
{
"id": 36,
"type": "[Comfy3D] Save 3D Mesh",
"pos": [
3398,
89
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 27,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESH",
"link": 81
}
],
"outputs": [
{
"name": "save_path",
"type": "STRING",
"links": [
55
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Save 3D Mesh"
},
"widgets_values": [
"MeshTest\\3DMesh_%Y-%m-%d-%M-%S-%f.obj"
]
},
{
"id": 37,
"type": "[Comfy3D] Preview 3DMesh",
"pos": [
3744,
89
],
"size": [
600,
500
],
"flags": {},
"order": 28,
"mode": 0,
"inputs": [
{
"name": "mesh_file_path",
"type": "STRING",
"link": 55,
"widget": {
"name": "mesh_file_path"
}
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Preview 3DMesh"
},
"widgets_values": [
"",
null
]
}
],
"links": [
@@ -1391,14 +1354,6 @@
0,
"POINTCLOUD"
],
[
54,
38,
0,
36,
0,
"MESH"
],
[
55,
36,
@@ -1407,14 +1362,6 @@
0,
"STRING"
],
[
56,
34,
0,
38,
3,
"MESH"
],
[
58,
7,
@@ -1447,14 +1394,6 @@
0,
"*"
],
[
63,
42,
0,
38,
1,
"MASK"
],
[
64,
31,
@@ -1463,30 +1402,6 @@
0,
"*"
],
[
65,
43,
0,
38,
2,
"ORBIT_CAMPOSES"
],
[
68,
34,
0,
46,
0,
"MESH"
],
[
69,
40,
0,
38,
0,
"IMAGE"
],
[
70,
7,
@@ -1534,6 +1449,46 @@
24,
0,
"GS_PLY"
],
[
77,
40,
0,
48,
0,
"IMAGE"
],
[
78,
42,
0,
48,
1,
"MASK"
],
[
79,
34,
0,
48,
3,
"MESH"
],
[
80,
43,
0,
48,
2,
"ORBIT_CAMPOSES"
],
[
81,
48,
0,
36,
0,
"MESH"
]
],
"groups": [
_Example_Workflows/Multi-View-Images_to_Instant-NGP_to_3DMesh.json
+949
View File
@@ -0,0 +1,949 @@
{
"last_node_id": 22,
"last_link_id": 41,
"nodes": [
{
"id": 6,
"type": "MaskToImage",
"pos": [
1490,
740
],
"size": {
"0": 210,
"1": 26
},
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "mask",
"type": "MASK",
"link": 6
}
],
"outputs": [
{
"name": "IMAGE",
"type": "IMAGE",
"links": [
4
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "MaskToImage"
}
},
{
"id": 7,
"type": "PreviewImage",
"pos": [
1740,
740
],
"size": {
"0": 215.39028930664062,
"1": 246
},
"flags": {
"collapsed": false
},
"order": 12,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 4
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 13,
"type": "PreviewImage",
"pos": [
1485,
-69
],
"size": {
"0": 215.39028930664062,
"1": 246
},
"flags": {
"collapsed": false
},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "images",
"type": "IMAGE",
"link": 17
}
],
"properties": {
"Node name for S&R": "PreviewImage"
}
},
{
"id": 10,
"type": "[Comfy3D] Preview 3DMesh",
"pos": [
2400,
684
],
"size": [
600,
500
],
"flags": {},
"order": 15,
"mode": 0,
"inputs": [
{
"name": "mesh_file_path",
"type": "STRING",
"link": 14,
"widget": {
"name": "mesh_file_path"
}
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Preview 3DMesh"
},
"widgets_values": [
"",
null
]
},
{
"id": 9,
"type": "Reroute",
"pos": [
1237,
553
],
"size": [
75,
26
],
"flags": {},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 11
}
],
"outputs": [
{
"name": "",
"type": "ORBIT_CAMPOSES",
"links": [
19,
34
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 2,
"type": "[Comfy3D] Gaussian Splatting Orbit Renderer",
"pos": [
807,
289
],
"size": {
"0": 556.800048828125,
"1": 198
},
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"link": 16
},
{
"name": "render_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 5
}
],
"outputs": [
{
"name": "rendered_gs_images",
"type": "IMAGE",
"links": [
17,
21,
32
],
"shape": 3,
"slot_index": 0
},
{
"name": "rendered_gs_masks",
"type": "MASK",
"links": [
6,
23,
33
],
"shape": 3,
"slot_index": 1
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Gaussian Splatting Orbit Renderer"
},
"widgets_values": [
1024,
1024,
49.1,
1,
1,
1
]
},
{
"id": 5,
"type": "[Comfy3D] Stack Orbit Camera Poses",
"pos": [
362,
389
],
"size": {
"0": 315,
"1": 586
},
"flags": {},
"order": 0,
"mode": 0,
"outputs": [
{
"name": "orbit_camposes",
"type": "ORBIT_CAMPOSES",
"links": [
5,
9
],
"shape": 3,
"slot_index": 0
},
{
"name": "orbit_radius_list",
"type": "FLOAT",
"links": null,
"shape": 6
},
{
"name": "elevation_list",
"type": "FLOAT",
"links": null,
"shape": 6
},
{
"name": "azimuth_list",
"type": "FLOAT",
"links": null,
"shape": 6
},
{
"name": "orbit_center_X_list",
"type": "FLOAT",
"links": null,
"shape": 6
},
{
"name": "orbit_center_Y_list",
"type": "FLOAT",
"links": null,
"shape": 6
},
{
"name": "orbit_center_Z_list",
"type": "FLOAT",
"links": null,
"shape": 6
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Stack Orbit Camera Poses"
},
"widgets_values": [
1.75,
1.75,
0.1,
-45,
45,
45,
0,
-0.01,
6,
0,
0,
0.1,
0,
0,
0.1,
0,
0,
0.1
]
},
{
"id": 12,
"type": "[Comfy3D] Switch 3DGS Axis",
"pos": [
466,
219
],
"size": {
"0": 210,
"1": 106
},
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"link": 15
}
],
"outputs": [
{
"name": "switched_gs_ply",
"type": "GS_PLY",
"links": [
16
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Switch 3DGS Axis"
},
"widgets_values": [
"+x",
"-y",
"-z"
]
},
{
"id": 1,
"type": "[Comfy3D] Load 3DGS",
"pos": [
104,
218
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 1,
"mode": 0,
"outputs": [
{
"name": "gs_ply",
"type": "GS_PLY",
"links": [
15
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Load 3DGS"
},
"widgets_values": [
"C:\\Users\\reall\\Softwares\\ComfyUI_windows_portable\\ComfyUI\\output\\LGMTest\\CatCoffee.ply"
]
},
{
"id": 8,
"type": "Reroute",
"pos": [
808,
553
],
"size": [
75,
26
],
"flags": {},
"order": 2,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 9
}
],
"outputs": [
{
"name": "",
"type": "ORBIT_CAMPOSES",
"links": [
11
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 16,
"type": "Reroute",
"pos": [
1526,
189
],
"size": [
75,
26
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 21
}
],
"outputs": [
{
"name": "",
"type": "IMAGE",
"links": [
37
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 17,
"type": "Reroute",
"pos": [
1551,
226
],
"size": [
75,
26
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 23
}
],
"outputs": [
{
"name": "",
"type": "MASK",
"links": [
38
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 15,
"type": "Reroute",
"pos": [
2124,
552
],
"size": [
75,
26
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "",
"type": "*",
"link": 19
}
],
"outputs": [
{
"name": "",
"type": "ORBIT_CAMPOSES",
"links": [
40
],
"slot_index": 0
}
],
"properties": {
"showOutputText": false,
"horizontal": false
}
},
{
"id": 19,
"type": "[Comfy3D] Save 3D Mesh",
"pos": [
2945,
188
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 16,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESH",
"link": 41
}
],
"outputs": [
{
"name": "save_path",
"type": "STRING",
"links": [
25
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Save 3D Mesh"
},
"widgets_values": [
"Instant_NGP_Test/CatCoffee_1.obj"
]
},
{
"id": 11,
"type": "[Comfy3D] Save 3D Mesh",
"pos": [
2035,
684
],
"size": {
"0": 315,
"1": 58
},
"flags": {},
"order": 13,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESH",
"link": 36
}
],
"outputs": [
{
"name": "save_path",
"type": "STRING",
"links": [
14
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Save 3D Mesh"
},
"widgets_values": [
"Instant_NGP_Test/CatCoffee.obj"
]
},
{
"id": 18,
"type": "[Comfy3D] Preview 3DMesh",
"pos": [
3303,
189
],
"size": [
600,
500
],
"flags": {},
"order": 17,
"mode": 0,
"inputs": [
{
"name": "mesh_file_path",
"type": "STRING",
"link": 25,
"widget": {
"name": "mesh_file_path"
}
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Preview 3DMesh"
},
"widgets_values": [
"",
null
]
},
{
"id": 21,
"type": "[Comfy3D] Instant NGP",
"pos": [
1570,
288
],
"size": {
"0": 418.1999816894531,
"1": 266
},
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "reference_image",
"type": "IMAGE",
"link": 32
},
{
"name": "reference_mask",
"type": "MASK",
"link": 33
},
{
"name": "reference_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 34
}
],
"outputs": [
{
"name": "mesh",
"type": "MESH",
"links": [
36,
39
],
"shape": 3,
"slot_index": 0
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Instant NGP"
},
"widgets_values": [
49.1,
1024,
128,
256,
128,
10,
1024,
false
]
},
{
"id": 22,
"type": "[Comfy3D] Fitting Mesh With Multiview Images",
"pos": [
2358,
188
],
"size": {
"0": 531.5999755859375,
"1": 358
},
"flags": {},
"order": 14,
"mode": 0,
"inputs": [
{
"name": "reference_images",
"type": "IMAGE",
"link": 37
},
{
"name": "reference_masks",
"type": "MASK",
"link": 38
},
{
"name": "reference_orbit_camera_poses",
"type": "ORBIT_CAMPOSES",
"link": 40
},
{
"name": "mesh",
"type": "MESH",
"link": 39
}
],
"outputs": [
{
"name": "trained_mesh",
"type": "MESH",
"links": [
41
],
"shape": 3,
"slot_index": 0
},
{
"name": "baked_texture",
"type": "IMAGE",
"links": null,
"shape": 3
}
],
"properties": {
"Node name for S&R": "[Comfy3D] Fitting Mesh With Multiview Images"
},
"widgets_values": [
49.1,
1024,
1024,
2048,
1,
0.1,
false,
0.0001,
0.5,
512,
false
]
}
],
"links": [
[
4,
6,
0,
7,
0,
"IMAGE"
],
[
5,
5,
0,
2,
1,
"ORBIT_CAMPOSES"
],
[
6,
2,
1,
6,
0,
"MASK"
],
[
9,
5,
0,
8,
0,
"*"
],
[
11,
8,
0,
9,
0,
"*"
],
[
14,
11,
0,
10,
0,
"STRING"
],
[
15,
1,
0,
12,
0,
"GS_PLY"
],
[
16,
12,
0,
2,
0,
"GS_PLY"
],
[
17,
2,
0,
13,
0,
"IMAGE"
],
[
19,
9,
0,
15,
0,
"*"
],
[
21,
2,
0,
16,
0,
"*"
],
[
23,
2,
1,
17,
0,
"*"
],
[
25,
19,
0,
18,
0,
"STRING"
],
[
32,
2,
0,
21,
0,
"IMAGE"
],
[
33,
2,
1,
21,
1,
"MASK"
],
[
34,
9,
0,
21,
2,
"ORBIT_CAMPOSES"
],
[
36,
21,
0,
11,
0,
"MESH"
],
[
37,
16,
0,
22,
0,
"IMAGE"
],
[
38,
17,
0,
22,
1,
"MASK"
],
[
39,
21,
0,
22,
3,
"MESH"
],
[
40,
15,
0,
22,
2,
"ORBIT_CAMPOSES"
],
[
41,
22,
0,
19,
0,
"MESH"
]
],
"groups": [],
"config": {},
"extra": {},
"version": 0.4
}
_Example_Workflows/_Example_Inputs_Files/Meshes/CatCoffee.ply
BIN
View File
Binary file not shown.
algorithms/Instant_NGP.py
+206
View File
@@ -0,0 +1,206 @@
import tqdm
import random
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import nerfacc
import comfy.utils
from pytorch_msssim import SSIM, MS_SSIM
from kiui.op import safe_normalize
from kiui.cam import orbit_camera
from kiui.nn import MLP, trunc_exp
from ..shared_utils.image_utils import prepare_torch_img
class InstantNGP(nn.Module):
def __init__(self, resolution=128, device="cuda"):
super().__init__()
from kiui.gridencoder import GridEncoder
self.device = torch.device(device)
self.ref_size_H = resolution
self.ref_size_W = resolution
self.render_step_size = 5e-3
self.aabb = torch.tensor([-1.0, -1.0, -1.0, 1.0, 1.0, 1.0], device=self.device)
self.estimator = nerfacc.OccGridEstimator(roi_aabb=self.aabb, resolution=64, levels=1)
self.encoder_density = GridEncoder(num_levels=12) # VMEncoder(output_dim=16, mode='sum')
self.encoder = GridEncoder(num_levels=12)
self.mlp_density = MLP(self.encoder_density.output_dim, 1, 32, 2, bias=False)
self.mlp = MLP(self.encoder.output_dim, 3, 32, 2, bias=False)
def get_rays(self, pose, h, w, fovy, opengl=True):
x, y = torch.meshgrid(
torch.arange(w, device=pose.device),
torch.arange(h, device=pose.device),
indexing="xy",
)
x = x.flatten()
y = y.flatten()
cx = w * 0.5
cy = h * 0.5
focal = h * 0.5 / np.tan(0.5 * np.deg2rad(fovy))
camera_dirs = F.pad(
torch.stack(
[
(x - cx + 0.5) / focal,
(y - cy + 0.5) / focal * (-1.0 if opengl else 1.0),
],
dim=-1,
),
(0, 1),
value=(-1.0 if opengl else 1.0),
) # [hw, 3]
rays_d = camera_dirs @ pose[:3, :3].transpose(0, 1) # [hw, 3]
rays_o = pose[:3, 3].unsqueeze(0).expand_as(rays_d) # [hw, 3]
rays_o = rays_o.view(h, w, 3)
rays_d = safe_normalize(rays_d).view(h, w, 3)
return rays_o, rays_d
def get_color(self, xs):
return torch.sigmoid(self.mlp(self.encoder(xs.to(self.device))))
def get_density(self, xs):
# xs: [..., 3]
xs = xs.to(self.device)
prefix = xs.shape[:-1]
xs = xs.view(-1, 3)
feats = self.encoder_density(xs)
density = trunc_exp(self.mlp_density(feats))
density = density.view(*prefix, 1)
return density
def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
self.ref_imgs_num = len(reference_images)
self.all_ref_cam_poses = reference_orbit_camera_poses
self.ref_cam_fovy = reference_orbit_camera_fovy
# prepare reference images and masks
ref_imgs_torch_list = []
ref_masks_torch_list = []
for i in range(self.ref_imgs_num):
ref_imgs_torch_list.append(prepare_torch_img(reference_images[i].unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2).unsqueeze(0), self.ref_size_H, self.ref_size_W, self.device))
self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0) # [N, 3, H, W]
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0).squeeze(1) # [N, H, W]
def render_nerf(self, pose):
pose = torch.from_numpy(pose.astype(np.float32)).to(self.device)
# get rays
rays_o, rays_d = self.get_rays(pose, self.ref_size_H, self.ref_size_W, self.ref_cam_fovy)
hw = rays_o.shape[0] * rays_o.shape[1]
rays_o = rays_o.view(hw, 3)
rays_d = rays_d.view(hw, 3)
# update occ grid
if self.training:
def occ_eval_fn(xs):
sigmas = self.get_density(xs)
return self.render_step_size * sigmas
self.estimator.update_every_n_steps(self.render_step, occ_eval_fn=occ_eval_fn, occ_thre=0.01, n=8)
self.render_step += 1
# render
def sigma_fn(t_starts, t_ends, ray_indices):
t_origins = rays_o[ray_indices]
t_dirs = rays_d[ray_indices]
xs = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
sigmas = self.get_density(xs)
return sigmas.squeeze(-1)
with torch.no_grad():
ray_indices, t_starts, t_ends = self.estimator.sampling(
rays_o,
rays_d,
sigma_fn=sigma_fn,
near_plane=0.01,
far_plane=100,
render_step_size=self.render_step_size,
stratified=self.training,
cone_angle=0,
)
t_origins = rays_o[ray_indices]
t_dirs = rays_d[ray_indices]
xs = t_origins + t_dirs * (t_starts + t_ends)[:, None] / 2.0
sigmas = self.get_density(xs).squeeze(-1)
rgbs = torch.sigmoid(self.mlp(self.encoder(xs)))
n_rays=rays_o.shape[0]
weights, t, alphas = nerfacc.render_weight_from_density(t_starts, t_ends, sigmas, ray_indices=ray_indices, n_rays=n_rays)
color = nerfacc.accumulate_along_rays(weights, values=rgbs, ray_indices=ray_indices, n_rays=n_rays)
alpha = nerfacc.accumulate_along_rays(weights, values=None, ray_indices=ray_indices, n_rays=n_rays)
color = color + 1 * (1.0 - alpha)
color = color.view(self.ref_size_H, self.ref_size_W, 3).clamp(0, 1).permute(2, 0, 1).contiguous()
alpha = alpha.view(self.ref_size_H, self.ref_size_W).clamp(0, 1).contiguous()
return color, alpha
def fit_nerf(self, iters=512):
optimizer = torch.optim.Adam([
{'params': self.encoder_density.parameters(), 'lr': 1e-2},
{'params': self.encoder.parameters(), 'lr': 1e-2},
{'params': self.mlp_density.parameters(), 'lr': 1e-3},
{'params': self.mlp.parameters(), 'lr': 1e-3},
])
print(f"[INFO] fitting nerf...")
self.render_step = 0
ref_imgs_num_minus_1 = self.ref_imgs_num-1
comfy_pbar = comfy.utils.ProgressBar(iters)
pbar = tqdm.trange(iters)
for step in pbar:
i = random.randint(0, ref_imgs_num_minus_1)
radius, elevation, azimuth, center_X, center_Y, center_Z = self.all_ref_cam_poses[i]
orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
image_gt = self.ref_imgs_torch[i] # [3, H, W]
alpha_gt = self.ref_masks_torch[i] # [H, W]
image_pred, alpha_pred = self.render_nerf(pose)
# if i % 200 == 0:
# kiui.vis.plot_image(image_gt, alpha_gt, image_pred, alpha_pred)
loss_mse = F.mse_loss(image_pred, image_gt) + 0.1 * F.mse_loss(alpha_pred, alpha_gt)
loss = loss_mse #+ 0.1 * self.encoder_density.tv_loss() #+ 0.0001 * self.encoder_density.density_loss()
#loss += self.lambda_ssim * (1 - self.ms_ssim_loss(image_gt, image_pred))
loss.backward()
self.encoder_density.grad_total_variation(1e-8)
optimizer.step()
optimizer.zero_grad()
pbar.set_description(f"NeRF Fitting Loss = {loss_mse.item():.6f}")
comfy_pbar.update_absolute(step + 1)
torch.cuda.synchronize()
print(f"[INFO] finished fitting nerf!")
algorithms/diff_texturing.py → algorithms/diff_mesh.py
+27 -3
View File
@@ -2,9 +2,12 @@ import random
import tqdm
import torch
import torch.nn as nn
import torch.nn.functional as F
import numpy as np
from kiui.mesh_utils import clean_mesh, decimate_mesh
from kiui.cam import orbit_camera
from kiui.mesh_utils import laplacian_smooth_loss, normal_consistency
from pytorch_msssim import SSIM, MS_SSIM
import comfy.utils
@@ -13,11 +16,14 @@ from .diff_mesh_renderer import DiffRastRenderer
from ..shared_utils.camera_utils import OrbitCamera
from ..shared_utils.image_utils import prepare_torch_img
class DiffTextureBaker:
class DiffMesh:
def __init__(self, mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, force_cuda_rasterize):
def __init__(self, mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, remesh_after_n_iteration, force_cuda_rasterize):
self.device = torch.device("cuda")
self.train_mesh_geometry = train_mesh_geometry
self.remesh_after_n_iteration = remesh_after_n_iteration
# prepare main components for optimization
self.renderer = DiffRastRenderer(mesh, force_cuda_rasterize).to(self.device)
@@ -51,7 +57,7 @@ class DiffTextureBaker:
self.ref_imgs_torch = torch.cat(ref_imgs_torch_list, dim=0)
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
def training(self):
def training(self, decimate_target=5e4):
starter = torch.cuda.Event(enable_timing=True)
ender = torch.cuda.Event(enable_timing=True)
starter.record()
@@ -100,6 +106,24 @@ class DiffTextureBaker:
# [1, 3, H, W] in [0, 1]
#loss += self.lambda_ssim * (1 - self.ssim_loss(X, Y))
loss += self.lambda_ssim * (1 - self.ms_ssim_loss(masked_ref_img_batch_torch, masked_rendered_img_batch_torch))
# Regularization loss
if self.train_mesh_geometry:
current_v = self.renderer.mesh.v + self.renderer.v_offsets
loss += 0.01 * laplacian_smooth_loss(current_v, self.renderer.mesh.f)
loss += 0.001 * normal_consistency(current_v, self.renderer.mesh.f)
loss += 0.1 * (self.renderer.v_offsets ** 2).sum(-1).mean()
# remesh periodically
if step > 0 and step % self.remesh_after_n_iteration == 0:
vertices = (self.renderer.mesh.v + self.renderer.v_offsets).detach().cpu().numpy()
triangles = self.renderer.mesh.f.detach().cpu().numpy()
vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.01)
if triangles.shape[0] > decimate_target:
vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False)
self.renderer.mesh.v = torch.from_numpy(vertices).contiguous().float().to(self.device)
self.renderer.mesh.f = torch.from_numpy(triangles).contiguous().int().to(self.device)
self.renderer.v_offsets = nn.Parameter(torch.zeros_like(self.renderer.mesh.v)).to(self.device)
# optimize step
loss.backward()
algorithms/diff_mesh_renderer.py
+3 -5
View File
@@ -7,6 +7,8 @@ import torch.nn as nn
import torch.nn.functional as F
import nvdiffrast.torch as dr
from kiui.op import inverse_sigmoid
from ..mesh_processer.mesh import safe_normalize
def scale_img_nhwc(x, size, mag='bilinear', min='bilinear'):
@@ -30,10 +32,6 @@ def scale_img_nhw(x, size, mag='bilinear', min='bilinear'):
def scale_img_hw(x, size, mag='bilinear', min='bilinear'):
return scale_img_nhwc(x[None, ..., None], size, mag, min)[0, ..., 0]
def trunc_rev_sigmoid(x, eps=1e-6):
x = x.clamp(eps, 1 - eps)
return torch.log(x / (1 - x))
def make_divisible(x, m=8):
return int(math.ceil(x / m) * m)
@@ -51,7 +49,7 @@ class DiffRastRenderer(nn.Module):
# extract trainable parameters
self.v_offsets = nn.Parameter(torch.zeros_like(self.mesh.v), requires_grad=True)
self.raw_albedo = nn.Parameter(trunc_rev_sigmoid(self.mesh.albedo), requires_grad=True)
self.raw_albedo = nn.Parameter(inverse_sigmoid(self.mesh.albedo), requires_grad=True)
def get_params(self, texture_lr, train_geo, geom_lr):
algorithms/gaussian_model.txt
+397
View File
@@ -0,0 +1,397 @@
import torch
import numpy as np
from utils.general_utils import inverse_sigmoid, get_expon_lr_func, build_rotation
from torch import nn
import os
from utils.system_utils import mkdir_p
from plyfile import PlyData, PlyElement
from utils.sh_utils import RGB2SH
from simple_knn._C import distCUDA2
from utils.graphics_utils import BasicPointCloud
from utils.general_utils import strip_symmetric, build_scaling_rotation
class GaussianModel:
def setup_functions(self):
def build_covariance_from_scaling_rotation(scaling, scaling_modifier, rotation):
L = build_scaling_rotation(scaling_modifier * scaling, rotation)
actual_covariance = L @ L.transpose(1, 2)
symm = strip_symmetric(actual_covariance)
return symm
self.scaling_activation = torch.exp
self.scaling_inverse_activation = torch.log
self.covariance_activation = build_covariance_from_scaling_rotation
self.opacity_activation = torch.sigmoid
self.inverse_opacity_activation = inverse_sigmoid
self.rotation_activation = torch.nn.functional.normalize
def __init__(self, sh_degree : int, divide_ratio : float):
self.active_sh_degree = 0
self.max_sh_degree = sh_degree
self._xyz = torch.empty(0)
self._features_dc = torch.empty(0)
self._features_rest = torch.empty(0)
self._scaling = torch.empty(0)
self._rotation = torch.empty(0)
self._opacity = torch.empty(0)
self.max_radii2D = torch.empty(0)
self.xyz_gradient_accum = torch.empty(0)
self.denom = torch.empty(0)
self.optimizer = None
self.percent_dense = 0
self.spatial_lr_scale = 0
self.setup_functions()
self.divide_ratio = divide_ratio
def capture(self):
return (
self.active_sh_degree,
self._xyz,
self._features_dc,
self._features_rest,
self._scaling,
self._rotation,
self._opacity,
self.max_radii2D,
self.xyz_gradient_accum,
self.denom,
self.optimizer.state_dict(),
self.spatial_lr_scale,
)
def restore(self, model_args, training_args):
(self.active_sh_degree,
self._xyz,
self._features_dc,
self._features_rest,
self._scaling,
self._rotation,
self._opacity,
self.max_radii2D,
xyz_gradient_accum,
denom,
opt_dict,
self.spatial_lr_scale) = model_args
self.training_setup(training_args)
self.xyz_gradient_accum = xyz_gradient_accum
self.denom = denom
self.optimizer.load_state_dict(opt_dict)
@property
def get_scaling(self):
return self.scaling_activation(self._scaling)
@property
def get_rotation(self):
return self.rotation_activation(self._rotation)
@property
def get_xyz(self):
return self._xyz
@property
def get_features(self):
features_dc = self._features_dc
features_rest = self._features_rest
return torch.cat((features_dc, features_rest), dim=1)
@property
def get_opacity(self):
return self.opacity_activation(self._opacity)
def get_covariance(self, scaling_modifier = 1):
return self.covariance_activation(self.get_scaling, scaling_modifier, self._rotation)
def oneupSHdegree(self):
if self.active_sh_degree < self.max_sh_degree:
self.active_sh_degree += 1
def create_from_pcd(self, pcd : BasicPointCloud, spatial_lr_scale : float):
self.spatial_lr_scale = spatial_lr_scale
fused_point_cloud = torch.tensor(np.asarray(pcd.points)).float().cuda()
fused_color = RGB2SH(torch.tensor(np.asarray(pcd.colors)).float().cuda())
features = torch.zeros((fused_color.shape[0], 3, (self.max_sh_degree + 1) ** 2)).float().cuda()
features[:, :3, 0 ] = fused_color
features[:, 3:, 1:] = 0.0
print("Number of points at initialisation : ", fused_point_cloud.shape[0])
dist2 = torch.clamp_min(distCUDA2(torch.from_numpy(np.asarray(pcd.points)).float().cuda()), 0.0000001)
scales = torch.log(torch.sqrt(dist2))[...,None].repeat(1, 3)
rots = torch.zeros((fused_point_cloud.shape[0], 4), device="cuda")
rots[:, 0] = 1
opacities = inverse_sigmoid(0.1 * torch.ones((fused_point_cloud.shape[0], 1), dtype=torch.float, device="cuda"))
self._xyz = nn.Parameter(fused_point_cloud.requires_grad_(True))
self._features_dc = nn.Parameter(features[:,:,0:1].transpose(1, 2).contiguous().requires_grad_(True))
self._features_rest = nn.Parameter(features[:,:,1:].transpose(1, 2).contiguous().requires_grad_(True))
self._scaling = nn.Parameter(scales.requires_grad_(True))
self._rotation = nn.Parameter(rots.requires_grad_(True))
self._opacity = nn.Parameter(opacities.requires_grad_(True))
self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
def training_setup(self, training_args):
self.percent_dense = training_args.percent_dense
self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
l = [
{'params': [self._xyz], 'lr': training_args.position_lr_init * self.spatial_lr_scale, "name": "xyz"},
{'params': [self._features_dc], 'lr': training_args.feature_lr, "name": "f_dc"},
{'params': [self._features_rest], 'lr': training_args.feature_lr / 20.0, "name": "f_rest"},
{'params': [self._opacity], 'lr': training_args.opacity_lr, "name": "opacity"},
{'params': [self._scaling], 'lr': training_args.scaling_lr, "name": "scaling"},
{'params': [self._rotation], 'lr': training_args.rotation_lr, "name": "rotation"}
]
self.optimizer = torch.optim.Adam(l, lr=0.0, eps=1e-15)
self.xyz_scheduler_args = get_expon_lr_func(lr_init=training_args.position_lr_init*self.spatial_lr_scale,
lr_final=training_args.position_lr_final*self.spatial_lr_scale,
lr_delay_mult=training_args.position_lr_delay_mult,
max_steps=training_args.position_lr_max_steps)
def update_learning_rate(self, iteration):
''' Learning rate scheduling per step '''
for param_group in self.optimizer.param_groups:
if param_group["name"] == "xyz":
lr = self.xyz_scheduler_args(iteration)
param_group['lr'] = lr
return lr
def construct_list_of_attributes(self):
l = ['x', 'y', 'z', 'nx', 'ny', 'nz']
for i in range(self._features_dc.shape[1]*self._features_dc.shape[2]):
l.append('f_dc_{}'.format(i))
for i in range(self._features_rest.shape[1]*self._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 save_ply(self, path):
mkdir_p(os.path.dirname(path))
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()
f_rest = self._features_rest.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
opacities = self._opacity.detach().cpu().numpy()
scale = 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')
PlyData([el]).write(path)
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]
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_")]
extra_f_names = sorted(extra_f_names, key = lambda x: int(x.split('_')[-1]))
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])
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_")]
scale_names = sorted(scale_names, key = lambda x: int(x.split('_')[-1]))
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")]
rot_names = sorted(rot_names, key = lambda x: int(x.split('_')[-1]))
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])
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))
self._features_rest = nn.Parameter(torch.tensor(features_extra, dtype=torch.float, device="cuda").transpose(1, 2).contiguous().requires_grad_(True))
self._opacity = nn.Parameter(torch.tensor(opacities, dtype=torch.float, device="cuda").requires_grad_(True))
self._scaling = nn.Parameter(torch.tensor(scales, dtype=torch.float, device="cuda").requires_grad_(True))
self._rotation = nn.Parameter(torch.tensor(rots, dtype=torch.float, device="cuda").requires_grad_(True))
self.active_sh_degree = self.max_sh_degree
def replace_tensor_to_optimizer(self, tensor, name):
optimizable_tensors = {}
for group in self.optimizer.param_groups:
if group["name"] == name:
stored_state = self.optimizer.state.get(group['params'][0], None)
stored_state["exp_avg"] = torch.zeros_like(tensor)
stored_state["exp_avg_sq"] = torch.zeros_like(tensor)
del self.optimizer.state[group['params'][0]]
group["params"][0] = nn.Parameter(tensor.requires_grad_(True))
self.optimizer.state[group['params'][0]] = stored_state
optimizable_tensors[group["name"]] = group["params"][0]
return optimizable_tensors
def _prune_optimizer(self, mask):
optimizable_tensors = {}
for group in self.optimizer.param_groups:
stored_state = self.optimizer.state.get(group['params'][0], None)
if stored_state is not None:
stored_state["exp_avg"] = stored_state["exp_avg"][mask]
stored_state["exp_avg_sq"] = stored_state["exp_avg_sq"][mask]
del self.optimizer.state[group['params'][0]]
group["params"][0] = nn.Parameter((group["params"][0][mask].requires_grad_(True)))
self.optimizer.state[group['params'][0]] = stored_state
optimizable_tensors[group["name"]] = group["params"][0]
else:
group["params"][0] = nn.Parameter(group["params"][0][mask].requires_grad_(True))
optimizable_tensors[group["name"]] = group["params"][0]
return optimizable_tensors
def prune_points(self, mask):
valid_points_mask = ~mask
optimizable_tensors = self._prune_optimizer(valid_points_mask)
self._xyz = optimizable_tensors["xyz"]
self._features_dc = optimizable_tensors["f_dc"]
self._features_rest = optimizable_tensors["f_rest"]
self._opacity = optimizable_tensors["opacity"]
self._scaling = optimizable_tensors["scaling"]
self._rotation = optimizable_tensors["rotation"]
self.xyz_gradient_accum = self.xyz_gradient_accum[valid_points_mask]
self.denom = self.denom[valid_points_mask]
self.max_radii2D = self.max_radii2D[valid_points_mask]
def cat_tensors_to_optimizer(self, tensors_dict):
optimizable_tensors = {}
for group in self.optimizer.param_groups:
assert len(group["params"]) == 1
extension_tensor = tensors_dict[group["name"]]
stored_state = self.optimizer.state.get(group['params'][0], None)
if stored_state is not None:
stored_state["exp_avg"] = torch.cat((stored_state["exp_avg"], torch.zeros_like(extension_tensor)), dim=0)
stored_state["exp_avg_sq"] = torch.cat((stored_state["exp_avg_sq"], torch.zeros_like(extension_tensor)), dim=0)
del self.optimizer.state[group['params'][0]]
group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
self.optimizer.state[group['params'][0]] = stored_state
optimizable_tensors[group["name"]] = group["params"][0]
else:
group["params"][0] = nn.Parameter(torch.cat((group["params"][0], extension_tensor), dim=0).requires_grad_(True))
optimizable_tensors[group["name"]] = group["params"][0]
return optimizable_tensors
def densification_postfix(self, new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation):
d = {"xyz": new_xyz,
"f_dc": new_features_dc,
"f_rest": new_features_rest,
"opacity": new_opacities,
"scaling" : new_scaling,
"rotation" : new_rotation}
optimizable_tensors = self.cat_tensors_to_optimizer(d)
self._xyz = optimizable_tensors["xyz"]
self._features_dc = optimizable_tensors["f_dc"]
self._features_rest = optimizable_tensors["f_rest"]
self._opacity = optimizable_tensors["opacity"]
self._scaling = optimizable_tensors["scaling"]
self._rotation = optimizable_tensors["rotation"]
self.xyz_gradient_accum = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
self.denom = torch.zeros((self.get_xyz.shape[0], 1), device="cuda")
self.max_radii2D = torch.zeros((self.get_xyz.shape[0]), device="cuda")
def densify_and_split(self, grads, grad_threshold, scene_extent, N=2):
n_init_points = self.get_xyz.shape[0]
padded_grad = torch.zeros((n_init_points), device="cuda")
padded_grad[:grads.shape[0]] = grads.squeeze()
selected_pts_mask = torch.where(padded_grad >= grad_threshold, True, False)
selected_pts_mask = torch.logical_and(selected_pts_mask,
torch.max(self.get_scaling, dim=1).values > self.percent_dense*scene_extent)
stds = self.get_scaling[selected_pts_mask].repeat(N,1)
means =torch.zeros((stds.size(0), 3),device="cuda")
samples = torch.normal(mean=means, std=stds)
rots = build_rotation(self._rotation[selected_pts_mask]).repeat(N,1,1)
new_xyz = torch.bmm(rots, samples.unsqueeze(-1)).squeeze(-1) + self.get_xyz[selected_pts_mask].repeat(N, 1)
new_scaling = self.scaling_inverse_activation(self.get_scaling[selected_pts_mask].repeat(N,1) / (self.divide_ratio*N))
new_rotation = self._rotation[selected_pts_mask].repeat(N,1)
new_features_dc = self._features_dc[selected_pts_mask].repeat(N,1,1)
new_features_rest = self._features_rest[selected_pts_mask].repeat(N,1,1)
new_opacity = self._opacity[selected_pts_mask].repeat(N,1)
self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacity, new_scaling, new_rotation)
prune_filter = torch.cat((selected_pts_mask, torch.zeros(N * selected_pts_mask.sum(), device="cuda", dtype=bool)))
self.prune_points(prune_filter)
def densify_and_clone(self, grads, grad_threshold, scene_extent):
selected_pts_mask = torch.where(torch.norm(grads, dim=-1) >= grad_threshold, True, False)
selected_pts_mask = torch.logical_and(selected_pts_mask,
torch.max(self.get_scaling, dim=1).values <= self.percent_dense*scene_extent)
new_xyz = self._xyz[selected_pts_mask]
new_features_dc = self._features_dc[selected_pts_mask]
new_features_rest = self._features_rest[selected_pts_mask]
new_opacities = self._opacity[selected_pts_mask]
new_scaling = self._scaling[selected_pts_mask]
new_rotation = self._rotation[selected_pts_mask]
self.densification_postfix(new_xyz, new_features_dc, new_features_rest, new_opacities, new_scaling, new_rotation)
def densify_and_prune(self, max_grad, min_opacity, extent, max_screen_size):
grads = self.xyz_gradient_accum / self.denom
grads[grads.isnan()] = 0.0
self.densify_and_clone(grads, max_grad, extent)
self.densify_and_split(grads, max_grad, extent)
prune_mask = (self.get_opacity < min_opacity).squeeze()
if max_screen_size:
big_points_vs = self.max_radii2D > max_screen_size
big_points_ws = self.get_scaling.max(dim=1).values > 0.1 * extent
prune_mask = torch.logical_or(torch.logical_or(prune_mask, big_points_vs), big_points_ws)
self.prune_points(prune_mask)
torch.cuda.empty_cache()
def add_densification_stats(self, viewspace_point_tensor, update_filter):
self.xyz_gradient_accum[update_filter] += torch.norm(viewspace_point_tensor.grad[update_filter,:2], dim=-1, keepdim=True)
self.denom[update_filter] += 1
algorithms/main_3DGS_renderer.py
+2 -3
View File
@@ -27,13 +27,12 @@ from diff_gaussian_rasterization import (
)
from simple_knn._C import distCUDA2
from kiui.op import inverse_sigmoid
from ..shared_utils.sh_utils import eval_sh, SH2RGB, RGB2SH
from ..mesh_processer.mesh import Mesh, PointCloud
from ..mesh_processer.mesh_utils import construct_list_of_gs_attributes, write_gs_ply, read_gs_ply
def inverse_sigmoid(x):
return torch.log(x/(1-x))
def get_expon_lr_func(
lr_init, lr_final, lr_delay_steps=0, lr_delay_mult=1.0, max_steps=1000000
):
algorithms/nerf_marching_cubes_converter.py
+10
View File
@@ -162,6 +162,7 @@ class GSConverterNeRFMarchingCubes(nn.Module):
])
print(f"[INFO] fitting nerf...")
imgs, alphas = [], []
pbar = tqdm.trange(iters)
for i in pbar:
@@ -172,6 +173,8 @@ class GSConverterNeRFMarchingCubes(nn.Module):
pose = orbit_camera(ver, hor, rad)
image_gt, alpha_gt = self.render_gs(pose)
imgs.append(image_gt.permute(1, 2, 0).unsqueeze(0))
alphas.append(alpha_gt.unsqueeze(0))
image_pred, alpha_pred = self.render_nerf(pose)
# if i % 200 == 0:
@@ -189,6 +192,8 @@ class GSConverterNeRFMarchingCubes(nn.Module):
pbar.set_description(f"MSE = {loss_mse.item():.6f}")
print(f"[INFO] finished fitting nerf!")
return torch.cat(imgs, dim=0), torch.cat(alphas, dim=0)
def render_mesh(self, pose):
@@ -265,6 +270,11 @@ class GSConverterNeRFMarchingCubes(nn.Module):
self.v = torch.from_numpy(vertices).contiguous().float().to(self.device)
self.f = torch.from_numpy(triangles).contiguous().int().to(self.device)
self.deform = nn.Parameter(torch.zeros_like(self.v)).to(self.device)
# Here coarse mesh already resemble the target shape closely
#mcube_mesh = Mesh(v=self.v, f=self.f, albedo=None, device=self.device)
#mcube_mesh.auto_normal()
#mcube_mesh.auto_uv()
# fit mesh from gs
lr_factor = 1
mesh_processer/mesh_utils.py
+79 -1
View File
@@ -459,4 +459,82 @@ def switch_mesh_axis_and_scale(mesh, target_axis, target_scale, flip_normal=Fals
if flip_normal:
mesh.vn *= -1
return mesh
def marching_cubes_density_to_mesh(get_density_func, grid_size=256, S=128, density_thresh=10, decimate_target=5e4):
import mcubes
from kiui.mesh_utils import clean_mesh, decimate_mesh
sigmas = np.zeros([grid_size, grid_size, grid_size], dtype=np.float32)
X = torch.linspace(-1, 1, grid_size).split(S)
Y = torch.linspace(-1, 1, grid_size).split(S)
Z = torch.linspace(-1, 1, grid_size).split(S)
for xi, xs in enumerate(X):
for yi, ys in enumerate(Y):
for zi, zs in enumerate(Z):
xx, yy, zz = torch.meshgrid(xs, ys, zs, indexing='ij')
pts = torch.cat([xx.reshape(-1, 1), yy.reshape(-1, 1), zz.reshape(-1, 1)], dim=-1) # [S, 3]
val = get_density_func(pts)
sigmas[xi * S: xi * S + len(xs), yi * S: yi * S + len(ys), zi * S: zi * S + len(zs)] = val.reshape(len(xs), len(ys), len(zs)).detach().cpu().numpy() # [S, 1] --> [x, y, z]
print(f'[INFO] marching cubes thresh: {density_thresh} ({sigmas.min()} ~ {sigmas.max()})')
vertices, triangles = mcubes.marching_cubes(sigmas, density_thresh)
vertices = vertices / (grid_size - 1.0) * 2 - 1
# clean
vertices = vertices.astype(np.float32)
triangles = triangles.astype(np.int32)
vertices, triangles = clean_mesh(vertices, triangles, remesh=True, remesh_size=0.01)
if triangles.shape[0] > decimate_target:
vertices, triangles = decimate_mesh(vertices, triangles, decimate_target, optimalplacement=False)
return vertices, triangles
def color_func_to_albedo(mesh, get_rgb_func, texture_resolution=1024, padding=2, batch_size=640000, device="cuda", force_cuda_rast=False):
import nvdiffrast.torch as dr
from kiui.op import uv_padding
if force_cuda_rast:
glctx = dr.RasterizeCudaContext()
else:
glctx = dr.RasterizeGLContext()
# render uv maps
h = w = texture_resolution
uv = mesh.vt * 2.0 - 1.0 # uvs to range [-1, 1]
uv = torch.cat((uv, torch.zeros_like(uv[..., :1]), torch.ones_like(uv[..., :1])), dim=-1) # [N, 4]
rast, _ = dr.rasterize(glctx, uv.unsqueeze(0), mesh.ft, (h, w)) # [1, h, w, 4]
xyzs, _ = dr.interpolate(mesh.v.unsqueeze(0), rast, mesh.f) # [1, h, w, 3]
mask, _ = dr.interpolate(torch.ones_like(mesh.v[:, :1]).unsqueeze(0), rast, mesh.f) # [1, h, w, 1]
# masked query
xyzs = xyzs.view(-1, 3)
mask = (mask > 0).view(-1)
albedo = torch.zeros(h * w, 3, device=device, dtype=torch.float32)
if mask.any():
print(f"[INFO] querying texture...")
xyzs = xyzs[mask] # [M, 3]
# batched inference to avoid OOM
batch = []
head = 0
while head < xyzs.shape[0]:
tail = min(head + batch_size, xyzs.shape[0])
batch.append(get_rgb_func(xyzs[head:tail]).float())
head += batch_size
albedo[mask] = torch.cat(batch, dim=0)
albedo = albedo.view(h, w, -1)
mask = mask.view(h, w)
albedo = uv_padding(albedo, mask, padding)
return albedo
nodes.py
+87 -14
View File
@@ -26,16 +26,20 @@ from .mesh_processer.mesh_utils import (
switch_ply_axis_and_scale,
switch_mesh_axis_and_scale,
calculate_max_sh_degree_from_gs_ply,
marching_cubes_density_to_mesh,
color_func_to_albedo,
)
from .algorithms.main_3DGS import GaussianSplatting, GaussianSplattingCameraController, GSParams
from .algorithms.main_3DGS_renderer import GaussianSplattingRenderer
from .algorithms.diff_texturing import DiffTextureBaker
from .algorithms.diff_mesh import DiffMesh
from .algorithms.dmtet import DMTetMesh
from .algorithms.triplane_gaussian_transformers import TGS
from .algorithms.large_multiview_gaussian_model import LGM
from .algorithms.nerf_marching_cubes_converter import GSConverterNeRFMarchingCubes
from .algorithms.NeuS_runner import NeuSParams, NeuSRunner
from .algorithms.convolutional_reconstruction_model import CRMSampler
from .algorithms.Instant_NGP import InstantNGP
from .tgs.utils.config import ExperimentConfig, load_config as load_config_tgs
from .tgs.data import CustomImageOrbitDataset
from .tgs.utils.misc import todevice, get_device
@@ -876,7 +880,7 @@ class Gaussian_Splatting:
return (gs_ply, )
class Bake_Texture_To_Mesh:
class Fitting_Mesh_With_Multiview_Images:
def __init__(self):
self.need_update = False
@@ -892,12 +896,13 @@ class Bake_Texture_To_Mesh:
"mesh": ("MESH",),
"mesh_albedo_width": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"mesh_albedo_height": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"training_iterations": ("INT", {"default": 1000, "min": 1, "max": 100000}),
"training_iterations": ("INT", {"default": 1024, "min": 1, "max": 100000}),
"batch_size": ("INT", {"default": 3, "min": 1, "max": 0xffffffffffffffff}),
"texture_learning_rate": ("FLOAT", {"default": 0.1, "min": 0.00001, "step": 0.00001}),
"texture_learning_rate": ("FLOAT", {"default": 0.001, "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}),
"geometry_learning_rate": ("FLOAT", {"default": 0.0001, "min": 0.00001, "step": 0.00001}),
"ms_ssim_loss_weight": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}),
"remesh_after_n_iteration": ("INT", {"default": 512, "min": 128, "max": 100000}),
"force_cuda_rasterize": ("BOOLEAN", {"default": False},),
},
}
@@ -910,10 +915,10 @@ class Bake_Texture_To_Mesh:
"trained_mesh",
"baked_texture", # [1, H, W, 3]
)
FUNCTION = "bake_texture"
FUNCTION = "fitting_mesh"
CATEGORY = "Comfy3D/Algorithm"
def bake_texture(
def fitting_mesh(
self,
reference_images,
reference_masks,
@@ -928,6 +933,7 @@ class Bake_Texture_To_Mesh:
train_mesh_geometry,
geometry_learning_rate,
ms_ssim_loss_weight,
remesh_after_n_iteration,
force_cuda_rasterize
):
@@ -949,12 +955,12 @@ class Bake_Texture_To_Mesh:
with torch.inference_mode(False):
texture_baker = DiffTextureBaker(mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, force_cuda_rasterize)
mesh_fitter = DiffMesh(mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, remesh_after_n_iteration, force_cuda_rasterize)
texture_baker.prepare_training(reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy)
texture_baker.training()
mesh_fitter.prepare_training(reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy)
mesh_fitter.training()
trained_mesh, baked_texture = texture_baker.get_mesh_and_texture()
trained_mesh, baked_texture = mesh_fitter.get_mesh_and_texture()
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()
@@ -1444,9 +1450,13 @@ class Convert_3DGS_to_Mesh_with_NeRF_and_Marching_Cubes:
RETURN_TYPES = (
"MESH",
"IMAGE",
"MASK",
)
RETURN_NAMES = (
"mesh",
"imgs",
"alphas",
)
FUNCTION = "convert_gs_ply"
CATEGORY = "Comfy3D/Algorithm"
@@ -1456,11 +1466,11 @@ class Convert_3DGS_to_Mesh_with_NeRF_and_Marching_Cubes:
chosen_config = config_defaults[gs_config]
chosen_config.force_cuda_rast = force_cuda_rast
converter = GSConverterNeRFMarchingCubes(config_defaults[gs_config], gs_ply).cuda()
converter.fit_nerf()
imgs, alphas = converter.fit_nerf()
converter.fit_mesh()
converter.fit_mesh_uv()
return(converter.get_mesh(), )
return(converter.get_mesh(), imgs, alphas)
class NeuS:
NeuS_config_path = "configs/NeuS.conf"
@@ -1933,4 +1943,67 @@ class Convolutional_Reconstruction_Model:
mesh = CRMSampler.generate3d(crm_model, np_imgs, np_xyzs, device)
return (mesh,)
return (mesh,)
class Instant_NGP:
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
"reference_image": ("IMAGE",),
"reference_mask": ("MASK",),
"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": 512, "min": 1, "max": 0xffffffffffffffff}),
"training_resolution": ("INT", {"default": 128, "min": 128, "max": 8192}),
"marching_cude_grids_resolution": ("INT", {"default": 256, "min": 1, "max": 0xffffffffffffffff}),
"marching_cude_grids_batch_size": ("INT", {"default": 128, "min": 1, "max": 0xffffffffffffffff}),
"marching_cude_threshold": ("FLOAT", {"default": 10.0, "min": 0.0, "step": 0.01}),
"texture_resolution": ("INT", {"default": 1024, "min": 128, "max": 8192}),
"force_cuda_rast": ("BOOLEAN", {"default": False}),
},
}
RETURN_TYPES = (
"MESH",
)
RETURN_NAMES = (
"mesh",
)
FUNCTION = "run_instant_ngp"
CATEGORY = "Comfy3D/Algorithm"
def run_instant_ngp(
self,
reference_image,
reference_mask,
reference_orbit_camera_poses,
reference_orbit_camera_fovy,
training_iterations,
training_resolution,
marching_cude_grids_resolution,
marching_cude_grids_batch_size,
marching_cude_threshold,
texture_resolution,
force_cuda_rast
):
with torch.inference_mode(False):
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
ngp = InstantNGP(training_resolution).to(device)
ngp.prepare_training(reference_image, reference_mask, reference_orbit_camera_poses, reference_orbit_camera_fovy)
ngp.fit_nerf(training_iterations)
vertices, triangles = marching_cubes_density_to_mesh(ngp.get_density, marching_cude_grids_resolution, marching_cude_grids_batch_size, marching_cude_threshold)
v = torch.from_numpy(vertices).contiguous().float().to(device)
f = torch.from_numpy(triangles).contiguous().int().to(device)
mesh = Mesh(v=v, f=f, albedo=None, device=device)
mesh.auto_normal()
mesh.auto_uv()
mesh.albedo = color_func_to_albedo(mesh, ngp.get_color, texture_resolution, device=device, force_cuda_rast=force_cuda_rast)
return (mesh, )