diff --git a/README.md b/README.md
index dee4782..56386cd 100644
--- a/README.md
+++ b/README.md
@@ -12,26 +12,23 @@
### Note: this project is still a WIP and not been released into ComFyUI package database yet
## Currently support:
-- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 04/02/2024**]
+- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 09/02/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
+
+- **Large Multiview Gaussian Model**: [3DTopia/LGM](https://github.com/3DTopia/LGM)
+ - Enable single image to 3D Gaussian in less than 30 seconds on a RTX3080 GPU, later you can also convert 3D Gaussian to mesh
+
+
+
+- **Triplane Gaussian Transformers**: [VAST-AI-Research/TriplaneGaussian](https://github.com/VAST-AI-Research/TriplaneGaussian)
+ - Enable single image to 3D Gaussian in less than 10 seconds on a RTX3080 GPU, later you can also convert 3D Gaussian to mesh
+
+
+
- **Preview 3DGS and 3D Mesh**: 3D Visualization inside ComfyUI:
- Using [gsplat.js](https://github.com/huggingface/gsplat.js/tree/main) and [three.js](https://github.com/mrdoob/three.js/tree/dev) for 3DGS & 3D Mesh visualization respectively
- https://github.com/MrForExample/ComfyUI-3D-Pack/assets/62230687/9f3c56b1-afb3-4bf1-8845-ab1025a87463
-
-- **Triplane Gaussian Transformers**: [VAST-AI-Research/TriplaneGaussian](https://github.com/VAST-AI-Research/TriplaneGaussian)
- - Enable single image to 3D Gaussian, then you can use Deep Marching Tetrahedrons node to convert it to mesh
-
-
-
-
- 
- |
-
-
- |
-
-
+
- **Stack Orbit Camera Poses**: Automatically generate all range of camera pose combinations
- You can use it to conditioning the [StableZero123 (You need to Download the checkpoint first)](https://comfyanonymous.github.io/ComfyUI_examples/3d/), with full range of camera poses in one prompt pass
@@ -71,17 +68,20 @@
- Since different algorithms likely use different coordinate system, so the ability to re-mapping the axis of coordinate is crucial for passing generated result between differnt nodes.
## Roadmap:
-- [x] Add DMTet algorithm to allow convertion from points cloud(Gaussian/.ply) to mesh (.obj, .ply, .glb)
+- [x] Add DMTet algorithm to allow conversion from points cloud(Gaussian/.ply) to mesh (.obj, .ply, .glb)
-- [x] Integrate [Triplane Meets Gaussian Splatting: Fast and Generalizable Single-View 3D Reconstruction with Transformers](https://github.com/VAST-AI-Research/TriplaneGaussian)
+- [x] Integrate [Triplane Meets Gaussian Splatting: Fast and Generalizable Single-View 3D Reconstruction with Transformers](https://zouzx.github.io/TriplaneGaussian/)
- [x] Add interactive 3D UI inside ComfuUI to visulaize training and generated results for 3D representations
- [x] Add a new node to generate renderer image sequence given a 3D gaussians and orbit camera poses (So we can later feed it to the differentiable renderer to bake it onto a given mesh)
-- [ ] Add a general SDS/VSD Optimization algorithm to allow training 3D representations with diffusion model, *The real fun begins here* ;)
+- [x] Integrate [LGM: Large Multi-View Gaussian Model for High-Resolution 3D Content Creation](https://me.kiui.moe/lgm/)
-- [ ] Improve 3DGS to Mesh:
+- [ ] Add a general SDS/ISM Optimization algorithm to allow training 3D representations with diffusion model, *The real fun starts here* ;)
+ - Need to do some in-depth research on Interval Score Matching (ISM), since math behind it makes perfect sense and also there are so many ways we could improve upon the result obtained from [LucidDreamer](https://github.com/EnVision-Research/LucidDreamer)
+
+- [ ] Improve 3DGS to Mesh conversion algorithms:
- Support to training DMTet with images(RGB, Alpha, Normal Map)
- Find better methods to converts 3DGS or Points Cloud to Mesh (Normal maps reconstruction maybe?)
diff --git a/_Example_Workflows/Large_Multiview_Gaussian_Model_to_3DGS and_3DMesh.json b/_Example_Workflows/Large_Multiview_Gaussian_Model_to_3DGS and_3DMesh.json
new file mode 100644
index 0000000..27b84e1
--- /dev/null
+++ b/_Example_Workflows/Large_Multiview_Gaussian_Model_to_3DGS and_3DMesh.json
@@ -0,0 +1,627 @@
+{
+ "last_node_id": 16,
+ "last_link_id": 18,
+ "nodes": [
+ {
+ "id": 3,
+ "type": "[Comfy3D] MVDream Model",
+ "pos": [
+ 792,
+ 460
+ ],
+ "size": {
+ "0": 405.5999755859375,
+ "1": 238
+ },
+ "flags": {},
+ "order": 5,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "mvdream_pipe",
+ "type": "DIFFUSERS_PIPE",
+ "link": 1
+ },
+ {
+ "name": "reference_image",
+ "type": "IMAGE",
+ "link": 2
+ },
+ {
+ "name": "reference_mask",
+ "type": "MASK",
+ "link": 5
+ }
+ ],
+ "outputs": [
+ {
+ "name": "multiview_images",
+ "type": "IMAGE",
+ "links": [
+ 6
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] MVDream Model"
+ },
+ "widgets_values": [
+ "",
+ "ugly, blurry, pixelated obscure, unnatural colors, poor lighting, dull, unclear, cropped, lowres, low quality, artifacts, duplicate",
+ 5,
+ 30,
+ 0
+ ]
+ },
+ {
+ "id": 4,
+ "type": "LoadImage",
+ "pos": [
+ 177,
+ 479
+ ],
+ "size": {
+ "0": 315,
+ "1": 314
+ },
+ "flags": {},
+ "order": 0,
+ "mode": 0,
+ "outputs": [
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 2
+ ],
+ "shape": 3,
+ "slot_index": 0
+ },
+ {
+ "name": "MASK",
+ "type": "MASK",
+ "links": [
+ 3
+ ],
+ "shape": 3,
+ "slot_index": 1
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "LoadImage"
+ },
+ "widgets_values": [
+ "Cat_Coffee_rgba.png",
+ "image"
+ ]
+ },
+ {
+ "id": 5,
+ "type": "InvertMask",
+ "pos": [
+ 539,
+ 689
+ ],
+ "size": {
+ "0": 210,
+ "1": 26
+ },
+ "flags": {},
+ "order": 3,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "mask",
+ "type": "MASK",
+ "link": 3
+ }
+ ],
+ "outputs": [
+ {
+ "name": "MASK",
+ "type": "MASK",
+ "links": [
+ 5
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "InvertMask"
+ }
+ },
+ {
+ "id": 2,
+ "type": "[Comfy3D] Load Diffusers Pipeline",
+ "pos": [
+ 381,
+ 319
+ ],
+ "size": {
+ "0": 375.3948669433594,
+ "1": 83.33004760742188
+ },
+ "flags": {},
+ "order": 1,
+ "mode": 0,
+ "outputs": [
+ {
+ "name": "pipe",
+ "type": "DIFFUSERS_PIPE",
+ "links": [
+ 1,
+ 8
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Load Diffusers Pipeline"
+ },
+ "widgets_values": [
+ "MVDreamPipeline",
+ "ashawkey/imagedream-ipmv-diffusers"
+ ]
+ },
+ {
+ "id": 12,
+ "type": "[Comfy3D] Switch 3DGS Axis",
+ "pos": [
+ 1554,
+ 460
+ ],
+ "size": {
+ "0": 210,
+ "1": 106
+ },
+ "flags": {},
+ "order": 7,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "gs_ply",
+ "type": "GS_PLY",
+ "link": 12
+ }
+ ],
+ "outputs": [
+ {
+ "name": "switched_gs_ply",
+ "type": "GS_PLY",
+ "links": [
+ 13
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Switch 3DGS Axis"
+ },
+ "widgets_values": [
+ "+x",
+ "-y",
+ "-z"
+ ]
+ },
+ {
+ "id": 8,
+ "type": "Reroute",
+ "pos": [
+ 1150,
+ 319
+ ],
+ "size": [
+ 75,
+ 26
+ ],
+ "flags": {},
+ "order": 4,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "",
+ "type": "*",
+ "link": 8
+ }
+ ],
+ "outputs": [
+ {
+ "name": "",
+ "type": "DIFFUSERS_PIPE",
+ "links": [
+ 9
+ ],
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "showOutputText": false,
+ "horizontal": false
+ }
+ },
+ {
+ "id": 9,
+ "type": "[Comfy3D] Save 3DGS",
+ "pos": [
+ 1798,
+ 459
+ ],
+ "size": {
+ "0": 315,
+ "1": 58
+ },
+ "flags": {},
+ "order": 9,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "gs_ply",
+ "type": "GS_PLY",
+ "link": 13
+ }
+ ],
+ "outputs": [
+ {
+ "name": "save_path",
+ "type": "STRING",
+ "links": [
+ 11
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Save 3DGS"
+ },
+ "widgets_values": [
+ "LGMTest/FatCat.ply"
+ ]
+ },
+ {
+ "id": 10,
+ "type": "[Comfy3D] Preview 3DGS",
+ "pos": [
+ 2309.8418660992434,
+ 768.9263430496216
+ ],
+ "size": [
+ 600,
+ 500
+ ],
+ "flags": {},
+ "order": 11,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "gs_file_path",
+ "type": "STRING",
+ "link": 11,
+ "widget": {
+ "name": "gs_file_path"
+ }
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Preview 3DGS"
+ },
+ "widgets_values": [
+ "",
+ null
+ ]
+ },
+ {
+ "id": 1,
+ "type": "[Comfy3D] Load Large Multiview Gaussian Model",
+ "pos": [
+ 879,
+ 752
+ ],
+ "size": {
+ "0": 315,
+ "1": 82
+ },
+ "flags": {},
+ "order": 2,
+ "mode": 0,
+ "outputs": [
+ {
+ "name": "lgm_model",
+ "type": "LGM_MODEL",
+ "links": [
+ 7
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Load Large Multiview Gaussian Model"
+ },
+ "widgets_values": [
+ "model_fp16.safetensors",
+ "big"
+ ]
+ },
+ {
+ "id": 7,
+ "type": "[Comfy3D] Large Multiview Gaussian Model",
+ "pos": [
+ 1270,
+ 460
+ ],
+ "size": {
+ "0": 252,
+ "1": 66
+ },
+ "flags": {},
+ "order": 6,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "multiview_images",
+ "type": "IMAGE",
+ "link": 6
+ },
+ {
+ "name": "lgm_model",
+ "type": "LGM_MODEL",
+ "link": 7
+ },
+ {
+ "name": "mvdream_pipe",
+ "type": "DIFFUSERS_PIPE",
+ "link": 9
+ }
+ ],
+ "outputs": [
+ {
+ "name": "gs_ply",
+ "type": "GS_PLY",
+ "links": [
+ 12,
+ 18
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Large Multiview Gaussian Model"
+ }
+ },
+ {
+ "id": 15,
+ "type": "[Comfy3D] Preview 3DMesh",
+ "pos": [
+ 2312,
+ 226
+ ],
+ "size": [
+ 600,
+ 500
+ ],
+ "flags": {},
+ "order": 12,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "mesh_file_path",
+ "type": "STRING",
+ "link": 16,
+ "widget": {
+ "name": "mesh_file_path"
+ }
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Preview 3DMesh"
+ },
+ "widgets_values": [
+ "",
+ null
+ ]
+ },
+ {
+ "id": 16,
+ "type": "[Comfy3D] Convert 3DGS to Mesh with NeRF and Marching Cubes",
+ "pos": [
+ 1550,
+ 224
+ ],
+ "size": [
+ 411.6000061035156,
+ 65.46688079833984
+ ],
+ "flags": {},
+ "order": 8,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "gs_ply",
+ "type": "GS_PLY",
+ "link": 18
+ }
+ ],
+ "outputs": [
+ {
+ "name": "mesh",
+ "type": "MESH",
+ "links": [
+ 17
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Convert 3DGS to Mesh with NeRF and Marching Cubes"
+ },
+ "widgets_values": [
+ "big"
+ ]
+ },
+ {
+ "id": 14,
+ "type": "[Comfy3D] Save 3D Mesh",
+ "pos": [
+ 1980,
+ 225
+ ],
+ "size": {
+ "0": 307.8253173828125,
+ "1": 60.248291015625
+ },
+ "flags": {},
+ "order": 10,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "mesh",
+ "type": "MESH",
+ "link": 17
+ }
+ ],
+ "outputs": [
+ {
+ "name": "save_path",
+ "type": "STRING",
+ "links": [
+ 16
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Save 3D Mesh"
+ },
+ "widgets_values": [
+ "LGMTest/FatCat.obj"
+ ]
+ }
+ ],
+ "links": [
+ [
+ 1,
+ 2,
+ 0,
+ 3,
+ 0,
+ "DIFFUSERS_PIPE"
+ ],
+ [
+ 2,
+ 4,
+ 0,
+ 3,
+ 1,
+ "IMAGE"
+ ],
+ [
+ 3,
+ 4,
+ 1,
+ 5,
+ 0,
+ "MASK"
+ ],
+ [
+ 5,
+ 5,
+ 0,
+ 3,
+ 2,
+ "MASK"
+ ],
+ [
+ 6,
+ 3,
+ 0,
+ 7,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 7,
+ 1,
+ 0,
+ 7,
+ 1,
+ "LGM_MODEL"
+ ],
+ [
+ 8,
+ 2,
+ 0,
+ 8,
+ 0,
+ "*"
+ ],
+ [
+ 9,
+ 8,
+ 0,
+ 7,
+ 2,
+ "DIFFUSERS_PIPE"
+ ],
+ [
+ 11,
+ 9,
+ 0,
+ 10,
+ 0,
+ "STRING"
+ ],
+ [
+ 12,
+ 7,
+ 0,
+ 12,
+ 0,
+ "GS_PLY"
+ ],
+ [
+ 13,
+ 12,
+ 0,
+ 9,
+ 0,
+ "GS_PLY"
+ ],
+ [
+ 16,
+ 14,
+ 0,
+ 15,
+ 0,
+ "STRING"
+ ],
+ [
+ 17,
+ 16,
+ 0,
+ 14,
+ 0,
+ "MESH"
+ ],
+ [
+ 18,
+ 7,
+ 0,
+ 16,
+ 0,
+ "GS_PLY"
+ ]
+ ],
+ "groups": [],
+ "config": {},
+ "extra": {},
+ "version": 0.4
+}
\ No newline at end of file
diff --git a/_Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS.json b/_Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS(DMTet and DiffRast).json
similarity index 100%
rename from _Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS.json
rename to _Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS(DMTet and DiffRast).json
diff --git a/_Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS_and_3DMesh.json b/_Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS_and_3DMesh.json
new file mode 100644
index 0000000..506d65a
--- /dev/null
+++ b/_Example_Workflows/Triplane_Gaussian_Transformers_to_3DGS_and_3DMesh.json
@@ -0,0 +1,580 @@
+{
+ "last_node_id": 34,
+ "last_link_id": 52,
+ "nodes": [
+ {
+ "id": 9,
+ "type": "InvertMask",
+ "pos": [
+ 576,
+ 690
+ ],
+ "size": {
+ "0": 140,
+ "1": 26
+ },
+ "flags": {},
+ "order": 4,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "mask",
+ "type": "MASK",
+ "link": 10
+ }
+ ],
+ "outputs": [
+ {
+ "name": "MASK",
+ "type": "MASK",
+ "links": [
+ 15
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "InvertMask"
+ }
+ },
+ {
+ "id": 27,
+ "type": "ImpactFloat",
+ "pos": [
+ 915,
+ -357
+ ],
+ "size": {
+ "0": 315,
+ "1": 58
+ },
+ "flags": {},
+ "order": 0,
+ "mode": 0,
+ "outputs": [
+ {
+ "name": "FLOAT",
+ "type": "FLOAT",
+ "links": [],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "title": "Camera Fovy",
+ "properties": {
+ "Node name for S&R": "ImpactFloat"
+ },
+ "widgets_values": [
+ 49.1
+ ]
+ },
+ {
+ "id": 12,
+ "type": "[Comfy3D] Preview 3DGS",
+ "pos": [
+ 1590,
+ 800
+ ],
+ "size": [
+ 600,
+ 500
+ ],
+ "flags": {
+ "collapsed": false
+ },
+ "order": 9,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "gs_file_path",
+ "type": "STRING",
+ "link": 13,
+ "widget": {
+ "name": "gs_file_path"
+ }
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Preview 3DGS"
+ },
+ "widgets_values": [
+ "",
+ null
+ ]
+ },
+ {
+ "id": 14,
+ "type": "[Comfy3D] Triplane Gaussian Transformers",
+ "pos": [
+ 767,
+ 631
+ ],
+ "size": {
+ "0": 252,
+ "1": 98
+ },
+ "flags": {},
+ "order": 5,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "reference_image",
+ "type": "IMAGE",
+ "link": 14
+ },
+ {
+ "name": "reference_mask",
+ "type": "MASK",
+ "link": 15
+ },
+ {
+ "name": "tgs_model",
+ "type": "TGS_MODEL",
+ "link": 16
+ }
+ ],
+ "outputs": [
+ {
+ "name": "gs_ply",
+ "type": "GS_PLY",
+ "links": [
+ 17
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Triplane Gaussian Transformers"
+ },
+ "widgets_values": [
+ 1.9000000000000001
+ ]
+ },
+ {
+ "id": 26,
+ "type": "[Comfy3D] Preview 3DMesh",
+ "pos": [
+ 2328,
+ 796
+ ],
+ "size": [
+ 600,
+ 500
+ ],
+ "flags": {},
+ "order": 12,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "mesh_file_path",
+ "type": "STRING",
+ "link": 32,
+ "widget": {
+ "name": "mesh_file_path"
+ }
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Preview 3DMesh"
+ },
+ "widgets_values": [
+ "",
+ null
+ ]
+ },
+ {
+ "id": 11,
+ "type": "Note",
+ "pos": [
+ 953,
+ 429
+ ],
+ "size": {
+ "0": 347.9384460449219,
+ "1": 141.6276092529297
+ },
+ "flags": {},
+ "order": 1,
+ "mode": 0,
+ "title": "Switch 3DGS Axis Note",
+ "properties": {
+ "text": ""
+ },
+ "widgets_values": [
+ "Triplane Gaussian Transformers to Preview:\n axis x to: -y\n axis y to: +z\n axis z to: -x\n\nGaussian Splatting to Preview:\n axis x to: +x\n axis y to: -y\n axis z to: -z"
+ ],
+ "color": "#432",
+ "bgcolor": "#653"
+ },
+ {
+ "id": 13,
+ "type": "[Comfy3D] Switch 3DGS Axis",
+ "pos": [
+ 1052,
+ 631
+ ],
+ "size": {
+ "0": 210,
+ "1": 106
+ },
+ "flags": {},
+ "order": 6,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "gs_ply",
+ "type": "GS_PLY",
+ "link": 17
+ }
+ ],
+ "outputs": [
+ {
+ "name": "switched_gs_ply",
+ "type": "GS_PLY",
+ "links": [
+ 12,
+ 51
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Switch 3DGS Axis"
+ },
+ "widgets_values": [
+ "-y",
+ "+z",
+ "-x"
+ ]
+ },
+ {
+ "id": 10,
+ "type": "[Comfy3D] Save 3DGS",
+ "pos": [
+ 1288,
+ 800
+ ],
+ "size": {
+ "0": 275.57391357421875,
+ "1": 58
+ },
+ "flags": {},
+ "order": 7,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "gs_ply",
+ "type": "GS_PLY",
+ "link": 12
+ }
+ ],
+ "outputs": [
+ {
+ "name": "save_path",
+ "type": "STRING",
+ "links": [
+ 13
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Save 3DGS"
+ },
+ "widgets_values": [
+ "MeshTest\\3DGS_Picacho.ply"
+ ]
+ },
+ {
+ "id": 34,
+ "type": "[Comfy3D] Switch 3DGS Axis",
+ "pos": [
+ 1294,
+ 629
+ ],
+ "size": {
+ "0": 210,
+ "1": 106
+ },
+ "flags": {},
+ "order": 8,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "gs_ply",
+ "type": "GS_PLY",
+ "link": 51
+ }
+ ],
+ "outputs": [
+ {
+ "name": "switched_gs_ply",
+ "type": "GS_PLY",
+ "links": [
+ 52
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Switch 3DGS Axis"
+ },
+ "widgets_values": [
+ "+x",
+ "-y",
+ "-z"
+ ]
+ },
+ {
+ "id": 33,
+ "type": "[Comfy3D] Convert 3DGS to Mesh with NeRF and Marching Cubes",
+ "pos": [
+ 1531,
+ 629
+ ],
+ "size": {
+ "0": 411.6000061035156,
+ "1": 58
+ },
+ "flags": {},
+ "order": 10,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "gs_ply",
+ "type": "GS_PLY",
+ "link": 52
+ }
+ ],
+ "outputs": [
+ {
+ "name": "mesh",
+ "type": "MESH",
+ "links": [
+ 50
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Convert 3DGS to Mesh with NeRF and Marching Cubes"
+ },
+ "widgets_values": [
+ "big"
+ ]
+ },
+ {
+ "id": 25,
+ "type": "[Comfy3D] Save 3D Mesh",
+ "pos": [
+ 1974,
+ 629
+ ],
+ "size": {
+ "0": 315,
+ "1": 58
+ },
+ "flags": {},
+ "order": 11,
+ "mode": 0,
+ "inputs": [
+ {
+ "name": "mesh",
+ "type": "MESH",
+ "link": 50
+ }
+ ],
+ "outputs": [
+ {
+ "name": "save_path",
+ "type": "STRING",
+ "links": [
+ 32
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Save 3D Mesh"
+ },
+ "widgets_values": [
+ "MeshTest\\3DMesh_Picacho.obj"
+ ]
+ },
+ {
+ "id": 2,
+ "type": "[Comfy3D] Load Triplane Gaussian Transformers",
+ "pos": [
+ 444,
+ 768
+ ],
+ "size": {
+ "0": 294,
+ "1": 58
+ },
+ "flags": {},
+ "order": 2,
+ "mode": 0,
+ "outputs": [
+ {
+ "name": "tgs_model",
+ "type": "TGS_MODEL",
+ "links": [
+ 16
+ ],
+ "shape": 3,
+ "slot_index": 0
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "[Comfy3D] Load Triplane Gaussian Transformers"
+ },
+ "widgets_values": [
+ "model_lvis_rel.ckpt"
+ ]
+ },
+ {
+ "id": 4,
+ "type": "LoadImage",
+ "pos": [
+ 52,
+ 632
+ ],
+ "size": {
+ "0": 354.32330322265625,
+ "1": 315.4895935058594
+ },
+ "flags": {},
+ "order": 3,
+ "mode": 0,
+ "outputs": [
+ {
+ "name": "IMAGE",
+ "type": "IMAGE",
+ "links": [
+ 14
+ ],
+ "shape": 3,
+ "slot_index": 0
+ },
+ {
+ "name": "MASK",
+ "type": "MASK",
+ "links": [
+ 10
+ ],
+ "shape": 3,
+ "slot_index": 1
+ }
+ ],
+ "properties": {
+ "Node name for S&R": "LoadImage"
+ },
+ "widgets_values": [
+ "a_pikachu_with_smily_face-removebg.png",
+ "image"
+ ]
+ }
+ ],
+ "links": [
+ [
+ 10,
+ 4,
+ 1,
+ 9,
+ 0,
+ "MASK"
+ ],
+ [
+ 12,
+ 13,
+ 0,
+ 10,
+ 0,
+ "GS_PLY"
+ ],
+ [
+ 13,
+ 10,
+ 0,
+ 12,
+ 0,
+ "STRING"
+ ],
+ [
+ 14,
+ 4,
+ 0,
+ 14,
+ 0,
+ "IMAGE"
+ ],
+ [
+ 15,
+ 9,
+ 0,
+ 14,
+ 1,
+ "MASK"
+ ],
+ [
+ 16,
+ 2,
+ 0,
+ 14,
+ 2,
+ "TGS_MODEL"
+ ],
+ [
+ 17,
+ 14,
+ 0,
+ 13,
+ 0,
+ "GS_PLY"
+ ],
+ [
+ 32,
+ 25,
+ 0,
+ 26,
+ 0,
+ "STRING"
+ ],
+ [
+ 50,
+ 33,
+ 0,
+ 25,
+ 0,
+ "MESH"
+ ],
+ [
+ 51,
+ 13,
+ 0,
+ 34,
+ 0,
+ "GS_PLY"
+ ],
+ [
+ 52,
+ 34,
+ 0,
+ 33,
+ 0,
+ "GS_PLY"
+ ]
+ ],
+ "groups": [],
+ "config": {},
+ "extra": {},
+ "version": 0.4
+}
\ No newline at end of file
diff --git a/_Example_Workflows/_Example_Inputs_Files/Cat_Coffee_rgba.png b/_Example_Workflows/_Example_Inputs_Files/Cat_Coffee_rgba.png
new file mode 100644
index 0000000..1dc0b52
Binary files /dev/null and b/_Example_Workflows/_Example_Inputs_Files/Cat_Coffee_rgba.png differ
diff --git a/_Example_Workflows/_Example_Inputs_Files/Cucumber_Shark_rgba.png b/_Example_Workflows/_Example_Inputs_Files/Cucumber_Shark_rgba.png
new file mode 100644
index 0000000..70deed2
Binary files /dev/null and b/_Example_Workflows/_Example_Inputs_Files/Cucumber_Shark_rgba.png differ
diff --git a/_Example_Workflows/_Example_Inputs_Files/ghost_rgba.png b/_Example_Workflows/_Example_Inputs_Files/ghost_rgba.png
new file mode 100644
index 0000000..6ed0538
Binary files /dev/null and b/_Example_Workflows/_Example_Inputs_Files/ghost_rgba.png differ
diff --git a/algorithms/diff_texturing.py b/algorithms/diff_texturing.py
index 3388a78..2377830 100644
--- a/algorithms/diff_texturing.py
+++ b/algorithms/diff_texturing.py
@@ -4,12 +4,13 @@ import tqdm
import torch
import torch.nn.functional as F
import numpy as np
+from kiui.cam import orbit_camera
from pytorch_msssim import SSIM, MS_SSIM
import comfy.utils
from .diff_mesh_renderer import DiffRastRenderer
-from ..shared_utils.camera_utils import orbit_camera, OrbitCamera
+from ..shared_utils.camera_utils import OrbitCamera
from ..shared_utils.image_utils import prepare_torch_img
class DiffTextureBaker:
@@ -44,8 +45,8 @@ class DiffTextureBaker:
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], self.ref_size_H, self.ref_size_W, self.device))
- ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2), self.ref_size_H, self.ref_size_W, self.device))
+ 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)
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
diff --git a/algorithms/large_multiview_gaussian_model.py b/algorithms/large_multiview_gaussian_model.py
new file mode 100644
index 0000000..0e6fcc4
--- /dev/null
+++ b/algorithms/large_multiview_gaussian_model.py
@@ -0,0 +1,171 @@
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from kiui.lpips import LPIPS
+
+import kiui
+from kiui.cam import orbit_camera
+
+from ..lgm.core.unet import UNet
+from ..lgm.core.options import Options
+from ..lgm.core.gs import GaussianRenderer
+from ..lgm.core.utils import get_rays
+
+class LGM(nn.Module):
+ def __init__(
+ self,
+ opt: Options,
+ ):
+ super().__init__()
+
+ self.opt = opt
+
+ # unet
+ self.unet = UNet(
+ 9, 14,
+ down_channels=self.opt.down_channels,
+ down_attention=self.opt.down_attention,
+ mid_attention=self.opt.mid_attention,
+ up_channels=self.opt.up_channels,
+ up_attention=self.opt.up_attention,
+ )
+
+ # last conv
+ self.conv = nn.Conv2d(14, 14, kernel_size=1) # NOTE: maybe remove it if train again
+
+ # Gaussian Renderer
+ self.gs = GaussianRenderer(opt)
+
+ # activations...
+ self.pos_act = lambda x: x.clamp(-1, 1)
+ self.scale_act = lambda x: 0.1 * F.softplus(x)
+ self.opacity_act = lambda x: torch.sigmoid(x)
+ self.rot_act = F.normalize
+ self.rgb_act = lambda x: 0.5 * torch.tanh(x) + 0.5 # NOTE: may use sigmoid if train again
+
+ # LPIPS loss
+ if self.opt.lambda_lpips > 0:
+ self.lpips_loss = LPIPS(net='vgg')
+ self.lpips_loss.requires_grad_(False)
+
+
+ def state_dict(self, **kwargs):
+ # remove lpips_loss
+ state_dict = super().state_dict(**kwargs)
+ for k in list(state_dict.keys()):
+ if 'lpips_loss' in k:
+ del state_dict[k]
+ return state_dict
+
+
+ def prepare_default_rays(self, device, elevation=0):
+ cam_poses = np.stack([
+ orbit_camera(elevation, 0, radius=self.opt.cam_radius),
+ orbit_camera(elevation, 90, radius=self.opt.cam_radius),
+ orbit_camera(elevation, 180, radius=self.opt.cam_radius),
+ orbit_camera(elevation, 270, radius=self.opt.cam_radius),
+ ], axis=0) # [4, 4, 4]
+ cam_poses = torch.from_numpy(cam_poses)
+
+ rays_embeddings = []
+ for i in range(cam_poses.shape[0]):
+ rays_o, rays_d = get_rays(cam_poses[i], self.opt.input_size, self.opt.input_size, self.opt.fovy) # [h, w, 3]
+ rays_plucker = torch.cat([torch.cross(rays_o, rays_d, dim=-1), rays_d], dim=-1) # [h, w, 6]
+ rays_embeddings.append(rays_plucker)
+
+ ## visualize rays for plotting figure
+ # kiui.vis.plot_image(rays_d * 0.5 + 0.5, save=True)
+
+ rays_embeddings = torch.stack(rays_embeddings, dim=0).permute(0, 3, 1, 2).contiguous().to(device) # [V, 6, h, w]
+
+ return rays_embeddings
+
+
+ def forward_gaussians(self, images):
+ # images: [B, 4, 9, H, W]
+ # return: Gaussians: [B, dim_t]
+
+ B, V, C, H, W = images.shape
+ images = images.view(B*V, C, H, W)
+
+ x = self.unet(images) # [B*4, 14, h, w]
+ x = self.conv(x) # [B*4, 14, h, w]
+
+ x = x.reshape(B, 4, 14, self.opt.splat_size, self.opt.splat_size)
+
+ ## visualize multi-view gaussian features for plotting figure
+ # tmp_alpha = self.opacity_act(x[0, :, 3:4])
+ # tmp_img_rgb = self.rgb_act(x[0, :, 11:]) * tmp_alpha + (1 - tmp_alpha)
+ # tmp_img_pos = self.pos_act(x[0, :, 0:3]) * 0.5 + 0.5
+ # kiui.vis.plot_image(tmp_img_rgb, save=True)
+ # kiui.vis.plot_image(tmp_img_pos, save=True)
+
+ x = x.permute(0, 1, 3, 4, 2).reshape(B, -1, 14)
+
+ pos = self.pos_act(x[..., 0:3]) # [B, N, 3]
+ opacity = self.opacity_act(x[..., 3:4])
+ scale = self.scale_act(x[..., 4:7])
+ rotation = self.rot_act(x[..., 7:11])
+ rgbs = self.rgb_act(x[..., 11:])
+
+ gaussians = torch.cat([pos, opacity, scale, rotation, rgbs], dim=-1) # [B, N, 14]
+
+ return gaussians
+
+
+ def forward(self, data, step_ratio=1):
+ # data: output of the dataloader
+ # return: loss
+
+ results = {}
+ loss = 0
+
+ images = data['input'] # [B, 4, 9, h, W], input features
+
+ # use the first view to predict gaussians
+ gaussians = self.forward_gaussians(images) # [B, N, 14]
+
+ results['gaussians'] = gaussians
+
+ # random bg for training
+ if self.training:
+ bg_color = torch.rand(3, dtype=torch.float32, device=gaussians.device)
+ else:
+ bg_color = torch.ones(3, dtype=torch.float32, device=gaussians.device)
+
+ # use the other views for rendering and supervision
+ results = self.gs.render(gaussians, data['cam_view'], data['cam_view_proj'], data['cam_pos'], bg_color=bg_color)
+ pred_images = results['image'] # [B, V, C, output_size, output_size]
+ pred_alphas = results['alpha'] # [B, V, 1, output_size, output_size]
+
+ results['images_pred'] = pred_images
+ results['alphas_pred'] = pred_alphas
+
+ gt_images = data['images_output'] # [B, V, 3, output_size, output_size], ground-truth novel views
+ gt_masks = data['masks_output'] # [B, V, 1, output_size, output_size], ground-truth masks
+
+ gt_images = gt_images * gt_masks + bg_color.view(1, 1, 3, 1, 1) * (1 - gt_masks)
+
+ loss_mse = F.mse_loss(pred_images, gt_images) + F.mse_loss(pred_alphas, gt_masks)
+ loss = loss + loss_mse
+
+ if self.opt.lambda_lpips > 0:
+ loss_lpips = self.lpips_loss(
+ # gt_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1,
+ # pred_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1,
+ # downsampled to at most 256 to reduce memory cost
+ F.interpolate(gt_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False),
+ F.interpolate(pred_images.view(-1, 3, self.opt.output_size, self.opt.output_size) * 2 - 1, (256, 256), mode='bilinear', align_corners=False),
+ ).mean()
+ results['loss_lpips'] = loss_lpips
+ loss = loss + self.opt.lambda_lpips * loss_lpips
+
+ results['loss'] = loss
+
+ # metric
+ with torch.no_grad():
+ psnr = -10 * torch.log10(torch.mean((pred_images.detach() - gt_images) ** 2))
+ results['psnr'] = psnr
+
+ return results
\ No newline at end of file
diff --git a/algorithms/main_3DGS.py b/algorithms/main_3DGS.py
index 4c28f9e..a0604d5 100644
--- a/algorithms/main_3DGS.py
+++ b/algorithms/main_3DGS.py
@@ -4,12 +4,13 @@ import tqdm
import torch
import torch.nn.functional as F
import numpy as np
+from kiui.cam import orbit_camera
from pytorch_msssim import SSIM, MS_SSIM
import comfy.utils
from .main_3DGS_renderer import GaussianSplattingRenderer
-from ..shared_utils.camera_utils import orbit_camera, OrbitCamera, MiniCam, calculate_fovX, get_projection_matrix
+from ..shared_utils.camera_utils import OrbitCamera, MiniCam, calculate_fovX, get_projection_matrix
from ..shared_utils.image_utils import prepare_torch_img
class GSParams:
@@ -153,8 +154,8 @@ class GaussianSplatting:
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], self.ref_size_H, self.ref_size_W, self.device))
- ref_masks_torch_list.append(prepare_torch_img(reference_masks[i].unsqueeze(2), self.ref_size_H, self.ref_size_W, self.device))
+ 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)
self.ref_masks_torch = torch.cat(ref_masks_torch_list, dim=0)
diff --git a/algorithms/nerf_marching_cubes_converter.py b/algorithms/nerf_marching_cubes_converter.py
new file mode 100644
index 0000000..a402c06
--- /dev/null
+++ b/algorithms/nerf_marching_cubes_converter.py
@@ -0,0 +1,416 @@
+import tqdm
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import mcubes
+import nerfacc
+import nvdiffrast.torch as dr
+
+from kiui.mesh_utils import clean_mesh, decimate_mesh
+from kiui.mesh_utils import laplacian_smooth_loss, normal_consistency
+from kiui.op import uv_padding, inverse_sigmoid
+from kiui.cam import orbit_camera, get_perspective
+from kiui.nn import MLP, trunc_exp
+
+from ..mesh_processer.mesh import Mesh
+from ..lgm.core.options import Options
+from ..lgm.core.gs import GaussianRenderer
+from ..lgm.core.utils import get_rays
+
+# Triple renderer of gaussians, gaussian, and diso mesh.
+# gaussian --> nerf --> mesh
+class GSConverterNeRFMarchingCubes(nn.Module):
+ def __init__(self, opt: Options, gs_ply):
+ super().__init__()
+ from kiui.gridencoder import GridEncoder
+
+ self.opt = opt
+ self.device = torch.device("cuda")
+
+ # gs renderer
+ self.tan_half_fov = np.tan(0.5 * np.deg2rad(opt.fovy))
+ self.proj_matrix = torch.zeros(4, 4, dtype=torch.float32, device=self.device)
+ self.proj_matrix[0, 0] = 1 / self.tan_half_fov
+ self.proj_matrix[1, 1] = 1 / self.tan_half_fov
+ self.proj_matrix[2, 2] = (opt.zfar + opt.znear) / (opt.zfar - opt.znear)
+ self.proj_matrix[3, 2] = - (opt.zfar * opt.znear) / (opt.zfar - opt.znear)
+ self.proj_matrix[2, 3] = 1
+
+ self.gs_renderer = GaussianRenderer(opt)
+
+ self.gaussians = self.gs_renderer.create_from_ply(gs_ply).to(self.device)
+
+ # nerf renderer
+ if not self.opt.force_cuda_rast:
+ self.glctx = dr.RasterizeGLContext()
+ else:
+ self.glctx = dr.RasterizeCudaContext()
+
+ self.step = 0
+ 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)
+
+ # mesh renderer
+ self.proj = torch.from_numpy(get_perspective(self.opt.fovy)).float().to(self.device)
+ self.v = self.f = None
+ self.vt = self.ft = None
+ self.deform = None
+ self.albedo = None
+
+
+ @torch.no_grad()
+ def render_gs(self, pose):
+
+ cam_poses = torch.from_numpy(pose).unsqueeze(0).to(self.device)
+ cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+
+ # cameras needed by gaussian rasterizer
+ cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+ cam_view_proj = cam_view @ self.proj_matrix # [V, 4, 4]
+ cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+ out = self.gs_renderer.render(self.gaussians.unsqueeze(0), cam_view.unsqueeze(0), cam_view_proj.unsqueeze(0), cam_pos.unsqueeze(0))
+ image = out['image'].squeeze(1).squeeze(0) # [C, H, W]
+ alpha = out['alpha'].squeeze(2).squeeze(1).squeeze(0) # [H, W]
+
+ return image, alpha
+
+ def get_density(self, xs):
+ # xs: [..., 3]
+ 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 render_nerf(self, pose):
+
+ pose = torch.from_numpy(pose.astype(np.float32)).to(self.device)
+
+ # get rays
+ resolution = self.opt.output_size
+ rays_o, rays_d = get_rays(pose, resolution, resolution, self.opt.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.step, occ_eval_fn=occ_eval_fn, occ_thre=0.01, n=8)
+ self.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, trans, 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(resolution, resolution, 3).clamp(0, 1).permute(2, 0, 1).contiguous()
+ alpha = alpha.view(resolution, resolution).clamp(0, 1).contiguous()
+
+ return color, alpha
+
+ def fit_nerf(self, iters=512, resolution=128):
+
+ self.opt.output_size = resolution
+
+ 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...")
+ pbar = tqdm.trange(iters)
+ for i in pbar:
+
+ ver = np.random.randint(-45, 45)
+ hor = np.random.randint(-180, 180)
+ rad = np.random.uniform(1.5, 3.0)
+
+ pose = orbit_camera(ver, hor, rad)
+
+ image_gt, alpha_gt = self.render_gs(pose)
+ 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.backward()
+ self.encoder_density.grad_total_variation(1e-8)
+
+ optimizer.step()
+ optimizer.zero_grad()
+
+ pbar.set_description(f"MSE = {loss_mse.item():.6f}")
+
+ print(f"[INFO] finished fitting nerf!")
+
+ def render_mesh(self, pose):
+
+ h = w = self.opt.output_size
+
+ v = self.v + self.deform
+ f = self.f
+
+ pose = torch.from_numpy(pose.astype(np.float32)).to(v.device)
+
+ # get v_clip and render rgb
+ v_cam = torch.matmul(F.pad(v, pad=(0, 1), mode='constant', value=1.0), torch.inverse(pose).T).float().unsqueeze(0)
+ v_clip = v_cam @ self.proj.T
+
+ rast, rast_db = dr.rasterize(self.glctx, v_clip, f, (h, w))
+
+ alpha = torch.clamp(rast[..., -1:], 0, 1).contiguous() # [1, H, W, 1]
+ alpha = dr.antialias(alpha, rast, v_clip, f).clamp(0, 1).squeeze(-1).squeeze(0) # [H, W] important to enable gradients!
+
+ if self.albedo is None:
+ xyzs, _ = dr.interpolate(v.unsqueeze(0), rast, f) # [1, H, W, 3]
+ xyzs = xyzs.view(-1, 3)
+ mask = (alpha > 0).view(-1)
+ image = torch.zeros_like(xyzs, dtype=torch.float32)
+ if mask.any():
+ masked_albedo = torch.sigmoid(self.mlp(self.encoder(xyzs[mask].detach(), bound=1)))
+ image[mask] = masked_albedo.float()
+ else:
+ texc, texc_db = dr.interpolate(self.vt.unsqueeze(0), rast, self.ft, rast_db=rast_db, diff_attrs='all')
+ image = torch.sigmoid(dr.texture(self.albedo.unsqueeze(0), texc, uv_da=texc_db)) # [1, H, W, 3]
+
+ image = image.view(1, h, w, 3)
+ # image = dr.antialias(image, rast, v_clip, f).clamp(0, 1)
+ image = image.squeeze(0).permute(2, 0, 1).contiguous() # [3, H, W]
+ image = alpha * image + (1 - alpha)
+
+ return image, alpha
+
+ def fit_mesh(self, iters=2048, resolution=512, decimate_target=5e4):
+
+ self.opt.output_size = resolution
+
+ # init mesh from nerf
+ grid_size = 256
+ sigmas = np.zeros([grid_size, grid_size, grid_size], dtype=np.float32)
+
+ S = 128
+ density_thresh = 10
+
+ 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 = self.get_density(pts.to(self.device))
+ 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)
+
+ 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)
+
+ # fit mesh from gs
+ lr_factor = 1
+ optimizer = torch.optim.Adam([
+ {'params': self.encoder.parameters(), 'lr': 1e-3 * lr_factor},
+ {'params': self.mlp.parameters(), 'lr': 1e-3 * lr_factor},
+ {'params': self.deform, 'lr': 1e-4},
+ ])
+
+ print(f"[INFO] fitting mesh...")
+ pbar = tqdm.trange(iters)
+ for i in pbar:
+
+ ver = np.random.randint(-10, 10)
+ hor = np.random.randint(-180, 180)
+ rad = self.opt.cam_radius # np.random.uniform(1, 2)
+
+ pose = orbit_camera(ver, hor, rad)
+
+ image_gt, alpha_gt = self.render_gs(pose)
+ image_pred, alpha_pred = self.render_mesh(pose)
+
+ loss_mse = F.mse_loss(image_pred, image_gt) + 0.1 * F.mse_loss(alpha_pred, alpha_gt)
+ # loss_lap = laplacian_smooth_loss(self.v + self.deform, self.f)
+ loss_normal = normal_consistency(self.v + self.deform, self.f)
+ loss_offsets = (self.deform ** 2).sum(-1).mean()
+ loss = loss_mse + 0.001 * loss_normal + 0.1 * loss_offsets
+
+ loss.backward()
+
+ optimizer.step()
+ optimizer.zero_grad()
+
+ # remesh periodically
+ if i > 0 and i % 512 == 0:
+ vertices = (self.v + self.deform).detach().cpu().numpy()
+ triangles = self.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.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)
+ lr_factor *= 0.5
+ optimizer = torch.optim.Adam([
+ {'params': self.encoder.parameters(), 'lr': 1e-3 * lr_factor},
+ {'params': self.mlp.parameters(), 'lr': 1e-3 * lr_factor},
+ {'params': self.deform, 'lr': 1e-4},
+ ])
+
+ pbar.set_description(f"MSE = {loss_mse.item():.6f}")
+
+ # last clean
+ vertices = (self.v + self.deform).detach().cpu().numpy()
+ triangles = self.f.detach().cpu().numpy()
+ vertices, triangles = clean_mesh(vertices, triangles, remesh=False)
+ 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))
+
+ print(f"[INFO] finished fitting mesh!")
+
+ # uv mesh refine
+ def fit_mesh_uv(self, iters=512, resolution=512, texture_resolution=1024, padding=2):
+
+ self.opt.output_size = resolution
+
+ # unwrap uv
+ print(f"[INFO] uv unwrapping...")
+ mesh = Mesh(v=self.v, f=self.f, albedo=None, device=self.device)
+ mesh.auto_normal()
+ mesh.auto_uv()
+
+ self.vt = mesh.vt
+ self.ft = mesh.ft
+
+ # 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(self.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=self.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 + 640000, xyzs.shape[0])
+ batch.append(torch.sigmoid(self.mlp(self.encoder(xyzs[head:tail]))).float())
+ head += 640000
+
+ albedo[mask] = torch.cat(batch, dim=0)
+
+ albedo = albedo.view(h, w, -1)
+ mask = mask.view(h, w)
+ albedo = uv_padding(albedo, mask, padding)
+
+ # optimize texture
+ self.albedo = nn.Parameter(inverse_sigmoid(albedo)).to(self.device)
+
+ optimizer = torch.optim.Adam([
+ {'params': self.albedo, 'lr': 1e-3},
+ ])
+
+ print(f"[INFO] fitting mesh texture...")
+ pbar = tqdm.trange(iters)
+ for i in pbar:
+
+ # shrink to front view as we care more about it...
+ ver = np.random.randint(-5, 5)
+ hor = np.random.randint(-15, 15)
+ rad = self.opt.cam_radius # np.random.uniform(1, 2)
+
+ pose = orbit_camera(ver, hor, rad)
+
+ image_gt, alpha_gt = self.render_gs(pose)
+ image_pred, alpha_pred = self.render_mesh(pose)
+
+ loss_mse = F.mse_loss(image_pred, image_gt)
+ loss = loss_mse
+
+ loss.backward()
+
+ optimizer.step()
+ optimizer.zero_grad()
+
+ pbar.set_description(f"MSE = {loss_mse.item():.6f}")
+
+ print(f"[INFO] finished fitting mesh texture!")
+
+
+ @torch.no_grad()
+ def get_mesh(self):
+ mesh = Mesh(v=self.v, f=self.f, vt=self.vt, ft=self.ft, albedo=torch.sigmoid(self.albedo), device=self.device)
+ mesh.auto_normal()
+ return mesh
diff --git a/install.bat b/install.bat
index 96ddaab..6521da2 100644
--- a/install.bat
+++ b/install.bat
@@ -10,15 +10,6 @@ if exist "%python_miniconda_exec%" (
echo Installing with Miniconda Environment
"%python_miniconda_exec%" -s -m pip install -r "%requirements_txt%"
- git clone --recursive https://github.com/ashawkey/diff-gaussian-rasterization
- "%python_miniconda_exec%" -s -m pip install ./diff-gaussian-rasterization
- "%python_miniconda_exec%" -s -m pip install ./simple-knn
- git clone --recursive https://github.com/NVlabs/nvdiffrast
- "%python_miniconda_exec%" -s -m pip install ./nvdiffrast
- "%python_miniconda_exec%" -s -m pip install git+https://github.com/facebookresearch/pytorch3d.git@stable
-
- "%python_miniconda_exec%" -s -m pip install git+https://github.com/rusty1s/pytorch_scatter.git
-
set "python_miniconda_exec=..\..\..\..\..\..\..\python_miniconda_env\ComfyUI\python.exe"
cd tgs/models/snowflake/pointnet2_ops_lib && "%python_miniconda_exec%" -s setup.py install && cd ../../../../
diff --git a/lgm/acc_configs/gpu1.yaml b/lgm/acc_configs/gpu1.yaml
new file mode 100644
index 0000000..bbbabf2
--- /dev/null
+++ b/lgm/acc_configs/gpu1.yaml
@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: 'NO'
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 1
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
diff --git a/lgm/acc_configs/gpu4.yaml b/lgm/acc_configs/gpu4.yaml
new file mode 100644
index 0000000..6f4c538
--- /dev/null
+++ b/lgm/acc_configs/gpu4.yaml
@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 4
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
diff --git a/lgm/acc_configs/gpu6.yaml b/lgm/acc_configs/gpu6.yaml
new file mode 100644
index 0000000..32ef41c
--- /dev/null
+++ b/lgm/acc_configs/gpu6.yaml
@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: fp16
+num_machines: 1
+num_processes: 6
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
diff --git a/lgm/acc_configs/gpu8.yaml b/lgm/acc_configs/gpu8.yaml
new file mode 100644
index 0000000..46b5c35
--- /dev/null
+++ b/lgm/acc_configs/gpu8.yaml
@@ -0,0 +1,15 @@
+compute_environment: LOCAL_MACHINE
+debug: false
+distributed_type: MULTI_GPU
+downcast_bf16: 'no'
+machine_rank: 0
+main_training_function: main
+mixed_precision: bf16
+num_machines: 1
+num_processes: 8
+rdzv_backend: static
+same_network: true
+tpu_env: []
+tpu_use_cluster: false
+tpu_use_sudo: false
+use_cpu: false
diff --git a/lgm/core/__init__.py b/lgm/core/__init__.py
new file mode 100644
index 0000000..e69de29
diff --git a/lgm/core/attention.py b/lgm/core/attention.py
new file mode 100644
index 0000000..ea1382f
--- /dev/null
+++ b/lgm/core/attention.py
@@ -0,0 +1,156 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+#
+# This source code is licensed under the Apache License, Version 2.0
+# found in the LICENSE file in the root directory of this source tree.
+
+# References:
+# https://github.com/facebookresearch/dino/blob/master/vision_transformer.py
+# https://github.com/rwightman/pytorch-image-models/tree/master/timm/models/vision_transformer.py
+
+import os
+import warnings
+
+from torch import Tensor
+from torch import nn
+
+XFORMERS_ENABLED = os.environ.get("XFORMERS_DISABLED") is None
+try:
+ if XFORMERS_ENABLED:
+ from xformers.ops import memory_efficient_attention, unbind
+
+ XFORMERS_AVAILABLE = True
+ warnings.warn("xFormers is available (Attention)")
+ else:
+ warnings.warn("xFormers is disabled (Attention)")
+ raise ImportError
+except ImportError:
+ XFORMERS_AVAILABLE = False
+ warnings.warn("xFormers is not available (Attention)")
+
+
+class Attention(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int = 8,
+ qkv_bias: bool = False,
+ proj_bias: bool = True,
+ attn_drop: float = 0.0,
+ proj_drop: float = 0.0,
+ ) -> None:
+ super().__init__()
+ self.num_heads = num_heads
+ head_dim = dim // num_heads
+ self.scale = head_dim**-0.5
+
+ self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)
+ self.attn_drop = nn.Dropout(attn_drop)
+ self.proj = nn.Linear(dim, dim, bias=proj_bias)
+ self.proj_drop = nn.Dropout(proj_drop)
+
+ def forward(self, x: Tensor) -> Tensor:
+ B, N, C = x.shape
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads).permute(2, 0, 3, 1, 4)
+
+ q, k, v = qkv[0] * self.scale, qkv[1], qkv[2]
+ attn = q @ k.transpose(-2, -1)
+
+ attn = attn.softmax(dim=-1)
+ attn = self.attn_drop(attn)
+
+ x = (attn @ v).transpose(1, 2).reshape(B, N, C)
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ return x
+
+
+class MemEffAttention(Attention):
+ def forward(self, x: Tensor, attn_bias=None) -> Tensor:
+ if not XFORMERS_AVAILABLE:
+ if attn_bias is not None:
+ raise AssertionError("xFormers is required for using nested tensors")
+ return super().forward(x)
+
+ B, N, C = x.shape
+ qkv = self.qkv(x).reshape(B, N, 3, self.num_heads, C // self.num_heads)
+
+ q, k, v = unbind(qkv, 2)
+
+ x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+ x = x.reshape([B, N, C])
+
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ return x
+
+
+class CrossAttention(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ dim_q: int,
+ dim_k: int,
+ dim_v: int,
+ num_heads: int = 8,
+ qkv_bias: bool = False,
+ proj_bias: bool = True,
+ attn_drop: float = 0.0,
+ proj_drop: float = 0.0,
+ ) -> None:
+ super().__init__()
+ self.dim = dim
+ self.num_heads = num_heads
+ head_dim = dim // num_heads
+ self.scale = head_dim**-0.5
+
+ self.to_q = nn.Linear(dim_q, dim, bias=qkv_bias)
+ self.to_k = nn.Linear(dim_k, dim, bias=qkv_bias)
+ self.to_v = nn.Linear(dim_v, dim, bias=qkv_bias)
+ self.attn_drop = nn.Dropout(attn_drop)
+ self.proj = nn.Linear(dim, dim, bias=proj_bias)
+ self.proj_drop = nn.Dropout(proj_drop)
+
+ def forward(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
+ # q: [B, N, Cq]
+ # k: [B, M, Ck]
+ # v: [B, M, Cv]
+ # return: [B, N, C]
+
+ B, N, _ = q.shape
+ M = k.shape[1]
+
+ q = self.scale * self.to_q(q).reshape(B, N, self.num_heads, self.dim // self.num_heads).permute(0, 2, 1, 3) # [B, nh, N, C/nh]
+ k = self.to_k(k).reshape(B, M, self.num_heads, self.dim // self.num_heads).permute(0, 2, 1, 3) # [B, nh, M, C/nh]
+ v = self.to_v(v).reshape(B, M, self.num_heads, self.dim // self.num_heads).permute(0, 2, 1, 3) # [B, nh, M, C/nh]
+
+ attn = q @ k.transpose(-2, -1) # [B, nh, N, M]
+
+ attn = attn.softmax(dim=-1) # [B, nh, N, M]
+ attn = self.attn_drop(attn)
+
+ x = (attn @ v).transpose(1, 2).reshape(B, N, -1) # [B, nh, N, M] @ [B, nh, M, C/nh] --> [B, nh, N, C/nh] --> [B, N, nh, C/nh] --> [B, N, C]
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ return x
+
+
+class MemEffCrossAttention(CrossAttention):
+ def forward(self, q: Tensor, k: Tensor, v: Tensor, attn_bias=None) -> Tensor:
+ if not XFORMERS_AVAILABLE:
+ if attn_bias is not None:
+ raise AssertionError("xFormers is required for using nested tensors")
+ return super().forward(x)
+
+ B, N, _ = q.shape
+ M = k.shape[1]
+
+ q = self.scale * self.to_q(q).reshape(B, N, self.num_heads, self.dim // self.num_heads) # [B, N, nh, C/nh]
+ k = self.to_k(k).reshape(B, M, self.num_heads, self.dim // self.num_heads) # [B, M, nh, C/nh]
+ v = self.to_v(v).reshape(B, M, self.num_heads, self.dim // self.num_heads) # [B, M, nh, C/nh]
+
+ x = memory_efficient_attention(q, k, v, attn_bias=attn_bias)
+ x = x.reshape(B, N, -1)
+
+ x = self.proj(x)
+ x = self.proj_drop(x)
+ return x
diff --git a/lgm/core/gs.py b/lgm/core/gs.py
new file mode 100644
index 0000000..f84274d
--- /dev/null
+++ b/lgm/core/gs.py
@@ -0,0 +1,186 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from diff_gaussian_rasterization import (
+ GaussianRasterizationSettings,
+ GaussianRasterizer,
+)
+
+from lgm.core.options import Options
+
+import kiui
+
+class GaussianRenderer:
+ def __init__(self, opt: Options):
+
+ self.opt = opt
+ self.bg_color = torch.tensor([1, 1, 1], dtype=torch.float32, device="cuda")
+
+ # intrinsics
+ self.tan_half_fov = np.tan(0.5 * np.deg2rad(self.opt.fovy))
+ self.proj_matrix = torch.zeros(4, 4, dtype=torch.float32)
+ self.proj_matrix[0, 0] = 1 / self.tan_half_fov
+ self.proj_matrix[1, 1] = 1 / self.tan_half_fov
+ self.proj_matrix[2, 2] = (opt.zfar + opt.znear) / (opt.zfar - opt.znear)
+ self.proj_matrix[3, 2] = - (opt.zfar * opt.znear) / (opt.zfar - opt.znear)
+ self.proj_matrix[2, 3] = 1
+
+ def render(self, gaussians, cam_view, cam_view_proj, cam_pos, bg_color=None, scale_modifier=1):
+ # gaussians: [B, N, 14]
+ # cam_view, cam_view_proj: [B, V, 4, 4]
+ # cam_pos: [B, V, 3]
+
+ device = gaussians.device
+ B, V = cam_view.shape[:2]
+
+ # loop of loop...
+ images = []
+ alphas = []
+ for b in range(B):
+
+ # pos, opacity, scale, rotation, shs
+ means3D = gaussians[b, :, 0:3].contiguous().float()
+ opacity = gaussians[b, :, 3:4].contiguous().float()
+ scales = gaussians[b, :, 4:7].contiguous().float()
+ rotations = gaussians[b, :, 7:11].contiguous().float()
+ rgbs = gaussians[b, :, 11:].contiguous().float() # [N, 3]
+
+ for v in range(V):
+
+ # render novel views
+ view_matrix = cam_view[b, v].float()
+ view_proj_matrix = cam_view_proj[b, v].float()
+ campos = cam_pos[b, v].float()
+
+ raster_settings = GaussianRasterizationSettings(
+ image_height=self.opt.output_size,
+ image_width=self.opt.output_size,
+ tanfovx=self.tan_half_fov,
+ tanfovy=self.tan_half_fov,
+ bg=self.bg_color if bg_color is None else bg_color,
+ scale_modifier=scale_modifier,
+ viewmatrix=view_matrix,
+ projmatrix=view_proj_matrix,
+ sh_degree=0,
+ campos=campos,
+ prefiltered=False,
+ debug=False,
+ )
+
+ rasterizer = GaussianRasterizer(raster_settings=raster_settings)
+
+ # Rasterize visible Gaussians to image, obtain their radii (on screen).
+ rendered_image, radii, rendered_depth, rendered_alpha = rasterizer(
+ means3D=means3D,
+ means2D=torch.zeros_like(means3D, dtype=torch.float32, device=device),
+ shs=None,
+ colors_precomp=rgbs,
+ opacities=opacity,
+ scales=scales,
+ rotations=rotations,
+ cov3D_precomp=None,
+ )
+
+ rendered_image = rendered_image.clamp(0, 1)
+
+ images.append(rendered_image)
+ alphas.append(rendered_alpha)
+
+ images = torch.stack(images, dim=0).view(B, V, 3, self.opt.output_size, self.opt.output_size)
+ alphas = torch.stack(alphas, dim=0).view(B, V, 1, self.opt.output_size, self.opt.output_size)
+
+ return {
+ "image": images, # [B, V, 3, H, W]
+ "alpha": alphas, # [B, V, 1, H, W]
+ }
+
+
+ def to_ply(self, gaussians, compatible=True):
+ # gaussians: [B, N, 14]
+ # compatible: save pre-activated gaussians as in the original paper
+
+ assert gaussians.shape[0] == 1, 'only support batch size 1'
+
+ from plyfile import PlyData, PlyElement
+
+ means3D = gaussians[0, :, 0:3].contiguous().float()
+ opacity = gaussians[0, :, 3:4].contiguous().float()
+ scales = gaussians[0, :, 4:7].contiguous().float()
+ rotations = gaussians[0, :, 7:11].contiguous().float()
+ shs = gaussians[0, :, 11:].unsqueeze(1).contiguous().float() # [N, 1, 3]
+
+ # prune by opacity
+ mask = opacity.squeeze(-1) >= 0.005
+ means3D = means3D[mask]
+ opacity = opacity[mask]
+ scales = scales[mask]
+ rotations = rotations[mask]
+ shs = shs[mask]
+
+ # invert activation to make it compatible with the original ply format
+ if compatible:
+ opacity = kiui.op.inverse_sigmoid(opacity)
+ scales = torch.log(scales + 1e-8)
+ shs = (shs - 0.5) / 0.28209479177387814
+
+ xyzs = means3D.detach().cpu().numpy()
+ f_dc = shs.detach().transpose(1, 2).flatten(start_dim=1).contiguous().cpu().numpy()
+ opacities = opacity.detach().cpu().numpy()
+ scales = scales.detach().cpu().numpy()
+ rotations = rotations.detach().cpu().numpy()
+
+ l = ['x', 'y', 'z']
+ # All channels except the 3 DC
+ for i in range(f_dc.shape[1]):
+ l.append('f_dc_{}'.format(i))
+ l.append('opacity')
+ for i in range(scales.shape[1]):
+ l.append('scale_{}'.format(i))
+ for i in range(rotations.shape[1]):
+ l.append('rot_{}'.format(i))
+
+ dtype_full = [(attribute, 'f4') for attribute in l]
+
+ elements = np.empty(xyzs.shape[0], dtype=dtype_full)
+ attributes = np.concatenate((xyzs, f_dc, opacities, scales, rotations), axis=1)
+ elements[:] = list(map(tuple, attributes))
+ el = PlyElement.describe(elements, 'vertex')
+
+ return PlyData([el])
+
+ def create_from_ply(self, plydata, compatible=True):
+
+ xyz = np.stack((np.asarray(plydata.elements[0]["x"]),
+ np.asarray(plydata.elements[0]["y"]),
+ np.asarray(plydata.elements[0]["z"])), axis=1)
+ print("Number of points at loading : ", xyz.shape[0])
+
+ opacities = np.asarray(plydata.elements[0]["opacity"])[..., np.newaxis]
+
+ shs = np.zeros((xyz.shape[0], 3))
+ shs[:, 0] = np.asarray(plydata.elements[0]["f_dc_0"])
+ shs[:, 1] = np.asarray(plydata.elements[0]["f_dc_1"])
+ shs[:, 2] = np.asarray(plydata.elements[0]["f_dc_2"])
+
+ scale_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("scale_")]
+ scales = np.zeros((xyz.shape[0], len(scale_names)))
+ for idx, attr_name in enumerate(scale_names):
+ scales[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+ rot_names = [p.name for p in plydata.elements[0].properties if p.name.startswith("rot_")]
+ rots = np.zeros((xyz.shape[0], len(rot_names)))
+ for idx, attr_name in enumerate(rot_names):
+ rots[:, idx] = np.asarray(plydata.elements[0][attr_name])
+
+ gaussians = np.concatenate([xyz, opacities, scales, rots, shs], axis=1)
+ gaussians = torch.from_numpy(gaussians).float() # cpu
+
+ if compatible:
+ gaussians[..., 3:4] = torch.sigmoid(gaussians[..., 3:4])
+ gaussians[..., 4:7] = torch.exp(gaussians[..., 4:7])
+ gaussians[..., 11:] = 0.28209479177387814 * gaussians[..., 11:] + 0.5
+
+ return gaussians
\ No newline at end of file
diff --git a/lgm/core/options.py b/lgm/core/options.py
new file mode 100644
index 0000000..ba6e8f1
--- /dev/null
+++ b/lgm/core/options.py
@@ -0,0 +1,117 @@
+from dataclasses import dataclass
+from typing import Tuple, Literal, Dict, Optional
+
+
+@dataclass
+class Options:
+ ### model
+ # Unet image input size
+ input_size: int = 256
+ # Unet definition
+ down_channels: Tuple[int, ...] = (64, 128, 256, 512, 1024, 1024)
+ down_attention: Tuple[bool, ...] = (False, False, False, True, True, True)
+ mid_attention: bool = True
+ up_channels: Tuple[int, ...] = (1024, 1024, 512, 256)
+ up_attention: Tuple[bool, ...] = (True, True, True, False)
+ # Unet output size, dependent on the input_size and U-Net structure!
+ splat_size: int = 64
+ # gaussian render size
+ output_size: int = 256
+
+ ### dataset
+ # data mode (only support s3 now)
+ data_mode: Literal['s3'] = 's3'
+ # fovy of the dataset
+ fovy: float = 49.1
+ # camera near plane
+ znear: float = 0.5
+ # camera far plane
+ zfar: float = 2.5
+ # number of all views (input + output)
+ num_views: int = 12
+ # number of views
+ num_input_views: int = 4
+ # camera radius
+ cam_radius: float = 1.5 # to better use [-1, 1]^3 space
+ # num workers
+ num_workers: int = 8
+
+ ### training
+ # workspace
+ workspace: str = './workspace'
+ # resume
+ resume: Optional[str] = None
+ # batch size (per-GPU)
+ batch_size: int = 8
+ # gradient accumulation
+ gradient_accumulation_steps: int = 1
+ # training epochs
+ num_epochs: int = 30
+ # lpips loss weight
+ lambda_lpips: float = 1.0
+ # gradient clip
+ gradient_clip: float = 1.0
+ # mixed precision
+ mixed_precision: str = 'bf16'
+ # learning rate
+ lr: float = 4e-4
+ # augmentation prob for grid distortion
+ prob_grid_distortion: float = 0.5
+ # augmentation prob for camera jitter
+ prob_cam_jitter: float = 0.5
+
+ ### testing
+ # test image path
+ test_path: Optional[str] = None
+
+ ### misc
+ # nvdiffrast backend setting
+ force_cuda_rast: bool = False
+ # render fancy video with gaussian scaling effect
+ fancy_video: bool = False
+
+
+# all the default settings
+config_defaults: Dict[str, Options] = {}
+config_doc: Dict[str, str] = {}
+
+config_doc['default'] = 'the default settings for LGM'
+config_defaults['default'] = Options()
+
+config_doc['small'] = 'small model with lower resolution Gaussians'
+config_defaults['small'] = Options(
+ input_size=256,
+ splat_size=64,
+ output_size=256,
+ batch_size=8,
+ gradient_accumulation_steps=1,
+ mixed_precision='bf16',
+)
+
+config_doc['big'] = 'big model with higher resolution Gaussians'
+config_defaults['big'] = Options(
+ input_size=256,
+ up_channels=(1024, 1024, 512, 256, 128), # one more decoder
+ up_attention=(True, True, True, False, False),
+ splat_size=128,
+ output_size=512, # render & supervise Gaussians at a higher resolution.
+ batch_size=8,
+ num_views=8,
+ gradient_accumulation_steps=1,
+ mixed_precision='bf16',
+)
+
+config_doc['tiny'] = 'tiny model for ablation'
+config_defaults['tiny'] = Options(
+ input_size=256,
+ down_channels=(32, 64, 128, 256, 512),
+ down_attention=(False, False, False, False, True),
+ up_channels=(512, 256, 128),
+ up_attention=(True, False, False, False),
+ splat_size=64,
+ output_size=256,
+ batch_size=16,
+ num_views=8,
+ gradient_accumulation_steps=1,
+ mixed_precision='bf16',
+)
diff --git a/lgm/core/provider_objaverse.py b/lgm/core/provider_objaverse.py
new file mode 100644
index 0000000..d73f646
--- /dev/null
+++ b/lgm/core/provider_objaverse.py
@@ -0,0 +1,172 @@
+import os
+import cv2
+import random
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torchvision.transforms.functional as TF
+from torch.utils.data import Dataset
+
+import kiui
+from lgm.core.options import Options
+from lgm.core.utils import get_rays, grid_distortion, orbit_camera_jitter
+
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+
+class ObjaverseDataset(Dataset):
+
+ def _warn(self):
+ raise NotImplementedError('this dataset is just an example and cannot be used directly, you should modify it to your own setting! (search keyword TODO)')
+
+ def __init__(self, opt: Options, training=True):
+
+ self.opt = opt
+ self.training = training
+
+ # TODO: remove this barrier
+ self._warn()
+
+ # TODO: load the list of objects for training
+ self.items = []
+ with open('TODO: file containing the list', 'r') as f:
+ for line in f.readlines():
+ self.items.append(line.strip())
+
+ # naive split
+ if self.training:
+ self.items = self.items[:-self.opt.batch_size]
+ else:
+ self.items = self.items[-self.opt.batch_size:]
+
+ # default camera intrinsics
+ self.tan_half_fov = np.tan(0.5 * np.deg2rad(self.opt.fovy))
+ self.proj_matrix = torch.zeros(4, 4, dtype=torch.float32)
+ self.proj_matrix[0, 0] = 1 / self.tan_half_fov
+ self.proj_matrix[1, 1] = 1 / self.tan_half_fov
+ self.proj_matrix[2, 2] = (self.opt.zfar + self.opt.znear) / (self.opt.zfar - self.opt.znear)
+ self.proj_matrix[3, 2] = - (self.opt.zfar * self.opt.znear) / (self.opt.zfar - self.opt.znear)
+ self.proj_matrix[2, 3] = 1
+
+
+ def __len__(self):
+ return len(self.items)
+
+ def __getitem__(self, idx):
+
+ uid = self.items[idx]
+ results = {}
+
+ # load num_views images
+ images = []
+ masks = []
+ cam_poses = []
+
+ vid_cnt = 0
+
+ # TODO: choose views, based on your rendering settings
+ if self.training:
+ # input views are in (36, 72), other views are randomly selected
+ vids = np.random.permutation(np.arange(36, 73))[:self.opt.num_input_views].tolist() + np.random.permutation(100).tolist()
+ else:
+ # fixed views
+ vids = np.arange(36, 73, 4).tolist() + np.arange(100).tolist()
+
+ for vid in vids:
+
+ image_path = os.path.join(uid, 'rgb', f'{vid:03d}.png')
+ camera_path = os.path.join(uid, 'pose', f'{vid:03d}.txt')
+
+ try:
+ # TODO: load data (modify self.client here)
+ image = np.frombuffer(self.client.get(image_path), np.uint8)
+ image = torch.from_numpy(cv2.imdecode(image, cv2.IMREAD_UNCHANGED).astype(np.float32) / 255) # [512, 512, 4] in [0, 1]
+ c2w = [float(t) for t in self.client.get(camera_path).decode().strip().split(' ')]
+ c2w = torch.tensor(c2w, dtype=torch.float32).reshape(4, 4)
+ except Exception as e:
+ # print(f'[WARN] dataset {uid} {vid}: {e}')
+ continue
+
+ # TODO: you may have a different camera system
+ # blender world + opencv cam --> opengl world & cam
+ c2w[1] *= -1
+ c2w[[1, 2]] = c2w[[2, 1]]
+ c2w[:3, 1:3] *= -1 # invert up and forward direction
+
+ # scale up radius to fully use the [-1, 1]^3 space!
+ c2w[:3, 3] *= self.opt.cam_radius / 1.5 # 1.5 is the default scale
+
+ image = image.permute(2, 0, 1) # [4, 512, 512]
+ mask = image[3:4] # [1, 512, 512]
+ image = image[:3] * mask + (1 - mask) # [3, 512, 512], to white bg
+ image = image[[2,1,0]].contiguous() # bgr to rgb
+
+ images.append(image)
+ masks.append(mask.squeeze(0))
+ cam_poses.append(c2w)
+
+ vid_cnt += 1
+ if vid_cnt == self.opt.num_views:
+ break
+
+ if vid_cnt < self.opt.num_views:
+ print(f'[WARN] dataset {uid}: not enough valid views, only {vid_cnt} views found!')
+ n = self.opt.num_views - vid_cnt
+ images = images + [images[-1]] * n
+ masks = masks + [masks[-1]] * n
+ cam_poses = cam_poses + [cam_poses[-1]] * n
+
+ images = torch.stack(images, dim=0) # [V, C, H, W]
+ masks = torch.stack(masks, dim=0) # [V, H, W]
+ cam_poses = torch.stack(cam_poses, dim=0) # [V, 4, 4]
+
+ # normalized camera feats as in paper (transform the first pose to a fixed position)
+ transform = torch.tensor([[1, 0, 0, 0], [0, 1, 0, 0], [0, 0, 1, self.opt.cam_radius], [0, 0, 0, 1]], dtype=torch.float32) @ torch.inverse(cam_poses[0])
+ cam_poses = transform.unsqueeze(0) @ cam_poses # [V, 4, 4]
+
+ images_input = F.interpolate(images[:self.opt.num_input_views].clone(), size=(self.opt.input_size, self.opt.input_size), mode='bilinear', align_corners=False) # [V, C, H, W]
+ cam_poses_input = cam_poses[:self.opt.num_input_views].clone()
+
+ # data augmentation
+ if self.training:
+ # apply random grid distortion to simulate 3D inconsistency
+ if random.random() < self.opt.prob_grid_distortion:
+ images_input[1:] = grid_distortion(images_input[1:])
+ # apply camera jittering (only to input!)
+ if random.random() < self.opt.prob_cam_jitter:
+ cam_poses_input[1:] = orbit_camera_jitter(cam_poses_input[1:])
+
+ images_input = TF.normalize(images_input, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
+
+ # resize render ground-truth images, range still in [0, 1]
+ results['images_output'] = F.interpolate(images, size=(self.opt.output_size, self.opt.output_size), mode='bilinear', align_corners=False) # [V, C, output_size, output_size]
+ results['masks_output'] = F.interpolate(masks.unsqueeze(1), size=(self.opt.output_size, self.opt.output_size), mode='bilinear', align_corners=False) # [V, 1, output_size, output_size]
+
+ # build rays for input views
+ rays_embeddings = []
+ for i in range(self.opt.num_input_views):
+ rays_o, rays_d = get_rays(cam_poses_input[i], self.opt.input_size, self.opt.input_size, self.opt.fovy) # [h, w, 3]
+ rays_plucker = torch.cat([torch.cross(rays_o, rays_d, dim=-1), rays_d], dim=-1) # [h, w, 6]
+ rays_embeddings.append(rays_plucker)
+
+
+ rays_embeddings = torch.stack(rays_embeddings, dim=0).permute(0, 3, 1, 2).contiguous() # [V, 6, h, w]
+ final_input = torch.cat([images_input, rays_embeddings], dim=1) # [V=4, 9, H, W]
+ results['input'] = final_input
+
+ # opengl to colmap camera for gaussian renderer
+ cam_poses[:, :3, 1:3] *= -1 # invert up & forward direction
+
+ # cameras needed by gaussian rasterizer
+ cam_view = torch.inverse(cam_poses).transpose(1, 2) # [V, 4, 4]
+ cam_view_proj = cam_view @ self.proj_matrix # [V, 4, 4]
+ cam_pos = - cam_poses[:, :3, 3] # [V, 3]
+
+ results['cam_view'] = cam_view
+ results['cam_view_proj'] = cam_view_proj
+ results['cam_pos'] = cam_pos
+
+ return results
\ No newline at end of file
diff --git a/lgm/core/unet.py b/lgm/core/unet.py
new file mode 100644
index 0000000..dcf75cd
--- /dev/null
+++ b/lgm/core/unet.py
@@ -0,0 +1,319 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import numpy as np
+from typing import Tuple, Literal
+from functools import partial
+
+from lgm.core.attention import MemEffAttention
+
+class MVAttention(nn.Module):
+ def __init__(
+ self,
+ dim: int,
+ num_heads: int = 8,
+ qkv_bias: bool = False,
+ proj_bias: bool = True,
+ attn_drop: float = 0.0,
+ proj_drop: float = 0.0,
+ groups: int = 32,
+ eps: float = 1e-5,
+ residual: bool = True,
+ skip_scale: float = 1,
+ num_frames: int = 4, # WARN: hardcoded!
+ ):
+ super().__init__()
+
+ self.residual = residual
+ self.skip_scale = skip_scale
+ self.num_frames = num_frames
+
+ self.norm = nn.GroupNorm(num_groups=groups, num_channels=dim, eps=eps, affine=True)
+ self.attn = MemEffAttention(dim, num_heads, qkv_bias, proj_bias, attn_drop, proj_drop)
+
+ def forward(self, x):
+ # x: [B*V, C, H, W]
+ BV, C, H, W = x.shape
+ B = BV // self.num_frames # assert BV % self.num_frames == 0
+
+ res = x
+ x = self.norm(x)
+
+ x = x.reshape(B, self.num_frames, C, H, W).permute(0, 1, 3, 4, 2).reshape(B, -1, C)
+ x = self.attn(x)
+ x = x.reshape(B, self.num_frames, H, W, C).permute(0, 1, 4, 2, 3).reshape(BV, C, H, W)
+
+ if self.residual:
+ x = (x + res) * self.skip_scale
+ return x
+
+class ResnetBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ resample: Literal['default', 'up', 'down'] = 'default',
+ groups: int = 32,
+ eps: float = 1e-5,
+ skip_scale: float = 1, # multiplied to output
+ ):
+ super().__init__()
+
+ self.in_channels = in_channels
+ self.out_channels = out_channels
+ self.skip_scale = skip_scale
+
+ self.norm1 = nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)
+ self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+ self.norm2 = nn.GroupNorm(num_groups=groups, num_channels=out_channels, eps=eps, affine=True)
+ self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+ self.act = F.silu
+
+ self.resample = None
+ if resample == 'up':
+ self.resample = partial(F.interpolate, scale_factor=2.0, mode="nearest")
+ elif resample == 'down':
+ self.resample = nn.AvgPool2d(kernel_size=2, stride=2)
+
+ self.shortcut = nn.Identity()
+ if self.in_channels != self.out_channels:
+ self.shortcut = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=True)
+
+
+ def forward(self, x):
+ res = x
+
+ x = self.norm1(x)
+ x = self.act(x)
+
+ if self.resample:
+ res = self.resample(res)
+ x = self.resample(x)
+
+ x = self.conv1(x)
+ x = self.norm2(x)
+ x = self.act(x)
+ x = self.conv2(x)
+
+ x = (x + self.shortcut(res)) * self.skip_scale
+
+ return x
+
+class DownBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ out_channels: int,
+ num_layers: int = 1,
+ downsample: bool = True,
+ attention: bool = True,
+ attention_heads: int = 16,
+ skip_scale: float = 1,
+ ):
+ super().__init__()
+
+ nets = []
+ attns = []
+ for i in range(num_layers):
+ in_channels = in_channels if i == 0 else out_channels
+ nets.append(ResnetBlock(in_channels, out_channels, skip_scale=skip_scale))
+ if attention:
+ attns.append(MVAttention(out_channels, attention_heads, skip_scale=skip_scale))
+ else:
+ attns.append(None)
+ self.nets = nn.ModuleList(nets)
+ self.attns = nn.ModuleList(attns)
+
+ self.downsample = None
+ if downsample:
+ self.downsample = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=2, padding=1)
+
+ def forward(self, x):
+ xs = []
+
+ for attn, net in zip(self.attns, self.nets):
+ x = net(x)
+ if attn:
+ x = attn(x)
+ xs.append(x)
+
+ if self.downsample:
+ x = self.downsample(x)
+ xs.append(x)
+
+ return x, xs
+
+
+class MidBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ num_layers: int = 1,
+ attention: bool = True,
+ attention_heads: int = 16,
+ skip_scale: float = 1,
+ ):
+ super().__init__()
+
+ nets = []
+ attns = []
+ # first layer
+ nets.append(ResnetBlock(in_channels, in_channels, skip_scale=skip_scale))
+ # more layers
+ for i in range(num_layers):
+ nets.append(ResnetBlock(in_channels, in_channels, skip_scale=skip_scale))
+ if attention:
+ attns.append(MVAttention(in_channels, attention_heads, skip_scale=skip_scale))
+ else:
+ attns.append(None)
+ self.nets = nn.ModuleList(nets)
+ self.attns = nn.ModuleList(attns)
+
+ def forward(self, x):
+ x = self.nets[0](x)
+ for attn, net in zip(self.attns, self.nets[1:]):
+ if attn:
+ x = attn(x)
+ x = net(x)
+ return x
+
+
+class UpBlock(nn.Module):
+ def __init__(
+ self,
+ in_channels: int,
+ prev_out_channels: int,
+ out_channels: int,
+ num_layers: int = 1,
+ upsample: bool = True,
+ attention: bool = True,
+ attention_heads: int = 16,
+ skip_scale: float = 1,
+ ):
+ super().__init__()
+
+ nets = []
+ attns = []
+ for i in range(num_layers):
+ cin = in_channels if i == 0 else out_channels
+ cskip = prev_out_channels if (i == num_layers - 1) else out_channels
+
+ nets.append(ResnetBlock(cin + cskip, out_channels, skip_scale=skip_scale))
+ if attention:
+ attns.append(MVAttention(out_channels, attention_heads, skip_scale=skip_scale))
+ else:
+ attns.append(None)
+ self.nets = nn.ModuleList(nets)
+ self.attns = nn.ModuleList(attns)
+
+ self.upsample = None
+ if upsample:
+ self.upsample = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1)
+
+ def forward(self, x, xs):
+
+ for attn, net in zip(self.attns, self.nets):
+ res_x = xs[-1]
+ xs = xs[:-1]
+ x = torch.cat([x, res_x], dim=1)
+ x = net(x)
+ if attn:
+ x = attn(x)
+
+ if self.upsample:
+ x = F.interpolate(x, scale_factor=2.0, mode='nearest')
+ x = self.upsample(x)
+
+ return x
+
+
+# it could be asymmetric!
+class UNet(nn.Module):
+ def __init__(
+ self,
+ in_channels: int = 3,
+ out_channels: int = 3,
+ down_channels: Tuple[int, ...] = (64, 128, 256, 512, 1024),
+ down_attention: Tuple[bool, ...] = (False, False, False, True, True),
+ mid_attention: bool = True,
+ up_channels: Tuple[int, ...] = (1024, 512, 256),
+ up_attention: Tuple[bool, ...] = (True, True, False),
+ layers_per_block: int = 2,
+ skip_scale: float = np.sqrt(0.5),
+ ):
+ super().__init__()
+
+ # first
+ self.conv_in = nn.Conv2d(in_channels, down_channels[0], kernel_size=3, stride=1, padding=1)
+
+ # down
+ down_blocks = []
+ cout = down_channels[0]
+ for i in range(len(down_channels)):
+ cin = cout
+ cout = down_channels[i]
+
+ down_blocks.append(DownBlock(
+ cin, cout,
+ num_layers=layers_per_block,
+ downsample=(i != len(down_channels) - 1), # not final layer
+ attention=down_attention[i],
+ skip_scale=skip_scale,
+ ))
+ self.down_blocks = nn.ModuleList(down_blocks)
+
+ # mid
+ self.mid_block = MidBlock(down_channels[-1], attention=mid_attention, skip_scale=skip_scale)
+
+ # up
+ up_blocks = []
+ cout = up_channels[0]
+ for i in range(len(up_channels)):
+ cin = cout
+ cout = up_channels[i]
+ cskip = down_channels[max(-2 - i, -len(down_channels))] # for assymetric
+
+ up_blocks.append(UpBlock(
+ cin, cskip, cout,
+ num_layers=layers_per_block + 1, # one more layer for up
+ upsample=(i != len(up_channels) - 1), # not final layer
+ attention=up_attention[i],
+ skip_scale=skip_scale,
+ ))
+ self.up_blocks = nn.ModuleList(up_blocks)
+
+ # last
+ self.norm_out = nn.GroupNorm(num_channels=up_channels[-1], num_groups=32, eps=1e-5)
+ self.conv_out = nn.Conv2d(up_channels[-1], out_channels, kernel_size=3, stride=1, padding=1)
+
+
+ def forward(self, x):
+ # x: [B, Cin, H, W]
+
+ # first
+ x = self.conv_in(x)
+
+ # down
+ xss = [x]
+ for block in self.down_blocks:
+ x, xs = block(x)
+ xss.extend(xs)
+
+ # mid
+ x = self.mid_block(x)
+
+ # up
+ for block in self.up_blocks:
+ xs = xss[-len(block.nets):]
+ xss = xss[:-len(block.nets)]
+ x = block(x, xs)
+
+ # last
+ x = self.norm_out(x)
+ x = F.silu(x)
+ x = self.conv_out(x) # [B, Cout, H', W']
+
+ return x
diff --git a/lgm/core/utils.py b/lgm/core/utils.py
new file mode 100644
index 0000000..17ab4f5
--- /dev/null
+++ b/lgm/core/utils.py
@@ -0,0 +1,109 @@
+import numpy as np
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+import roma
+from kiui.op import safe_normalize
+
+def get_rays(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 orbit_camera_jitter(poses, strength=0.1):
+ # poses: [B, 4, 4], assume orbit camera in opengl format
+ # random orbital rotate
+
+ B = poses.shape[0]
+ rotvec_x = poses[:, :3, 1] * strength * np.pi * (torch.rand(B, 1, device=poses.device) * 2 - 1)
+ rotvec_y = poses[:, :3, 0] * strength * np.pi / 2 * (torch.rand(B, 1, device=poses.device) * 2 - 1)
+
+ rot = roma.rotvec_to_rotmat(rotvec_x) @ roma.rotvec_to_rotmat(rotvec_y)
+ R = rot @ poses[:, :3, :3]
+ T = rot @ poses[:, :3, 3:]
+
+ new_poses = poses.clone()
+ new_poses[:, :3, :3] = R
+ new_poses[:, :3, 3:] = T
+
+ return new_poses
+
+def grid_distortion(images, strength=0.5):
+ # images: [B, C, H, W]
+ # num_steps: int, grid resolution for distortion
+ # strength: float in [0, 1], strength of distortion
+
+ B, C, H, W = images.shape
+
+ num_steps = np.random.randint(8, 17)
+ grid_steps = torch.linspace(-1, 1, num_steps)
+
+ # have to loop batch...
+ grids = []
+ for b in range(B):
+ # construct displacement
+ x_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive
+ x_steps = (x_steps + strength * (torch.rand_like(x_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb
+ x_steps = (x_steps * W).long() # [num_steps]
+ x_steps[0] = 0
+ x_steps[-1] = W
+ xs = []
+ for i in range(num_steps - 1):
+ xs.append(torch.linspace(grid_steps[i], grid_steps[i + 1], x_steps[i + 1] - x_steps[i]))
+ xs = torch.cat(xs, dim=0) # [W]
+
+ y_steps = torch.linspace(0, 1, num_steps) # [num_steps], inclusive
+ y_steps = (y_steps + strength * (torch.rand_like(y_steps) - 0.5) / (num_steps - 1)).clamp(0, 1) # perturb
+ y_steps = (y_steps * H).long() # [num_steps]
+ y_steps[0] = 0
+ y_steps[-1] = H
+ ys = []
+ for i in range(num_steps - 1):
+ ys.append(torch.linspace(grid_steps[i], grid_steps[i + 1], y_steps[i + 1] - y_steps[i]))
+ ys = torch.cat(ys, dim=0) # [H]
+
+ # construct grid
+ grid_x, grid_y = torch.meshgrid(xs, ys, indexing='xy') # [H, W]
+ grid = torch.stack([grid_x, grid_y], dim=-1) # [H, W, 2]
+
+ grids.append(grid)
+
+ grids = torch.stack(grids, dim=0).to(images.device) # [B, H, W, 2]
+
+ # grid sample
+ images = F.grid_sample(images, grids, align_corners=False)
+
+ return images
+
diff --git a/lgm/mvdream/mv_unet.py b/lgm/mvdream/mv_unet.py
new file mode 100644
index 0000000..7d9ad4d
--- /dev/null
+++ b/lgm/mvdream/mv_unet.py
@@ -0,0 +1,1005 @@
+import math
+import numpy as np
+from inspect import isfunction
+from typing import Optional, Any, List
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from diffusers.configuration_utils import ConfigMixin
+from diffusers.models.modeling_utils import ModelMixin
+
+# require xformers!
+import xformers
+import xformers.ops
+
+from kiui.cam import orbit_camera
+
+def get_camera(
+ num_frames, elevation=0, azimuth_start=0, azimuth_span=360, blender_coord=True, extra_view=False,
+):
+ angle_gap = azimuth_span / num_frames
+ cameras = []
+ for azimuth in np.arange(azimuth_start, azimuth_span + azimuth_start, angle_gap):
+
+ pose = orbit_camera(elevation, azimuth, radius=1) # [4, 4]
+
+ # opengl to blender
+ if blender_coord:
+ pose[2] *= -1
+ pose[[1, 2]] = pose[[2, 1]]
+
+ cameras.append(pose.flatten())
+
+ if extra_view:
+ cameras.append(np.zeros_like(cameras[0]))
+
+ return torch.from_numpy(np.stack(cameras, axis=0)).float() # [num_frames, 16]
+
+
+def timestep_embedding(timesteps, dim, max_period=10000, repeat_only=False):
+ """
+ Create sinusoidal timestep embeddings.
+ :param timesteps: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an [N x dim] Tensor of positional embeddings.
+ """
+ if not repeat_only:
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period)
+ * torch.arange(start=0, end=half, dtype=torch.float32)
+ / half
+ ).to(device=timesteps.device)
+ args = timesteps[:, None] * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat(
+ [embedding, torch.zeros_like(embedding[:, :1])], dim=-1
+ )
+ else:
+ embedding = repeat(timesteps, "b -> b d", d=dim)
+ # import pdb; pdb.set_trace()
+ return embedding
+
+
+def zero_module(module):
+ """
+ Zero out the parameters of a module and return it.
+ """
+ for p in module.parameters():
+ p.detach().zero_()
+ return module
+
+
+def conv_nd(dims, *args, **kwargs):
+ """
+ Create a 1D, 2D, or 3D convolution module.
+ """
+ if dims == 1:
+ return nn.Conv1d(*args, **kwargs)
+ elif dims == 2:
+ return nn.Conv2d(*args, **kwargs)
+ elif dims == 3:
+ return nn.Conv3d(*args, **kwargs)
+ raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def avg_pool_nd(dims, *args, **kwargs):
+ """
+ Create a 1D, 2D, or 3D average pooling module.
+ """
+ if dims == 1:
+ return nn.AvgPool1d(*args, **kwargs)
+ elif dims == 2:
+ return nn.AvgPool2d(*args, **kwargs)
+ elif dims == 3:
+ return nn.AvgPool3d(*args, **kwargs)
+ raise ValueError(f"unsupported dimensions: {dims}")
+
+
+def default(val, d):
+ if val is not None:
+ return val
+ return d() if isfunction(d) else d
+
+
+class GEGLU(nn.Module):
+ def __init__(self, dim_in, dim_out):
+ super().__init__()
+ self.proj = nn.Linear(dim_in, dim_out * 2)
+
+ def forward(self, x):
+ x, gate = self.proj(x).chunk(2, dim=-1)
+ return x * F.gelu(gate)
+
+
+class FeedForward(nn.Module):
+ def __init__(self, dim, dim_out=None, mult=4, glu=False, dropout=0.0):
+ super().__init__()
+ inner_dim = int(dim * mult)
+ dim_out = default(dim_out, dim)
+ project_in = (
+ nn.Sequential(nn.Linear(dim, inner_dim), nn.GELU())
+ if not glu
+ else GEGLU(dim, inner_dim)
+ )
+
+ self.net = nn.Sequential(
+ project_in, nn.Dropout(dropout), nn.Linear(inner_dim, dim_out)
+ )
+
+ def forward(self, x):
+ return self.net(x)
+
+
+class MemoryEfficientCrossAttention(nn.Module):
+ # https://github.com/MatthieuTPHR/diffusers/blob/d80b531ff8060ec1ea982b65a1b8df70f73aa67c/src/diffusers/models/attention.py#L223
+ def __init__(
+ self,
+ query_dim,
+ context_dim=None,
+ heads=8,
+ dim_head=64,
+ dropout=0.0,
+ ip_dim=0,
+ ip_weight=1,
+ ):
+ super().__init__()
+
+ inner_dim = dim_head * heads
+ context_dim = default(context_dim, query_dim)
+
+ self.heads = heads
+ self.dim_head = dim_head
+
+ self.ip_dim = ip_dim
+ self.ip_weight = ip_weight
+
+ if self.ip_dim > 0:
+ self.to_k_ip = nn.Linear(context_dim, inner_dim, bias=False)
+ self.to_v_ip = nn.Linear(context_dim, inner_dim, bias=False)
+
+ self.to_q = nn.Linear(query_dim, inner_dim, bias=False)
+ self.to_k = nn.Linear(context_dim, inner_dim, bias=False)
+ self.to_v = nn.Linear(context_dim, inner_dim, bias=False)
+
+ self.to_out = nn.Sequential(
+ nn.Linear(inner_dim, query_dim), nn.Dropout(dropout)
+ )
+ self.attention_op: Optional[Any] = None
+
+ def forward(self, x, context=None):
+ q = self.to_q(x)
+ context = default(context, x)
+
+ if self.ip_dim > 0:
+ # context： [B, 77 + 16(ip), 1024]
+ token_len = context.shape[1]
+ context_ip = context[:, -self.ip_dim :, :]
+ k_ip = self.to_k_ip(context_ip)
+ v_ip = self.to_v_ip(context_ip)
+ context = context[:, : (token_len - self.ip_dim), :]
+
+ k = self.to_k(context)
+ v = self.to_v(context)
+
+ b, _, _ = q.shape
+ q, k, v = map(
+ lambda t: t.unsqueeze(3)
+ .reshape(b, t.shape[1], self.heads, self.dim_head)
+ .permute(0, 2, 1, 3)
+ .reshape(b * self.heads, t.shape[1], self.dim_head)
+ .contiguous(),
+ (q, k, v),
+ )
+
+ # actually compute the attention, what we cannot get enough of
+ out = xformers.ops.memory_efficient_attention(
+ q, k, v, attn_bias=None, op=self.attention_op
+ )
+
+ if self.ip_dim > 0:
+ k_ip, v_ip = map(
+ lambda t: t.unsqueeze(3)
+ .reshape(b, t.shape[1], self.heads, self.dim_head)
+ .permute(0, 2, 1, 3)
+ .reshape(b * self.heads, t.shape[1], self.dim_head)
+ .contiguous(),
+ (k_ip, v_ip),
+ )
+ # actually compute the attention, what we cannot get enough of
+ out_ip = xformers.ops.memory_efficient_attention(
+ q, k_ip, v_ip, attn_bias=None, op=self.attention_op
+ )
+ out = out + self.ip_weight * out_ip
+
+ out = (
+ out.unsqueeze(0)
+ .reshape(b, self.heads, out.shape[1], self.dim_head)
+ .permute(0, 2, 1, 3)
+ .reshape(b, out.shape[1], self.heads * self.dim_head)
+ )
+ return self.to_out(out)
+
+
+class BasicTransformerBlock3D(nn.Module):
+
+ def __init__(
+ self,
+ dim,
+ n_heads,
+ d_head,
+ context_dim,
+ dropout=0.0,
+ gated_ff=True,
+ ip_dim=0,
+ ip_weight=1,
+ ):
+ super().__init__()
+
+ self.attn1 = MemoryEfficientCrossAttention(
+ query_dim=dim,
+ context_dim=None, # self-attention
+ heads=n_heads,
+ dim_head=d_head,
+ dropout=dropout,
+ )
+ self.ff = FeedForward(dim, dropout=dropout, glu=gated_ff)
+ self.attn2 = MemoryEfficientCrossAttention(
+ query_dim=dim,
+ context_dim=context_dim,
+ heads=n_heads,
+ dim_head=d_head,
+ dropout=dropout,
+ # ip only applies to cross-attention
+ ip_dim=ip_dim,
+ ip_weight=ip_weight,
+ )
+ self.norm1 = nn.LayerNorm(dim)
+ self.norm2 = nn.LayerNorm(dim)
+ self.norm3 = nn.LayerNorm(dim)
+
+ def forward(self, x, context=None, num_frames=1):
+ x = rearrange(x, "(b f) l c -> b (f l) c", f=num_frames).contiguous()
+ x = self.attn1(self.norm1(x), context=None) + x
+ x = rearrange(x, "b (f l) c -> (b f) l c", f=num_frames).contiguous()
+ x = self.attn2(self.norm2(x), context=context) + x
+ x = self.ff(self.norm3(x)) + x
+ return x
+
+
+class SpatialTransformer3D(nn.Module):
+
+ def __init__(
+ self,
+ in_channels,
+ n_heads,
+ d_head,
+ context_dim, # cross attention input dim
+ depth=1,
+ dropout=0.0,
+ ip_dim=0,
+ ip_weight=1,
+ ):
+ super().__init__()
+
+ if not isinstance(context_dim, list):
+ context_dim = [context_dim]
+
+ self.in_channels = in_channels
+
+ inner_dim = n_heads * d_head
+ self.norm = nn.GroupNorm(num_groups=32, num_channels=in_channels, eps=1e-6, affine=True)
+ self.proj_in = nn.Linear(in_channels, inner_dim)
+
+ self.transformer_blocks = nn.ModuleList(
+ [
+ BasicTransformerBlock3D(
+ inner_dim,
+ n_heads,
+ d_head,
+ context_dim=context_dim[d],
+ dropout=dropout,
+ ip_dim=ip_dim,
+ ip_weight=ip_weight,
+ )
+ for d in range(depth)
+ ]
+ )
+
+ self.proj_out = zero_module(nn.Linear(in_channels, inner_dim))
+
+
+ def forward(self, x, context=None, num_frames=1):
+ # note: if no context is given, cross-attention defaults to self-attention
+ if not isinstance(context, list):
+ context = [context]
+ b, c, h, w = x.shape
+ x_in = x
+ x = self.norm(x)
+ x = rearrange(x, "b c h w -> b (h w) c").contiguous()
+ x = self.proj_in(x)
+ for i, block in enumerate(self.transformer_blocks):
+ x = block(x, context=context[i], num_frames=num_frames)
+ x = self.proj_out(x)
+ x = rearrange(x, "b (h w) c -> b c h w", h=h, w=w).contiguous()
+
+ return x + x_in
+
+
+class PerceiverAttention(nn.Module):
+ def __init__(self, *, dim, dim_head=64, heads=8):
+ super().__init__()
+ self.scale = dim_head ** -0.5
+ self.dim_head = dim_head
+ self.heads = heads
+ inner_dim = dim_head * heads
+
+ self.norm1 = nn.LayerNorm(dim)
+ self.norm2 = nn.LayerNorm(dim)
+
+ self.to_q = nn.Linear(dim, inner_dim, bias=False)
+ self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+ self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+ def forward(self, x, latents):
+ """
+ Args:
+ x (torch.Tensor): image features
+ shape (b, n1, D)
+ latent (torch.Tensor): latent features
+ shape (b, n2, D)
+ """
+ x = self.norm1(x)
+ latents = self.norm2(latents)
+
+ b, l, _ = latents.shape
+
+ q = self.to_q(latents)
+ kv_input = torch.cat((x, latents), dim=-2)
+ k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+
+ q, k, v = map(
+ lambda t: t.reshape(b, t.shape[1], self.heads, -1)
+ .transpose(1, 2)
+ .reshape(b, self.heads, t.shape[1], -1)
+ .contiguous(),
+ (q, k, v),
+ )
+
+ # attention
+ scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+ weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
+ weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+ out = weight @ v
+
+ out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+ return self.to_out(out)
+
+
+class Resampler(nn.Module):
+ def __init__(
+ self,
+ dim=1024,
+ depth=8,
+ dim_head=64,
+ heads=16,
+ num_queries=8,
+ embedding_dim=768,
+ output_dim=1024,
+ ff_mult=4,
+ ):
+ super().__init__()
+ self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim ** 0.5)
+ self.proj_in = nn.Linear(embedding_dim, dim)
+ self.proj_out = nn.Linear(dim, output_dim)
+ self.norm_out = nn.LayerNorm(output_dim)
+
+ self.layers = nn.ModuleList([])
+ for _ in range(depth):
+ self.layers.append(
+ nn.ModuleList(
+ [
+ PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+ nn.Sequential(
+ nn.LayerNorm(dim),
+ nn.Linear(dim, dim * ff_mult, bias=False),
+ nn.GELU(),
+ nn.Linear(dim * ff_mult, dim, bias=False),
+ )
+ ]
+ )
+ )
+
+ def forward(self, x):
+ latents = self.latents.repeat(x.size(0), 1, 1)
+ x = self.proj_in(x)
+ for attn, ff in self.layers:
+ latents = attn(x, latents) + latents
+ latents = ff(latents) + latents
+
+ latents = self.proj_out(latents)
+ return self.norm_out(latents)
+
+
+class CondSequential(nn.Sequential):
+ """
+ A sequential module that passes timestep embeddings to the children that
+ support it as an extra input.
+ """
+
+ def forward(self, x, emb, context=None, num_frames=1):
+ for layer in self:
+ if isinstance(layer, ResBlock):
+ x = layer(x, emb)
+ elif isinstance(layer, SpatialTransformer3D):
+ x = layer(x, context, num_frames=num_frames)
+ else:
+ x = layer(x)
+ return x
+
+
+class Upsample(nn.Module):
+ """
+ An upsampling layer with an optional convolution.
+ :param channels: channels in the inputs and outputs.
+ :param use_conv: a bool determining if a convolution is applied.
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+ upsampling occurs in the inner-two dimensions.
+ """
+
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.dims = dims
+ if use_conv:
+ self.conv = conv_nd(
+ dims, self.channels, self.out_channels, 3, padding=padding
+ )
+
+ def forward(self, x):
+ assert x.shape[1] == self.channels
+ if self.dims == 3:
+ x = F.interpolate(
+ x, (x.shape[2], x.shape[3] * 2, x.shape[4] * 2), mode="nearest"
+ )
+ else:
+ x = F.interpolate(x, scale_factor=2, mode="nearest")
+ if self.use_conv:
+ x = self.conv(x)
+ return x
+
+
+class Downsample(nn.Module):
+ """
+ A downsampling layer with an optional convolution.
+ :param channels: channels in the inputs and outputs.
+ :param use_conv: a bool determining if a convolution is applied.
+ :param dims: determines if the signal is 1D, 2D, or 3D. If 3D, then
+ downsampling occurs in the inner-two dimensions.
+ """
+
+ def __init__(self, channels, use_conv, dims=2, out_channels=None, padding=1):
+ super().__init__()
+ self.channels = channels
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.dims = dims
+ stride = 2 if dims != 3 else (1, 2, 2)
+ if use_conv:
+ self.op = conv_nd(
+ dims,
+ self.channels,
+ self.out_channels,
+ 3,
+ stride=stride,
+ padding=padding,
+ )
+ else:
+ assert self.channels == self.out_channels
+ self.op = avg_pool_nd(dims, kernel_size=stride, stride=stride)
+
+ def forward(self, x):
+ assert x.shape[1] == self.channels
+ return self.op(x)
+
+
+class ResBlock(nn.Module):
+ """
+ A residual block that can optionally change the number of channels.
+ :param channels: the number of input channels.
+ :param emb_channels: the number of timestep embedding channels.
+ :param dropout: the rate of dropout.
+ :param out_channels: if specified, the number of out channels.
+ :param use_conv: if True and out_channels is specified, use a spatial
+ convolution instead of a smaller 1x1 convolution to change the
+ channels in the skip connection.
+ :param dims: determines if the signal is 1D, 2D, or 3D.
+ :param up: if True, use this block for upsampling.
+ :param down: if True, use this block for downsampling.
+ """
+
+ def __init__(
+ self,
+ channels,
+ emb_channels,
+ dropout,
+ out_channels=None,
+ use_conv=False,
+ use_scale_shift_norm=False,
+ dims=2,
+ up=False,
+ down=False,
+ ):
+ super().__init__()
+ self.channels = channels
+ self.emb_channels = emb_channels
+ self.dropout = dropout
+ self.out_channels = out_channels or channels
+ self.use_conv = use_conv
+ self.use_scale_shift_norm = use_scale_shift_norm
+
+ self.in_layers = nn.Sequential(
+ nn.GroupNorm(32, channels),
+ nn.SiLU(),
+ conv_nd(dims, channels, self.out_channels, 3, padding=1),
+ )
+
+ self.updown = up or down
+
+ if up:
+ self.h_upd = Upsample(channels, False, dims)
+ self.x_upd = Upsample(channels, False, dims)
+ elif down:
+ self.h_upd = Downsample(channels, False, dims)
+ self.x_upd = Downsample(channels, False, dims)
+ else:
+ self.h_upd = self.x_upd = nn.Identity()
+
+ self.emb_layers = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(
+ emb_channels,
+ 2 * self.out_channels if use_scale_shift_norm else self.out_channels,
+ ),
+ )
+ self.out_layers = nn.Sequential(
+ nn.GroupNorm(32, self.out_channels),
+ nn.SiLU(),
+ nn.Dropout(p=dropout),
+ zero_module(
+ conv_nd(dims, self.out_channels, self.out_channels, 3, padding=1)
+ ),
+ )
+
+ if self.out_channels == channels:
+ self.skip_connection = nn.Identity()
+ elif use_conv:
+ self.skip_connection = conv_nd(
+ dims, channels, self.out_channels, 3, padding=1
+ )
+ else:
+ self.skip_connection = conv_nd(dims, channels, self.out_channels, 1)
+
+ def forward(self, x, emb):
+ if self.updown:
+ in_rest, in_conv = self.in_layers[:-1], self.in_layers[-1]
+ h = in_rest(x)
+ h = self.h_upd(h)
+ x = self.x_upd(x)
+ h = in_conv(h)
+ else:
+ h = self.in_layers(x)
+ emb_out = self.emb_layers(emb).type(h.dtype)
+ while len(emb_out.shape) < len(h.shape):
+ emb_out = emb_out[..., None]
+ if self.use_scale_shift_norm:
+ out_norm, out_rest = self.out_layers[0], self.out_layers[1:]
+ scale, shift = torch.chunk(emb_out, 2, dim=1)
+ h = out_norm(h) * (1 + scale) + shift
+ h = out_rest(h)
+ else:
+ h = h + emb_out
+ h = self.out_layers(h)
+ return self.skip_connection(x) + h
+
+
+class MultiViewUNetModel(ModelMixin, ConfigMixin):
+ """
+ The full multi-view UNet model with attention, timestep embedding and camera embedding.
+ :param in_channels: channels in the input Tensor.
+ :param model_channels: base channel count for the model.
+ :param out_channels: channels in the output Tensor.
+ :param num_res_blocks: number of residual blocks per downsample.
+ :param attention_resolutions: a collection of downsample rates at which
+ attention will take place. May be a set, list, or tuple.
+ For example, if this contains 4, then at 4x downsampling, attention
+ will be used.
+ :param dropout: the dropout probability.
+ :param channel_mult: channel multiplier for each level of the UNet.
+ :param conv_resample: if True, use learned convolutions for upsampling and
+ downsampling.
+ :param dims: determines if the signal is 1D, 2D, or 3D.
+ :param num_classes: if specified (as an int), then this model will be
+ class-conditional with `num_classes` classes.
+ :param num_heads: the number of attention heads in each attention layer.
+ :param num_heads_channels: if specified, ignore num_heads and instead use
+ a fixed channel width per attention head.
+ :param num_heads_upsample: works with num_heads to set a different number
+ of heads for upsampling. Deprecated.
+ :param use_scale_shift_norm: use a FiLM-like conditioning mechanism.
+ :param resblock_updown: use residual blocks for up/downsampling.
+ :param use_new_attention_order: use a different attention pattern for potentially
+ increased efficiency.
+ :param camera_dim: dimensionality of camera input.
+ """
+
+ def __init__(
+ self,
+ image_size,
+ in_channels,
+ model_channels,
+ out_channels,
+ num_res_blocks,
+ attention_resolutions,
+ dropout=0,
+ channel_mult=(1, 2, 4, 8),
+ conv_resample=True,
+ dims=2,
+ num_classes=None,
+ num_heads=-1,
+ num_head_channels=-1,
+ num_heads_upsample=-1,
+ use_scale_shift_norm=False,
+ resblock_updown=False,
+ transformer_depth=1,
+ context_dim=None,
+ n_embed=None,
+ num_attention_blocks=None,
+ adm_in_channels=None,
+ camera_dim=None,
+ ip_dim=0, # imagedream uses ip_dim > 0
+ ip_weight=1.0,
+ **kwargs,
+ ):
+ super().__init__()
+ assert context_dim is not None
+
+ if num_heads_upsample == -1:
+ num_heads_upsample = num_heads
+
+ if num_heads == -1:
+ assert (
+ num_head_channels != -1
+ ), "Either num_heads or num_head_channels has to be set"
+
+ if num_head_channels == -1:
+ assert (
+ num_heads != -1
+ ), "Either num_heads or num_head_channels has to be set"
+
+ self.image_size = image_size
+ self.in_channels = in_channels
+ self.model_channels = model_channels
+ self.out_channels = out_channels
+ if isinstance(num_res_blocks, int):
+ self.num_res_blocks = len(channel_mult) * [num_res_blocks]
+ else:
+ if len(num_res_blocks) != len(channel_mult):
+ raise ValueError(
+ "provide num_res_blocks either as an int (globally constant) or "
+ "as a list/tuple (per-level) with the same length as channel_mult"
+ )
+ self.num_res_blocks = num_res_blocks
+
+ if num_attention_blocks is not None:
+ assert len(num_attention_blocks) == len(self.num_res_blocks)
+ assert all(
+ map(
+ lambda i: self.num_res_blocks[i] >= num_attention_blocks[i],
+ range(len(num_attention_blocks)),
+ )
+ )
+ print(
+ f"Constructor of UNetModel received num_attention_blocks={num_attention_blocks}. "
+ f"This option has LESS priority than attention_resolutions {attention_resolutions}, "
+ f"i.e., in cases where num_attention_blocks[i] > 0 but 2**i not in attention_resolutions, "
+ f"attention will still not be set."
+ )
+
+ self.attention_resolutions = attention_resolutions
+ self.dropout = dropout
+ self.channel_mult = channel_mult
+ self.conv_resample = conv_resample
+ self.num_classes = num_classes
+ self.num_heads = num_heads
+ self.num_head_channels = num_head_channels
+ self.num_heads_upsample = num_heads_upsample
+ self.predict_codebook_ids = n_embed is not None
+
+ self.ip_dim = ip_dim
+ self.ip_weight = ip_weight
+
+ if self.ip_dim > 0:
+ self.image_embed = Resampler(
+ dim=context_dim,
+ depth=4,
+ dim_head=64,
+ heads=12,
+ num_queries=ip_dim, # num token
+ embedding_dim=1280,
+ output_dim=context_dim,
+ ff_mult=4,
+ )
+
+ time_embed_dim = model_channels * 4
+ self.time_embed = nn.Sequential(
+ nn.Linear(model_channels, time_embed_dim),
+ nn.SiLU(),
+ nn.Linear(time_embed_dim, time_embed_dim),
+ )
+
+ if camera_dim is not None:
+ time_embed_dim = model_channels * 4
+ self.camera_embed = nn.Sequential(
+ nn.Linear(camera_dim, time_embed_dim),
+ nn.SiLU(),
+ nn.Linear(time_embed_dim, time_embed_dim),
+ )
+
+ if self.num_classes is not None:
+ if isinstance(self.num_classes, int):
+ self.label_emb = nn.Embedding(self.num_classes, time_embed_dim)
+ elif self.num_classes == "continuous":
+ # print("setting up linear c_adm embedding layer")
+ self.label_emb = nn.Linear(1, time_embed_dim)
+ elif self.num_classes == "sequential":
+ assert adm_in_channels is not None
+ self.label_emb = nn.Sequential(
+ nn.Sequential(
+ nn.Linear(adm_in_channels, time_embed_dim),
+ nn.SiLU(),
+ nn.Linear(time_embed_dim, time_embed_dim),
+ )
+ )
+ else:
+ raise ValueError()
+
+ self.input_blocks = nn.ModuleList(
+ [
+ CondSequential(
+ conv_nd(dims, in_channels, model_channels, 3, padding=1)
+ )
+ ]
+ )
+ self._feature_size = model_channels
+ input_block_chans = [model_channels]
+ ch = model_channels
+ ds = 1
+ for level, mult in enumerate(channel_mult):
+ for nr in range(self.num_res_blocks[level]):
+ layers: List[Any] = [
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=mult * model_channels,
+ dims=dims,
+ use_scale_shift_norm=use_scale_shift_norm,
+ )
+ ]
+ ch = mult * model_channels
+ if ds in attention_resolutions:
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+
+ if num_attention_blocks is None or nr < num_attention_blocks[level]:
+ layers.append(
+ SpatialTransformer3D(
+ ch,
+ num_heads,
+ dim_head,
+ context_dim=context_dim,
+ depth=transformer_depth,
+ ip_dim=self.ip_dim,
+ ip_weight=self.ip_weight,
+ )
+ )
+ self.input_blocks.append(CondSequential(*layers))
+ self._feature_size += ch
+ input_block_chans.append(ch)
+ if level != len(channel_mult) - 1:
+ out_ch = ch
+ self.input_blocks.append(
+ CondSequential(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=out_ch,
+ dims=dims,
+ use_scale_shift_norm=use_scale_shift_norm,
+ down=True,
+ )
+ if resblock_updown
+ else Downsample(
+ ch, conv_resample, dims=dims, out_channels=out_ch
+ )
+ )
+ )
+ ch = out_ch
+ input_block_chans.append(ch)
+ ds *= 2
+ self._feature_size += ch
+
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+
+ self.middle_block = CondSequential(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ dims=dims,
+ use_scale_shift_norm=use_scale_shift_norm,
+ ),
+ SpatialTransformer3D(
+ ch,
+ num_heads,
+ dim_head,
+ context_dim=context_dim,
+ depth=transformer_depth,
+ ip_dim=self.ip_dim,
+ ip_weight=self.ip_weight,
+ ),
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ dims=dims,
+ use_scale_shift_norm=use_scale_shift_norm,
+ ),
+ )
+ self._feature_size += ch
+
+ self.output_blocks = nn.ModuleList([])
+ for level, mult in list(enumerate(channel_mult))[::-1]:
+ for i in range(self.num_res_blocks[level] + 1):
+ ich = input_block_chans.pop()
+ layers = [
+ ResBlock(
+ ch + ich,
+ time_embed_dim,
+ dropout,
+ out_channels=model_channels * mult,
+ dims=dims,
+ use_scale_shift_norm=use_scale_shift_norm,
+ )
+ ]
+ ch = model_channels * mult
+ if ds in attention_resolutions:
+ if num_head_channels == -1:
+ dim_head = ch // num_heads
+ else:
+ num_heads = ch // num_head_channels
+ dim_head = num_head_channels
+
+ if num_attention_blocks is None or i < num_attention_blocks[level]:
+ layers.append(
+ SpatialTransformer3D(
+ ch,
+ num_heads,
+ dim_head,
+ context_dim=context_dim,
+ depth=transformer_depth,
+ ip_dim=self.ip_dim,
+ ip_weight=self.ip_weight,
+ )
+ )
+ if level and i == self.num_res_blocks[level]:
+ out_ch = ch
+ layers.append(
+ ResBlock(
+ ch,
+ time_embed_dim,
+ dropout,
+ out_channels=out_ch,
+ dims=dims,
+ use_scale_shift_norm=use_scale_shift_norm,
+ up=True,
+ )
+ if resblock_updown
+ else Upsample(ch, conv_resample, dims=dims, out_channels=out_ch)
+ )
+ ds //= 2
+ self.output_blocks.append(CondSequential(*layers))
+ self._feature_size += ch
+
+ self.out = nn.Sequential(
+ nn.GroupNorm(32, ch),
+ nn.SiLU(),
+ zero_module(conv_nd(dims, model_channels, out_channels, 3, padding=1)),
+ )
+ if self.predict_codebook_ids:
+ self.id_predictor = nn.Sequential(
+ nn.GroupNorm(32, ch),
+ conv_nd(dims, model_channels, n_embed, 1),
+ # nn.LogSoftmax(dim=1) # change to cross_entropy and produce non-normalized logits
+ )
+
+ def forward(
+ self,
+ x,
+ timesteps=None,
+ context=None,
+ y=None,
+ camera=None,
+ num_frames=1,
+ ip=None,
+ ip_img=None,
+ **kwargs,
+ ):
+ """
+ Apply the model to an input batch.
+ :param x: an [(N x F) x C x ...] Tensor of inputs. F is the number of frames (views).
+ :param timesteps: a 1-D batch of timesteps.
+ :param context: conditioning plugged in via crossattn
+ :param y: an [N] Tensor of labels, if class-conditional.
+ :param num_frames: a integer indicating number of frames for tensor reshaping.
+ :return: an [(N x F) x C x ...] Tensor of outputs. F is the number of frames (views).
+ """
+ assert (
+ x.shape[0] % num_frames == 0
+ ), "input batch size must be dividable by num_frames!"
+ assert (y is not None) == (
+ self.num_classes is not None
+ ), "must specify y if and only if the model is class-conditional"
+
+ hs = []
+
+ t_emb = timestep_embedding(timesteps, self.model_channels, repeat_only=False).to(x.dtype)
+
+ emb = self.time_embed(t_emb)
+
+ if self.num_classes is not None:
+ assert y is not None
+ assert y.shape[0] == x.shape[0]
+ emb = emb + self.label_emb(y)
+
+ # Add camera embeddings
+ if camera is not None:
+ emb = emb + self.camera_embed(camera)
+
+ # imagedream variant
+ if self.ip_dim > 0:
+ x[(num_frames - 1) :: num_frames, :, :, :] = ip_img # place at [4, 9]
+ ip_emb = self.image_embed(ip)
+ context = torch.cat((context, ip_emb), 1)
+
+ h = x
+ for module in self.input_blocks:
+ h = module(h, emb, context, num_frames=num_frames)
+ hs.append(h)
+ h = self.middle_block(h, emb, context, num_frames=num_frames)
+ for module in self.output_blocks:
+ h = torch.cat([h, hs.pop()], dim=1)
+ h = module(h, emb, context, num_frames=num_frames)
+ h = h.type(x.dtype)
+ if self.predict_codebook_ids:
+ return self.id_predictor(h)
+ else:
+ return self.out(h)
\ No newline at end of file
diff --git a/lgm/mvdream/pipeline_mvdream.py b/lgm/mvdream/pipeline_mvdream.py
new file mode 100644
index 0000000..1cf37f4
--- /dev/null
+++ b/lgm/mvdream/pipeline_mvdream.py
@@ -0,0 +1,559 @@
+import torch
+import torch.nn.functional as F
+import inspect
+import numpy as np
+from typing import Callable, List, Optional, Union
+from transformers import CLIPTextModel, CLIPTokenizer, CLIPVisionModel, CLIPImageProcessor
+from diffusers import AutoencoderKL, DiffusionPipeline
+from diffusers.utils import (
+ deprecate,
+ is_accelerate_available,
+ is_accelerate_version,
+ logging,
+)
+from diffusers.configuration_utils import FrozenDict
+from diffusers.schedulers import DDIMScheduler
+from diffusers.utils.torch_utils import randn_tensor
+
+from lgm.mvdream.mv_unet import MultiViewUNetModel, get_camera
+
+logger = logging.get_logger(__name__) # pylint: disable=invalid-name
+
+
+class MVDreamPipeline(DiffusionPipeline):
+
+ _optional_components = ["feature_extractor", "image_encoder"]
+
+ def __init__(
+ self,
+ vae: AutoencoderKL,
+ unet: MultiViewUNetModel,
+ tokenizer: CLIPTokenizer,
+ text_encoder: CLIPTextModel,
+ scheduler: DDIMScheduler,
+ # imagedream variant
+ feature_extractor: CLIPImageProcessor,
+ image_encoder: CLIPVisionModel,
+ requires_safety_checker: bool = False,
+ ):
+ super().__init__()
+
+ if hasattr(scheduler.config, "steps_offset") and scheduler.config.steps_offset != 1: # type: ignore
+ deprecation_message = (
+ f"The configuration file of this scheduler: {scheduler} is outdated. `steps_offset`"
+ f" should be set to 1 instead of {scheduler.config.steps_offset}. Please make sure " # type: ignore
+ "to update the config accordingly as leaving `steps_offset` might led to incorrect results"
+ " in future versions. If you have downloaded this checkpoint from the Hugging Face Hub,"
+ " it would be very nice if you could open a Pull request for the `scheduler/scheduler_config.json`"
+ " file"
+ )
+ deprecate(
+ "steps_offset!=1", "1.0.0", deprecation_message, standard_warn=False
+ )
+ new_config = dict(scheduler.config)
+ new_config["steps_offset"] = 1
+ scheduler._internal_dict = FrozenDict(new_config)
+
+ if hasattr(scheduler.config, "clip_sample") and scheduler.config.clip_sample is True: # type: ignore
+ deprecation_message = (
+ f"The configuration file of this scheduler: {scheduler} has not set the configuration `clip_sample`."
+ " `clip_sample` should be set to False in the configuration file. Please make sure to update the"
+ " config accordingly as not setting `clip_sample` in the config might lead to incorrect results in"
+ " future versions. If you have downloaded this checkpoint from the Hugging Face Hub, it would be very"
+ " nice if you could open a Pull request for the `scheduler/scheduler_config.json` file"
+ )
+ deprecate(
+ "clip_sample not set", "1.0.0", deprecation_message, standard_warn=False
+ )
+ new_config = dict(scheduler.config)
+ new_config["clip_sample"] = False
+ scheduler._internal_dict = FrozenDict(new_config)
+
+ self.register_modules(
+ vae=vae,
+ unet=unet,
+ scheduler=scheduler,
+ tokenizer=tokenizer,
+ text_encoder=text_encoder,
+ feature_extractor=feature_extractor,
+ image_encoder=image_encoder,
+ )
+ self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
+ self.register_to_config(requires_safety_checker=requires_safety_checker)
+
+ def enable_vae_slicing(self):
+ r"""
+ Enable sliced VAE decoding.
+
+ When this option is enabled, the VAE will split the input tensor in slices to compute decoding in several
+ steps. This is useful to save some memory and allow larger batch sizes.
+ """
+ self.vae.enable_slicing()
+
+ def disable_vae_slicing(self):
+ r"""
+ Disable sliced VAE decoding. If `enable_vae_slicing` was previously invoked, this method will go back to
+ computing decoding in one step.
+ """
+ self.vae.disable_slicing()
+
+ def enable_vae_tiling(self):
+ r"""
+ Enable tiled VAE decoding.
+
+ When this option is enabled, the VAE will split the input tensor into tiles to compute decoding and encoding in
+ several steps. This is useful to save a large amount of memory and to allow the processing of larger images.
+ """
+ self.vae.enable_tiling()
+
+ def disable_vae_tiling(self):
+ r"""
+ Disable tiled VAE decoding. If `enable_vae_tiling` was previously invoked, this method will go back to
+ computing decoding in one step.
+ """
+ self.vae.disable_tiling()
+
+ def enable_sequential_cpu_offload(self, gpu_id=0):
+ r"""
+ Offloads all models to CPU using accelerate, significantly reducing memory usage. When called, unet,
+ text_encoder, vae and safety checker have their state dicts saved to CPU and then are moved to a
+ `torch.device('meta') and loaded to GPU only when their specific submodule has its `forward` method called.
+ Note that offloading happens on a submodule basis. Memory savings are higher than with
+ `enable_model_cpu_offload`, but performance is lower.
+ """
+ if is_accelerate_available() and is_accelerate_version(">=", "0.14.0"):
+ from accelerate import cpu_offload
+ else:
+ raise ImportError(
+ "`enable_sequential_cpu_offload` requires `accelerate v0.14.0` or higher"
+ )
+
+ device = torch.device(f"cuda:{gpu_id}")
+
+ if self.device.type != "cpu":
+ self.to("cpu", silence_dtype_warnings=True)
+ torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
+
+ for cpu_offloaded_model in [self.unet, self.text_encoder, self.vae]:
+ cpu_offload(cpu_offloaded_model, device)
+
+ def enable_model_cpu_offload(self, gpu_id=0):
+ r"""
+ Offloads all models to CPU using accelerate, reducing memory usage with a low impact on performance. Compared
+ to `enable_sequential_cpu_offload`, this method moves one whole model at a time to the GPU when its `forward`
+ method is called, and the model remains in GPU until the next model runs. Memory savings are lower than with
+ `enable_sequential_cpu_offload`, but performance is much better due to the iterative execution of the `unet`.
+ """
+ if is_accelerate_available() and is_accelerate_version(">=", "0.17.0.dev0"):
+ from accelerate import cpu_offload_with_hook
+ else:
+ raise ImportError(
+ "`enable_model_offload` requires `accelerate v0.17.0` or higher."
+ )
+
+ device = torch.device(f"cuda:{gpu_id}")
+
+ if self.device.type != "cpu":
+ self.to("cpu", silence_dtype_warnings=True)
+ torch.cuda.empty_cache() # otherwise we don't see the memory savings (but they probably exist)
+
+ hook = None
+ for cpu_offloaded_model in [self.text_encoder, self.unet, self.vae]:
+ _, hook = cpu_offload_with_hook(
+ cpu_offloaded_model, device, prev_module_hook=hook
+ )
+
+ # We'll offload the last model manually.
+ self.final_offload_hook = hook
+
+ @property
+ def _execution_device(self):
+ r"""
+ Returns the device on which the pipeline's models will be executed. After calling
+ `pipeline.enable_sequential_cpu_offload()` the execution device can only be inferred from Accelerate's module
+ hooks.
+ """
+ if not hasattr(self.unet, "_hf_hook"):
+ return self.device
+ for module in self.unet.modules():
+ if (
+ hasattr(module, "_hf_hook")
+ and hasattr(module._hf_hook, "execution_device")
+ and module._hf_hook.execution_device is not None
+ ):
+ return torch.device(module._hf_hook.execution_device)
+ return self.device
+
+ def _encode_prompt(
+ self,
+ prompt,
+ device,
+ num_images_per_prompt,
+ do_classifier_free_guidance: bool,
+ negative_prompt=None,
+ ):
+ r"""
+ Encodes the prompt into text encoder hidden states.
+
+ Args:
+ prompt (`str` or `List[str]`, *optional*):
+ prompt to be encoded
+ device: (`torch.device`):
+ torch device
+ num_images_per_prompt (`int`):
+ number of images that should be generated per prompt
+ do_classifier_free_guidance (`bool`):
+ whether to use classifier free guidance or not
+ negative_prompt (`str` or `List[str]`, *optional*):
+ The prompt or prompts not to guide the image generation. If not defined, one has to pass
+ `negative_prompt_embeds`. instead. If not defined, one has to pass `negative_prompt_embeds`. instead.
+ Ignored when not using guidance (i.e., ignored if `guidance_scale` is less than `1`).
+ prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+ provided, text embeddings will be generated from `prompt` input argument.
+ negative_prompt_embeds (`torch.FloatTensor`, *optional*):
+ Pre-generated negative text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt
+ weighting. If not provided, negative_prompt_embeds will be generated from `negative_prompt` input
+ argument.
+ """
+ if prompt is not None and isinstance(prompt, str):
+ batch_size = 1
+ elif prompt is not None and isinstance(prompt, list):
+ batch_size = len(prompt)
+ else:
+ raise ValueError(
+ f"`prompt` should be either a string or a list of strings, but got {type(prompt)}."
+ )
+
+ text_inputs = self.tokenizer(
+ prompt,
+ padding="max_length",
+ max_length=self.tokenizer.model_max_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+ text_input_ids = text_inputs.input_ids
+ untruncated_ids = self.tokenizer(
+ prompt, padding="longest", return_tensors="pt"
+ ).input_ids
+
+ if untruncated_ids.shape[-1] >= text_input_ids.shape[-1] and not torch.equal(
+ text_input_ids, untruncated_ids
+ ):
+ removed_text = self.tokenizer.batch_decode(
+ untruncated_ids[:, self.tokenizer.model_max_length - 1 : -1]
+ )
+ logger.warning(
+ "The following part of your input was truncated because CLIP can only handle sequences up to"
+ f" {self.tokenizer.model_max_length} tokens: {removed_text}"
+ )
+
+ if (
+ hasattr(self.text_encoder.config, "use_attention_mask")
+ and self.text_encoder.config.use_attention_mask
+ ):
+ attention_mask = text_inputs.attention_mask.to(device)
+ else:
+ attention_mask = None
+
+ prompt_embeds = self.text_encoder(
+ text_input_ids.to(device),
+ attention_mask=attention_mask,
+ )
+ prompt_embeds = prompt_embeds[0]
+
+ prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+ bs_embed, seq_len, _ = prompt_embeds.shape
+ # duplicate text embeddings for each generation per prompt, using mps friendly method
+ prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt, 1)
+ prompt_embeds = prompt_embeds.view(
+ bs_embed * num_images_per_prompt, seq_len, -1
+ )
+
+ # get unconditional embeddings for classifier free guidance
+ if do_classifier_free_guidance:
+ uncond_tokens: List[str]
+ if negative_prompt is None:
+ uncond_tokens = [""] * batch_size
+ elif type(prompt) is not type(negative_prompt):
+ raise TypeError(
+ f"`negative_prompt` should be the same type to `prompt`, but got {type(negative_prompt)} !="
+ f" {type(prompt)}."
+ )
+ elif isinstance(negative_prompt, str):
+ uncond_tokens = [negative_prompt]
+ elif batch_size != len(negative_prompt):
+ raise ValueError(
+ f"`negative_prompt`: {negative_prompt} has batch size {len(negative_prompt)}, but `prompt`:"
+ f" {prompt} has batch size {batch_size}. Please make sure that passed `negative_prompt` matches"
+ " the batch size of `prompt`."
+ )
+ else:
+ uncond_tokens = negative_prompt
+
+ max_length = prompt_embeds.shape[1]
+ uncond_input = self.tokenizer(
+ uncond_tokens,
+ padding="max_length",
+ max_length=max_length,
+ truncation=True,
+ return_tensors="pt",
+ )
+
+ if (
+ hasattr(self.text_encoder.config, "use_attention_mask")
+ and self.text_encoder.config.use_attention_mask
+ ):
+ attention_mask = uncond_input.attention_mask.to(device)
+ else:
+ attention_mask = None
+
+ negative_prompt_embeds = self.text_encoder(
+ uncond_input.input_ids.to(device),
+ attention_mask=attention_mask,
+ )
+ negative_prompt_embeds = negative_prompt_embeds[0]
+
+ # duplicate unconditional embeddings for each generation per prompt, using mps friendly method
+ seq_len = negative_prompt_embeds.shape[1]
+
+ negative_prompt_embeds = negative_prompt_embeds.to(
+ dtype=self.text_encoder.dtype, device=device
+ )
+
+ negative_prompt_embeds = negative_prompt_embeds.repeat(
+ 1, num_images_per_prompt, 1
+ )
+ negative_prompt_embeds = negative_prompt_embeds.view(
+ batch_size * num_images_per_prompt, seq_len, -1
+ )
+
+ # For classifier free guidance, we need to do two forward passes.
+ # Here we concatenate the unconditional and text embeddings into a single batch
+ # to avoid doing two forward passes
+ prompt_embeds = torch.cat([negative_prompt_embeds, prompt_embeds])
+
+ return prompt_embeds
+
+ def decode_latents(self, latents):
+ latents = 1 / self.vae.config.scaling_factor * latents
+ image = self.vae.decode(latents).sample
+ image = (image / 2 + 0.5).clamp(0, 1)
+ # we always cast to float32 as this does not cause significant overhead and is compatible with bfloat16
+ image = image.cpu().permute(0, 2, 3, 1).float().numpy()
+ return image
+
+ def prepare_extra_step_kwargs(self, generator, eta):
+ # prepare extra kwargs for the scheduler step, since not all schedulers have the same signature
+ # eta (η) is only used with the DDIMScheduler, it will be ignored for other schedulers.
+ # eta corresponds to η in DDIM paper: https://arxiv.org/abs/2010.02502
+ # and should be between [0, 1]
+
+ accepts_eta = "eta" in set(
+ inspect.signature(self.scheduler.step).parameters.keys()
+ )
+ extra_step_kwargs = {}
+ if accepts_eta:
+ extra_step_kwargs["eta"] = eta
+
+ # check if the scheduler accepts generator
+ accepts_generator = "generator" in set(
+ inspect.signature(self.scheduler.step).parameters.keys()
+ )
+ if accepts_generator:
+ extra_step_kwargs["generator"] = generator
+ return extra_step_kwargs
+
+ def prepare_latents(
+ self,
+ batch_size,
+ num_channels_latents,
+ height,
+ width,
+ dtype,
+ device,
+ generator,
+ latents=None,
+ ):
+ shape = (
+ batch_size,
+ num_channels_latents,
+ height // self.vae_scale_factor,
+ width // self.vae_scale_factor,
+ )
+ if isinstance(generator, list) and len(generator) != batch_size:
+ raise ValueError(
+ f"You have passed a list of generators of length {len(generator)}, but requested an effective batch"
+ f" size of {batch_size}. Make sure the batch size matches the length of the generators."
+ )
+
+ if latents is None:
+ latents = randn_tensor(
+ shape, generator=generator, device=device, dtype=dtype
+ )
+ else:
+ latents = latents.to(device)
+
+ # scale the initial noise by the standard deviation required by the scheduler
+ latents = latents * self.scheduler.init_noise_sigma
+ return latents
+
+ def encode_image(self, image, device, num_images_per_prompt):
+ dtype = next(self.image_encoder.parameters()).dtype
+
+ if image.dtype == np.float32:
+ image = (image * 255).astype(np.uint8)
+
+ image = self.feature_extractor(image, return_tensors="pt").pixel_values
+ image = image.to(device=device, dtype=dtype)
+
+ image_embeds = self.image_encoder(image, output_hidden_states=True).hidden_states[-2]
+ image_embeds = image_embeds.repeat_interleave(num_images_per_prompt, dim=0)
+
+ return torch.zeros_like(image_embeds), image_embeds
+
+ def encode_image_latents(self, image, device, num_images_per_prompt):
+
+ dtype = next(self.image_encoder.parameters()).dtype
+
+ image = torch.from_numpy(image).unsqueeze(0).permute(0, 3, 1, 2).to(device=device) # [1, 3, H, W]
+ image = 2 * image - 1
+ image = F.interpolate(image, (256, 256), mode='bilinear', align_corners=False)
+ image = image.to(dtype=dtype)
+
+ posterior = self.vae.encode(image).latent_dist
+ latents = posterior.sample() * self.vae.config.scaling_factor # [B, C, H, W]
+ latents = latents.repeat_interleave(num_images_per_prompt, dim=0)
+
+ return torch.zeros_like(latents), latents
+
+ @torch.no_grad()
+ def __call__(
+ self,
+ prompt: str = "",
+ image: Optional[np.ndarray] = None,
+ height: int = 256,
+ width: int = 256,
+ elevation: float = 0,
+ num_inference_steps: int = 50,
+ guidance_scale: float = 7.0,
+ negative_prompt: str = "",
+ num_images_per_prompt: int = 1,
+ eta: float = 0.0,
+ generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+ output_type: Optional[str] = "numpy", # pil, numpy, latents
+ callback: Optional[Callable[[int, int, torch.FloatTensor], None]] = None,
+ callback_steps: int = 1,
+ num_frames: int = 4,
+ device=torch.device("cuda:0"),
+ ):
+ self.unet = self.unet.to(device=device)
+ self.vae = self.vae.to(device=device)
+ self.text_encoder = self.text_encoder.to(device=device)
+
+ # here `guidance_scale` is defined analog to the guidance weight `w` of equation (2)
+ # of the Imagen paper: https://arxiv.org/pdf/2205.11487.pdf . `guidance_scale = 1`
+ # corresponds to doing no classifier free guidance.
+ do_classifier_free_guidance = guidance_scale > 1.0
+
+ # Prepare timesteps
+ self.scheduler.set_timesteps(num_inference_steps, device=device)
+ timesteps = self.scheduler.timesteps
+
+ # imagedream variant
+ if image is not None:
+ assert isinstance(image, np.ndarray) and image.dtype == np.float32
+ self.image_encoder = self.image_encoder.to(device=device)
+ image_embeds_neg, image_embeds_pos = self.encode_image(image, device, num_images_per_prompt)
+ image_latents_neg, image_latents_pos = self.encode_image_latents(image, device, num_images_per_prompt)
+
+ _prompt_embeds = self._encode_prompt(
+ prompt=prompt,
+ device=device,
+ num_images_per_prompt=num_images_per_prompt,
+ do_classifier_free_guidance=do_classifier_free_guidance,
+ negative_prompt=negative_prompt,
+ ) # type: ignore
+ prompt_embeds_neg, prompt_embeds_pos = _prompt_embeds.chunk(2)
+
+ # Prepare latent variables
+ actual_num_frames = num_frames if image is None else num_frames + 1
+ latents: torch.Tensor = self.prepare_latents(
+ actual_num_frames * num_images_per_prompt,
+ 4,
+ height,
+ width,
+ prompt_embeds_pos.dtype,
+ device,
+ generator,
+ None,
+ )
+
+ if image is not None:
+ camera = get_camera(num_frames, elevation=elevation, extra_view=True).to(dtype=latents.dtype, device=device)
+ else:
+ camera = get_camera(num_frames, elevation=elevation, extra_view=False).to(dtype=latents.dtype, device=device)
+ camera = camera.repeat_interleave(num_images_per_prompt, dim=0)
+
+ # Prepare extra step kwargs.
+ extra_step_kwargs = self.prepare_extra_step_kwargs(generator, eta)
+
+ # Denoising loop
+ num_warmup_steps = len(timesteps) - num_inference_steps * self.scheduler.order
+ with self.progress_bar(total=num_inference_steps) as progress_bar:
+ for i, t in enumerate(timesteps):
+ # expand the latents if we are doing classifier free guidance
+ multiplier = 2 if do_classifier_free_guidance else 1
+ latent_model_input = torch.cat([latents] * multiplier)
+ latent_model_input = self.scheduler.scale_model_input(latent_model_input, t)
+
+ unet_inputs = {
+ 'x': latent_model_input,
+ 'timesteps': torch.tensor([t] * actual_num_frames * multiplier, dtype=latent_model_input.dtype, device=device),
+ 'context': torch.cat([prompt_embeds_neg] * actual_num_frames + [prompt_embeds_pos] * actual_num_frames),
+ 'num_frames': actual_num_frames,
+ 'camera': torch.cat([camera] * multiplier),
+ }
+
+ if image is not None:
+ unet_inputs['ip'] = torch.cat([image_embeds_neg] * actual_num_frames + [image_embeds_pos] * actual_num_frames)
+ unet_inputs['ip_img'] = torch.cat([image_latents_neg] + [image_latents_pos]) # no repeat
+
+ # predict the noise residual
+ noise_pred = self.unet.forward(**unet_inputs)
+
+ # perform guidance
+ if do_classifier_free_guidance:
+ noise_pred_uncond, noise_pred_text = noise_pred.chunk(2)
+ noise_pred = noise_pred_uncond + guidance_scale * (
+ noise_pred_text - noise_pred_uncond
+ )
+
+ # compute the previous noisy sample x_t -> x_t-1
+ latents: torch.Tensor = self.scheduler.step(
+ noise_pred, t, latents, **extra_step_kwargs, return_dict=False
+ )[0]
+
+ # call the callback, if provided
+ if i == len(timesteps) - 1 or (
+ (i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0
+ ):
+ progress_bar.update()
+ if callback is not None and i % callback_steps == 0:
+ callback(i, t, latents) # type: ignore
+
+ # Post-processing
+ if output_type == "latent":
+ image = latents
+ elif output_type == "pil":
+ image = self.decode_latents(latents)
+ image = self.numpy_to_pil(image)
+ else: # numpy
+ image = self.decode_latents(latents)
+
+ # Offload last model to CPU
+ if hasattr(self, "final_offload_hook") and self.final_offload_hook is not None:
+ self.final_offload_hook.offload()
+
+ return image
\ No newline at end of file
diff --git a/mesh_processer/mesh.py b/mesh_processer/mesh.py
index 6e2f85b..f3082a8 100644
--- a/mesh_processer/mesh.py
+++ b/mesh_processer/mesh.py
@@ -336,7 +336,7 @@ class Mesh:
texture = np.ones((res_H, res_W, 3), dtype=np.float32) * np.array([0.5, 0.5, 0.5])
self.albedo = torch.tensor(texture, dtype=torch.float32, device=self.device)
else:
- self.albedo = prepare_torch_img(self.albedo, res_H, res_W, self.device).squeeze(0).permute(1, 2, 0).contiguous() # (1, 3, H, W) -> (H, W, 3)
+ self.albedo = prepare_torch_img(self.albedo.unsqueeze(0), res_H, res_W, self.device).squeeze(0).permute(1, 2, 0).contiguous() # (1, 3, H, W) -> (H, W, 3)
# aabb
def aabb(self):
diff --git a/nodes.py b/nodes.py
index b88954f..c9a0602 100644
--- a/nodes.py
+++ b/nodes.py
@@ -4,15 +4,17 @@ import re
import math
import copy
from enum import Enum
+from collections import OrderedDict
import folder_paths as comfy_paths
import torch
from torch.utils.data import DataLoader
+import torchvision.transforms.functional as TF
import numpy as np
+from safetensors.torch import load_file
from plyfile import PlyData
-from .shared_utils.common_utils import cstr, parse_save_filename
from .mesh_processer.mesh import Mesh
from .mesh_processer.mesh_utils import (
ply_to_points_cloud,
@@ -26,9 +28,20 @@ from .algorithms.main_3DGS_renderer import GaussianSplattingRenderer
from .algorithms.diff_texturing import DiffTextureBaker
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 .tgs.utils.config import ExperimentConfig, load_config
from .tgs.data import CustomImageOrbitDataset
from .tgs.utils.misc import todevice, get_device
+from .lgm.core.options import config_defaults
+from .lgm.mvdream.pipeline_mvdream import MVDreamPipeline
+
+from .shared_utils.image_utils import prepare_torch_img
+from .shared_utils.common_utils import cstr, parse_save_filename, get_list_filenames
+
+DIFFUSERS_PIPE_DICT = OrderedDict([
+ ("MVDreamPipeline", MVDreamPipeline),
+])
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "ComfyUI-3D-Pack")
@@ -688,7 +701,7 @@ class Gaussian_Splatting_Orbit_Renderer:
all_rendered_images, all_rendered_masks = cam_controller.render_all_pose(render_orbit_camera_poses)
all_rendered_images = all_rendered_images.permute(0, 2, 3, 1) # [N, 3, H, W] -> [N, H, W, 3]
- all_rendered_masks = all_rendered_masks.permute(0, 2, 3, 1).squeeze(3) # [N, 1, H, W] -> [N, H, W]
+ all_rendered_masks = all_rendered_masks.squeeze(1) # [N, 1, H, W] -> [N, H, W]
return (all_rendered_images, all_rendered_masks)
@@ -990,14 +1003,15 @@ TGS_CONFIG_PATH = "configs/tgs_config.yaml"
class Load_Triplane_Gaussian_Transformers:
+ checkpoints_dir = "checkpoints/tgs"
+ default_ckpt_name = "model_lvis_rel.ckpt"
+
@classmethod
def INPUT_TYPES(cls):
return {
"required": {
+ "model_name": (get_list_filenames(os.path.join(ROOT_PATH, cls.checkpoints_dir)), ),
},
-
- "optional": {
- }
}
RETURN_TYPES = (
@@ -1009,7 +1023,7 @@ class Load_Triplane_Gaussian_Transformers:
FUNCTION = "load_TGS"
CATEGORY = "Comfy3D/Import|Export"
- def load_TGS(self):
+ def load_TGS(self, model_name):
device = get_device()
@@ -1017,19 +1031,20 @@ class Load_Triplane_Gaussian_Transformers:
cfg: ExperimentConfig = load_config(config_path)
# Download pre-trained model if it not exist locally
- model_name = "model_lvis_rel.ckpt"
- model_path = os.path.join(ROOT_PATH, f"checkpoints/{model_name}")
- if not os.path.isfile(model_path):
+ ckpt_path = os.path.join(ROOT_PATH, self.checkpoints_dir, model_name)
+ if not os.path.isfile(ckpt_path):
from huggingface_hub import hf_hub_download
- hf_hub_download(repo_id="VAST-AI/TriplaneGaussian", local_dir="./checkpoints", filename=model_name, repo_type="model")
-
- cfg.system.weights=model_path
+ hf_hub_download(repo_id="VAST-AI/TriplaneGaussian", local_dir=self.checkpoints_dir, filename=self.default_ckpt_name, repo_type="model")
+ cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo VAST-AI/TriplaneGaussian instead").warning.print()
+ ckpt_path = os.path.join(ROOT_PATH, self.checkpoints_dir, self.default_ckpt_name)
+
+ cfg.system.weights=ckpt_path
tgs_model = TGS(cfg=cfg.system).to(device)
save_path = os.path.join(ROOT_PATH, "outputs")
tgs_model.set_save_dir(save_path)
- cstr(f"[{self.__class__.__name__}] loaded model ckpt from {model_path}").msg.print()
+ cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
return (tgs_model, )
@@ -1046,7 +1061,6 @@ class Triplane_Gaussian_Transformers:
},
}
- OUTPUT_NODE = True
RETURN_TYPES = (
"GS_PLY",
)
@@ -1090,4 +1104,217 @@ class Triplane_Gaussian_Transformers:
"""
return (gs_ply[0],)
-
\ No newline at end of file
+
+class Load_Diffusers_Pipeline:
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "diffusers_pipeline_name": (list(DIFFUSERS_PIPE_DICT.keys()),),
+ "pretrained_model_name": ("STRING", {"default": "ashawkey/imagedream-ipmv-diffusers", "multiline": False}),
+ },
+ }
+
+ RETURN_TYPES = (
+ "DIFFUSERS_PIPE",
+ )
+ RETURN_NAMES = (
+ "pipe",
+ )
+ FUNCTION = "load_diffusers_pipe"
+ CATEGORY = "Comfy3D/Import|Export"
+
+ def load_diffusers_pipe(self, diffusers_pipeline_name, pretrained_model_name):
+
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+
+ diffusers_pipeline_class = DIFFUSERS_PIPE_DICT[diffusers_pipeline_name]
+
+ # load image dream
+ pipe = diffusers_pipeline_class.from_pretrained(
+ pretrained_model_name, # remote weights
+ torch_dtype=torch.float16,
+ trust_remote_code=True,
+ # local_files_only=True,
+ ).to(device)
+
+ return (pipe,)
+
+class MVDream_Model:
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "mvdream_pipe": ("DIFFUSERS_PIPE",),
+ "reference_image": ("IMAGE",),
+ "reference_mask": ("MASK",),
+ "prompt": ("STRING", {
+ "default": "",
+ "multiline": True
+ }),
+ "prompt_neg": ("STRING", {
+ "default": "ugly, blurry, pixelated obscure, unnatural colors, poor lighting, dull, unclear, cropped, lowres, low quality, artifacts, duplicate",
+ "multiline": True
+ }),
+ "mv_guidance_scale": ("FLOAT", {"default": 5.0, "min": 0.0, "step": 0.01}),
+ "num_inference_steps": ("INT", {"default": 30, "min": 1}),
+ "elevation": ("FLOAT", {"default": 0.0, "min": ELEVATION_MIN, "max": ELEVATION_MAX, 'step': 0.0001}),
+ },
+ }
+
+ RETURN_TYPES = (
+ "IMAGE",
+ )
+ RETURN_NAMES = (
+ "multiview_images",
+ )
+ FUNCTION = "run_mvdream"
+ CATEGORY = "Comfy3D/Algorithm"
+
+ def run_mvdream(
+ self,
+ mvdream_pipe,
+ reference_image, # [1, H, W, 3]
+ reference_mask, # [1, H, W]
+ prompt,
+ prompt_neg,
+ mv_guidance_scale,
+ num_inference_steps,
+ elevation,
+ ):
+ if len(reference_image.shape) == 4:
+ reference_image = reference_image.squeeze(0)
+ if len(reference_mask.shape) == 3:
+ reference_mask = reference_mask.squeeze(0)
+
+ reference_mask = reference_mask.unsqueeze(2)
+ # give the white background to reference_image
+ reference_image = (reference_image * reference_mask + (1 - reference_mask)).detach().cpu().numpy()
+
+ # generate multi-view images
+ mv_images = mvdream_pipe(prompt, reference_image, negative_prompt=prompt_neg, guidance_scale=mv_guidance_scale, num_inference_steps=num_inference_steps, elevation=elevation)
+ mv_images = torch.from_numpy(np.stack([mv_images[1], mv_images[2], mv_images[3], mv_images[0]], axis=0)).float() # [4, H, W, 3], float32
+
+ return (mv_images,)
+
+class Load_Large_Multiview_Gaussian_Model:
+
+ checkpoints_dir = "checkpoints/lgm"
+ default_ckpt_name = "model_fp16.safetensors"
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "model_name": (get_list_filenames(os.path.join(ROOT_PATH, cls.checkpoints_dir)), ),
+ "lgb_config": (['big', 'default', 'small', 'tiny'], )
+ },
+ }
+
+ RETURN_TYPES = (
+ "LGM_MODEL",
+ )
+ RETURN_NAMES = (
+ "lgm_model",
+ )
+ FUNCTION = "load_LGM"
+ CATEGORY = "Comfy3D/Import|Export"
+
+ def load_LGM(self, model_name, lgb_config):
+
+ lgm_model = LGM(config_defaults[lgb_config])
+
+ # resume pretrained checkpoint
+ ckpt_path = os.path.join(ROOT_PATH, self.checkpoints_dir, model_name)
+ if not os.path.isfile(ckpt_path):
+ from huggingface_hub import hf_hub_download
+ hf_hub_download(repo_id="ashawkey/LGM", local_dir=self.checkpoints_dir, filename=self.default_ckpt_name, repo_type="model")
+ cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo ashawkey/LGM instead").warning.print()
+ ckpt_path = os.path.join(ROOT_PATH, self.checkpoints_dir, self.default_ckpt_name)
+
+ elif ckpt_path.endswith('safetensors'):
+ ckpt = load_file(ckpt_path, device='cpu')
+ else:
+ ckpt = torch.load(ckpt_path, map_location='cpu')
+
+ lgm_model.load_state_dict(ckpt, strict=False)
+ cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
+
+ # device
+ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+ lgm_model = lgm_model.half().to(device)
+ lgm_model.eval()
+
+ return (lgm_model,)
+
+IMAGENET_DEFAULT_MEAN = (0.485, 0.456, 0.406)
+IMAGENET_DEFAULT_STD = (0.229, 0.224, 0.225)
+
+class Large_Multiview_Gaussian_Model:
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "multiview_images": ("IMAGE", ),
+ "lgm_model": ("LGM_MODEL", ),
+ "mvdream_pipe": ("DIFFUSERS_PIPE",),
+ },
+ }
+
+ OUTPUT_NODE = True
+ RETURN_TYPES = (
+ "GS_PLY",
+ )
+ RETURN_NAMES = (
+ "gs_ply",
+ )
+ FUNCTION = "run_LGM"
+ CATEGORY = "Comfy3D/Algorithm"
+
+ @torch.no_grad()
+ def run_LGM(self, multiview_images, lgm_model, mvdream_pipe):
+ device = "cuda"
+ ref_image_torch = prepare_torch_img(multiview_images, lgm_model.opt.input_size, lgm_model.opt.input_size, device) # [4, 3, 256, 256]
+ ref_image_torch = TF.normalize(ref_image_torch, IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD)
+ rays_embeddings = lgm_model.prepare_default_rays(device)
+ ref_image_torch = torch.cat([ref_image_torch, rays_embeddings], dim=1).unsqueeze(0) # [1, 4, 9, 256, 256]
+
+ with torch.autocast(device_type=device, dtype=torch.float16):
+ # generate gaussians
+ gaussians = lgm_model.forward_gaussians(ref_image_torch)
+
+ # convert gaussians to ply
+ gs_ply = lgm_model.gs.to_ply(gaussians)
+
+ return (gs_ply, )
+
+class Convert_3DGS_to_Mesh_with_NeRF_and_Marching_Cubes:
+
+ @classmethod
+ def INPUT_TYPES(cls):
+ return {
+ "required": {
+ "gs_ply": ("GS_PLY",),
+ "gs_config": (['big', 'default', 'small', 'tiny'], )
+ },
+ }
+
+ RETURN_TYPES = (
+ "MESH",
+ )
+ RETURN_NAMES = (
+ "mesh",
+ )
+ FUNCTION = "convert_gs_ply"
+ CATEGORY = "Comfy3D/Algorithm"
+
+ def convert_gs_ply(self, gs_ply, gs_config):
+ with torch.inference_mode(False):
+ converter = GSConverterNeRFMarchingCubes(config_defaults[gs_config], gs_ply).cuda()
+ converter.fit_nerf()
+ converter.fit_mesh()
+ converter.fit_mesh_uv()
+
+ return(converter.get_mesh(), )
\ No newline at end of file
diff --git a/requirements.txt b/requirements.txt
index 5cf47fc..ab334cf 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,16 +1,27 @@
-#base
+# base
cmake
ninja
-kornia
-OmegaConf
+
+# computational libraries
numpy
einops
scipy
+kornia
opencv-python
+roma
+nerfacc>=0.5.3
+PyMCubes
+git+https://github.com/ashawkey/diff-gaussian-rasterization
+git+https://github.com/NVlabs/nvdiffrast/
+git+https://github.com/facebookresearch/pytorch3d.git@stable
+git+https://github.com/rusty1s/pytorch_scatter.git
+git+https://github.com/YixunLiang/simple-knn.git
+git+https://github.com/ashawkey/kiuikit.git
# for use ML models
diffusers
transformers
+safetensors
# for training differentiable tensors
pytorch_msssim
@@ -26,7 +37,9 @@ pygltflib
xatlas
pymeshlab
-# extra
+# configs & extra
torchtyping
tqdm
jaxtyping
+packaging
+OmegaConf
diff --git a/shared_utils/camera_utils.py b/shared_utils/camera_utils.py
index d079496..af73e8d 100644
--- a/shared_utils/camera_utils.py
+++ b/shared_utils/camera_utils.py
@@ -41,27 +41,6 @@ def look_at(campos, target, opengl=True):
R = np.stack([right_vector, up_vector, forward_vector], axis=1)
return R
-
-# elevation & azimuth to pose (cam2world) matrix
-def orbit_camera(elevation, azimuth, radius=1, is_degree=True, target=None, opengl=True):
- # radius: scalar
- # elevation: scalar, in (-90, 90), from +y to -y is (-90, 90)
- # azimuth: scalar, in (-180, 180), from +z to +x is (0, 90)
- # return: [4, 4], camera pose matrix
- if is_degree:
- elevation = np.deg2rad(elevation)
- azimuth = np.deg2rad(azimuth)
- x = radius * np.cos(elevation) * np.sin(azimuth)
- y = - radius * np.sin(elevation)
- z = radius * np.cos(elevation) * np.cos(azimuth)
- if target is None:
- target = np.zeros([3], dtype=np.float32)
- campos = np.array([x, y, z]) + target # [3]
- T = np.eye(4, dtype=np.float32)
- T[:3, :3] = look_at(campos, target, opengl)
- T[:3, 3] = campos
- return T
-
def get_look_at_camera_pose(target, target_to_cam_offset, look_distance=0.1, opengl=True):
"""
Calculate the pose (cam2world) matrix from target position the camera suppose to look at and offset vector from target to camera
diff --git a/shared_utils/common_utils.py b/shared_utils/common_utils.py
index 699c856..bfa891f 100644
--- a/shared_utils/common_utils.py
+++ b/shared_utils/common_utils.py
@@ -1,4 +1,6 @@
import os
+from os import listdir
+from os.path import isfile, join
import sys
from datetime import datetime
@@ -83,9 +85,9 @@ cstr.color.add_code("error", f"{cstr.color.RED}[Comfy3D] [ERROR] {cstr.color.END
def get_persistent_directory(folder_name):
if sys.platform == "win32":
- folder = os.path.join(os.path.expanduser("~"), "AppData", "Local", folder_name)
+ folder = join(os.path.expanduser("~"), "AppData", "Local", folder_name)
else:
- folder = os.path.join(os.path.expanduser("~"), "." + folder_name)
+ folder = join(os.path.expanduser("~"), "." + folder_name)
os.makedirs(folder, exist_ok=True)
return folder
@@ -96,7 +98,7 @@ def parse_save_filename(save_path, output_directory, supported_extensions, class
filename, file_extension = os.path.splitext(filename)
if file_extension.lower() in supported_extensions:
if not os.path.isabs(save_path):
- folder_path = os.path.join(output_directory, folder_path)
+ folder_path = join(output_directory, folder_path)
os.makedirs(folder_path, exist_ok=True)
@@ -107,10 +109,14 @@ def parse_save_filename(save_path, output_directory, supported_extensions, class
if date_format in filename:
filename = filename.replace(date_format, now.strftime(date_format))
- save_path = os.path.join(folder_path, filename) + file_extension
+ save_path = join(folder_path, filename) + file_extension
cstr(f"[{class_name}] Saving model to {save_path}").msg.print()
return save_path
else:
cstr(f"[{class_name}] File name {filename} does not end with supported file extensions: {supported_extensions}").error.print()
- return None
\ No newline at end of file
+ return None
+
+# Get all files in a folder, if extension_filter is provided then only reture the files with extensions in extension_filter
+def get_list_filenames(directory, extension_filter=None):
+ return [f for f in listdir(directory) if isfile(join(directory, f)) and (extension_filter is None or f.lower().endswith(extension_filter))]
\ No newline at end of file
diff --git a/shared_utils/image_utils.py b/shared_utils/image_utils.py
index 3cc3673..a0ee7f5 100644
--- a/shared_utils/image_utils.py
+++ b/shared_utils/image_utils.py
@@ -3,7 +3,7 @@ import torch.nn.functional as F
import numpy as np
def prepare_torch_img(img, size_H, size_W, device="cuda"):
- # [H, W, C] -> [1, C, H, W]
- img_new = img.permute(2, 0, 1).unsqueeze(0).to(device)
+ # [N, H, W, C] -> [N, C, H, W]
+ img_new = img.permute(0, 3, 1, 2).to(device)
img_new = F.interpolate(img_new, (size_H, size_W), mode="bilinear", align_corners=False).contiguous()
return img_new
\ No newline at end of file