Integrate LGM & Update 3DGS to Mesh Conversion Algorithm
This commit is contained in:
MrForExample
@@ -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
|
||||
|
||||
<video controls autoplay loop src="https://github.com/MrForExample/ComfyUI-3D-Pack/assets/62230687/9f3c56b1-afb3-4bf1-8845-ab1025a87463"></video>
|
||||
|
||||
- **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
|
||||
|
||||
<video controls autoplay loop src="https://github.com/MrForExample/ComfyUI-3D-Pack/assets/62230687/9f3c56b1-afb3-4bf1-8845-ab1025a87463"></video>
|
||||
|
||||
- **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
|
||||
|
||||
<table class="center">
|
||||
<tr>
|
||||
<td width=50% style="border: none">
|
||||
<img src="_Example_Workflows/_Example_Inputs_Files/a_pikachu_with_smily_face-removebg.png"/>
|
||||
</td>
|
||||
<td width=50% style="border: none">
|
||||
<video controls autoplay loop src="https://github.com/MrForExample/ComfyUI-3D-Pack/assets/62230687/aae7c16b-8d57-460d-b5b0-362f105a9d02" muted="false" width="256"></video>
|
||||
</td>
|
||||
</tr>
|
||||
</table>
|
||||
<video controls autoplay loop src="https://github.com/MrForExample/ComfyUI-3D-Pack/assets/62230687/9f3c56b1-afb3-4bf1-8845-ab1025a87463"></video>
|
||||
|
||||
- **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?)
|
||||
|
||||
|
||||
@@ -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
|
||||
}
|
||||
@@ -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
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 781 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 562 KiB |
Binary file not shown.
|
After Width: | Height: | Size: 54 KiB |
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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)
|
||||
|
||||
@@ -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
|
||||
@@ -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 ../../../../
|
||||
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
+186
@@ -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
|
||||
@@ -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',
|
||||
)
|
||||
@@ -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
|
||||
@@ -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
|
||||
@@ -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
|
||||
|
||||
File diff suppressed because it is too large
Load Diff
@@ -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
|
||||
@@ -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):
|
||||
|
||||
@@ -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],)
|
||||
|
||||
|
||||
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(), )
|
||||
+17
-4
@@ -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
|
||||
|
||||
@@ -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
|
||||
|
||||
@@ -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
|
||||
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))]
|
||||
@@ -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
|
||||
Reference in New Issue
Block a user