diff --git a/README.md b/README.md index 4d78105..700c9ed 100644 --- a/README.md +++ b/README.md @@ -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 diff --git a/_Example_Workflows/Multi-View-Images_to_3DGS_to_3DMesh(DMTet and DiffRast).json b/_Example_Workflows/Multi-View-Images_to_3DGS_to_3DMesh(DMTet and DiffRast).json index b90a814..b06cdc1 100644 --- a/_Example_Workflows/Multi-View-Images_to_3DGS_to_3DMesh(DMTet and DiffRast).json +++ b/_Example_Workflows/Multi-View-Images_to_3DGS_to_3DMesh(DMTet and DiffRast).json @@ -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": [ diff --git a/_Example_Workflows/Multi-View-Images_to_Instant-NGP_to_3DMesh.json b/_Example_Workflows/Multi-View-Images_to_Instant-NGP_to_3DMesh.json new file mode 100644 index 0000000..8264dc0 --- /dev/null +++ b/_Example_Workflows/Multi-View-Images_to_Instant-NGP_to_3DMesh.json @@ -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 +} \ No newline at end of file diff --git a/_Example_Workflows/_Example_Inputs_Files/Meshes/CatCoffee.ply b/_Example_Workflows/_Example_Inputs_Files/Meshes/CatCoffee.ply new file mode 100644 index 0000000..88d4320 Binary files /dev/null and b/_Example_Workflows/_Example_Inputs_Files/Meshes/CatCoffee.ply differ diff --git a/algorithms/Instant_NGP.py b/algorithms/Instant_NGP.py new file mode 100644 index 0000000..b26b71d --- /dev/null +++ b/algorithms/Instant_NGP.py @@ -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!") \ No newline at end of file diff --git a/algorithms/diff_texturing.py b/algorithms/diff_mesh.py similarity index 74% rename from algorithms/diff_texturing.py rename to algorithms/diff_mesh.py index 2377830..76a5c7d 100644 --- a/algorithms/diff_texturing.py +++ b/algorithms/diff_mesh.py @@ -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() diff --git a/algorithms/diff_mesh_renderer.py b/algorithms/diff_mesh_renderer.py index 7a413eb..6f55573 100644 --- a/algorithms/diff_mesh_renderer.py +++ b/algorithms/diff_mesh_renderer.py @@ -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): diff --git a/algorithms/gaussian_model.txt b/algorithms/gaussian_model.txt new file mode 100644 index 0000000..e8de3f0 --- /dev/null +++ b/algorithms/gaussian_model.txt @@ -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 diff --git a/algorithms/main_3DGS_renderer.py b/algorithms/main_3DGS_renderer.py index 848e6be..3a138c5 100644 --- a/algorithms/main_3DGS_renderer.py +++ b/algorithms/main_3DGS_renderer.py @@ -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 ): diff --git a/algorithms/nerf_marching_cubes_converter.py b/algorithms/nerf_marching_cubes_converter.py index b0513b8..1f52ed7 100644 --- a/algorithms/nerf_marching_cubes_converter.py +++ b/algorithms/nerf_marching_cubes_converter.py @@ -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 diff --git a/mesh_processer/mesh_utils.py b/mesh_processer/mesh_utils.py index 34c1d75..29256ff 100644 --- a/mesh_processer/mesh_utils.py +++ b/mesh_processer/mesh_utils.py @@ -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 - \ No newline at end of file + + +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 \ No newline at end of file diff --git a/nodes.py b/nodes.py index baacc95..c6a6193 100644 --- a/nodes.py +++ b/nodes.py @@ -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,) \ No newline at end of file + 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, ) \ No newline at end of file