Integrate Zero123++ & InstantMesh, Add Mesh Orbit Renderer
This commit is contained in:
Mr. For Example
@@ -13,14 +13,21 @@
|
||||
### Note: this project is still a WIP
|
||||
|
||||
## Currently support:
|
||||
- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 12/04/2024**]
|
||||
- For use case please check [Example Workflows](./_Example_Workflows/). [**Last update: 23/05/2024**]
|
||||
- **Note:** you need to put [Example Inputs Files & Folders](_Example_Workflows/_Example_Inputs_Files/) under ComfyUI Root Directory\ComfyUI\input folder before you can run the example workflow
|
||||
- [tripoSR-layered-diffusion workflow](https://github.com/C0nsumption/Consume-ComfyUI-Workflows/tree/main/assets/tripo_sr/00) by [@Consumption](https://twitter.com/c0nsumption_)
|
||||
|
||||
- **InstantMesh Reconstruction Model**: [TencentARC/InstantMesh](https://github.com/TencentARC/InstantMesh)
|
||||
- Sparse multi-view images with white background to 3D Mesh with RGB texture
|
||||
- Works with arbitrary MVDiffusion models (Probably works best with Zero123++, but also works with CRM MVDiffusion model)
|
||||
|
||||
- **Zero123++**: [SUDO-AI-3D/zero123plus](https://github.com/SUDO-AI-3D/zero123plus)
|
||||
- Single image to 6 multi-view images
|
||||
|
||||
- **CRM**: [thu-ml/CRM](https://github.com/thu-ml/CRM)
|
||||
- Three stages pipeline:
|
||||
1. Image to 6 multi-view images (Front, Back, Left, Right, Top & Down)
|
||||
2. Image & 6 multi-view images to 6 same views CCMs (Canonical Coordinate Maps)
|
||||
1. Single image to 6 multi-view images (Front, Back, Left, Right, Top & Down)
|
||||
2. Single image & 6 multi-view images to 6 same views CCMs (Canonical Coordinate Maps)
|
||||
3. 6 multi-view images & CCMs to 3D mesh
|
||||
- *Note: For low vram pc, if you can't fit all three models for each stages into your GPU memory, then you can divide those three stages into different comfy workflow and run them separately*
|
||||
|
||||
@@ -94,6 +101,9 @@
|
||||
- **Gaussian Splatting Orbit Renderer**
|
||||
- Render 3DGS to images sequences or video, given a 3DGS file and camera poses generated by **Stack Orbit Camera Poses** node
|
||||
|
||||
- **Mesh Orbit Renderer**
|
||||
- Render 3D mesh to images sequences or video, given a mesh file and camera poses generated by **Stack Orbit Camera Poses** node
|
||||
|
||||
- **Fitting_Mesh_With_Multiview_Images**
|
||||
- Bake Multi-View images into UVTexture of given 3D mesh using [Nvdiffrast](https://github.com/NVlabs/nvdiffrast), supports:
|
||||
- Export to .obj, .ply, .glb
|
||||
|
||||
@@ -0,0 +1,652 @@
|
||||
{
|
||||
"last_node_id": 30,
|
||||
"last_link_id": 46,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 2,
|
||||
"type": "[Comfy3D] Load CRM MVDiffusion Model",
|
||||
"pos": [
|
||||
-102,
|
||||
131
|
||||
],
|
||||
"size": {
|
||||
"0": 489.21466064453125,
|
||||
"1": 95.50737762451172
|
||||
},
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "crm_mvdiffusion_sampler",
|
||||
"type": "CRM_MVDIFFUSION_SAMPLER",
|
||||
"links": [
|
||||
1
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load CRM MVDiffusion Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
"pixel-diffusion.pth",
|
||||
"configs/crm_configs/sd_v2_base_ipmv_zero_SNR.yaml"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 24,
|
||||
"type": "[Comfy3D] Save 3D Mesh",
|
||||
"pos": [
|
||||
2540,
|
||||
310
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 11,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"link": 34
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "save_path",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
35
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Save 3D Mesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"InstantMesh_test\\Mesh_Cat_CRM.obj"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 26,
|
||||
"type": "[Comfy3D] Switch Mesh Axis",
|
||||
"pos": [
|
||||
2170,
|
||||
310
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 130
|
||||
},
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"link": 36
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "switched_mesh",
|
||||
"type": "MESH",
|
||||
"links": [
|
||||
34
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Switch Mesh Axis"
|
||||
},
|
||||
"widgets_values": [
|
||||
"+y",
|
||||
"+z",
|
||||
"+x",
|
||||
false
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "[Comfy3D] CRM Images MVDiffusion Model",
|
||||
"pos": [
|
||||
600,
|
||||
300
|
||||
],
|
||||
"size": {
|
||||
"0": 506.4000244140625,
|
||||
"1": 262
|
||||
},
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "crm_mvdiffusion_sampler",
|
||||
"type": "CRM_MVDIFFUSION_SAMPLER",
|
||||
"link": 1
|
||||
},
|
||||
{
|
||||
"name": "reference_image",
|
||||
"type": "IMAGE",
|
||||
"link": 2
|
||||
},
|
||||
{
|
||||
"name": "reference_mask",
|
||||
"type": "MASK",
|
||||
"link": 28
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "multiview_images",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
39
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "orbit_camposes",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"links": [
|
||||
41
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] CRM Images MVDiffusion Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
"3D assets",
|
||||
"uniform low no texture ugly, boring, bad anatomy, blurry, pixelated, obscure, unnatural colors, poor lighting, dull, and unclear.",
|
||||
1234,
|
||||
"fixed",
|
||||
5.5,
|
||||
50
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 25,
|
||||
"type": "[Comfy3D] Preview 3DMesh",
|
||||
"pos": [
|
||||
2900,
|
||||
310
|
||||
],
|
||||
"size": [
|
||||
600,
|
||||
500
|
||||
],
|
||||
"flags": {},
|
||||
"order": 12,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh_file_path",
|
||||
"type": "STRING",
|
||||
"link": 35,
|
||||
"widget": {
|
||||
"name": "mesh_file_path"
|
||||
}
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Preview 3DMesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
null
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 27,
|
||||
"type": "GetImagesFromBatchIndexed",
|
||||
"pos": [
|
||||
1141,
|
||||
214
|
||||
],
|
||||
"size": {
|
||||
"0": 255.99082946777344,
|
||||
"1": 76.0000228881836
|
||||
},
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 39
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
44
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "GetImagesFromBatchIndexed"
|
||||
},
|
||||
"widgets_values": [
|
||||
"5, 3, 2, 0"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 22,
|
||||
"type": "[Comfy3D] Load InstantMesh Reconstruction Model",
|
||||
"pos": [
|
||||
1439,
|
||||
101
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "lrm_model",
|
||||
"type": "LRM_MODEL",
|
||||
"links": [
|
||||
33
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load InstantMesh Reconstruction Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
"instant_mesh_large.ckpt"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 29,
|
||||
"type": "[Comfy3D] Get Camposes From List Indexed",
|
||||
"pos": [
|
||||
1389,
|
||||
513
|
||||
],
|
||||
"size": {
|
||||
"0": 354.3999938964844,
|
||||
"1": 76
|
||||
},
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "original_orbit_camera_poses",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"link": 41
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "ORBIT_CAMPOSES",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"links": [
|
||||
42
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Get Camposes From List Indexed"
|
||||
},
|
||||
"widgets_values": [
|
||||
"5, 3, 2, 0"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 30,
|
||||
"type": "Image Rembg (Remove Background)",
|
||||
"pos": [
|
||||
1427,
|
||||
213
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 250
|
||||
},
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 44
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
45,
|
||||
46
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "Image Rembg (Remove Background)"
|
||||
},
|
||||
"widgets_values": [
|
||||
false,
|
||||
"u2net",
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
240,
|
||||
10,
|
||||
10,
|
||||
"white"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 23,
|
||||
"type": "[Comfy3D] InstantMesh Reconstruction Model",
|
||||
"pos": [
|
||||
1800,
|
||||
310
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 122
|
||||
},
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "lrm_model",
|
||||
"type": "LRM_MODEL",
|
||||
"link": 33
|
||||
},
|
||||
{
|
||||
"name": "multiview_images",
|
||||
"type": "IMAGE",
|
||||
"link": 46
|
||||
},
|
||||
{
|
||||
"name": "orbit_camera_poses",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"link": 42
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"links": [
|
||||
36
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] InstantMesh Reconstruction Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
30,
|
||||
1024
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 20,
|
||||
"type": "InvertMask",
|
||||
"pos": [
|
||||
341,
|
||||
400
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 26
|
||||
},
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 26
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
28
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "InvertMask"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
-191,
|
||||
320
|
||||
],
|
||||
"size": {
|
||||
"0": 473.3881530761719,
|
||||
"1": 545.7767944335938
|
||||
},
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
2
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
26
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"Mario_Cat.png",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 15,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
1805,
|
||||
494
|
||||
],
|
||||
"size": {
|
||||
"0": 770.9938354492188,
|
||||
"1": 246
|
||||
},
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 45
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
1,
|
||||
2,
|
||||
0,
|
||||
3,
|
||||
0,
|
||||
"CRM_MVDIFFUSION_SAMPLER"
|
||||
],
|
||||
[
|
||||
2,
|
||||
1,
|
||||
0,
|
||||
3,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
26,
|
||||
1,
|
||||
1,
|
||||
20,
|
||||
0,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
28,
|
||||
20,
|
||||
0,
|
||||
3,
|
||||
2,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
33,
|
||||
22,
|
||||
0,
|
||||
23,
|
||||
0,
|
||||
"LRM_MODEL"
|
||||
],
|
||||
[
|
||||
34,
|
||||
26,
|
||||
0,
|
||||
24,
|
||||
0,
|
||||
"MESH"
|
||||
],
|
||||
[
|
||||
35,
|
||||
24,
|
||||
0,
|
||||
25,
|
||||
0,
|
||||
"STRING"
|
||||
],
|
||||
[
|
||||
36,
|
||||
23,
|
||||
0,
|
||||
26,
|
||||
0,
|
||||
"MESH"
|
||||
],
|
||||
[
|
||||
39,
|
||||
3,
|
||||
0,
|
||||
27,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
41,
|
||||
3,
|
||||
1,
|
||||
29,
|
||||
0,
|
||||
"ORBIT_CAMPOSES"
|
||||
],
|
||||
[
|
||||
42,
|
||||
29,
|
||||
0,
|
||||
23,
|
||||
2,
|
||||
"ORBIT_CAMPOSES"
|
||||
],
|
||||
[
|
||||
44,
|
||||
27,
|
||||
0,
|
||||
30,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
45,
|
||||
30,
|
||||
0,
|
||||
15,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
46,
|
||||
30,
|
||||
0,
|
||||
23,
|
||||
1,
|
||||
"IMAGE"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.5644739300537773,
|
||||
"offset": [
|
||||
148.6926373068756,
|
||||
262.952479355259
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,383 @@
|
||||
{
|
||||
"last_node_id": 9,
|
||||
"last_link_id": 6,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 6,
|
||||
"type": "[Comfy3D] Load 3D Mesh",
|
||||
"pos": [
|
||||
337,
|
||||
232
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 154
|
||||
},
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"links": [
|
||||
3
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load 3D Mesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"C:\\Users\\reall\\Softwares\\ComfyUI_windows_portable\\ComfyUI\\input\\Meshes\\Cammy_Uniform.obj",
|
||||
false,
|
||||
false,
|
||||
false,
|
||||
true
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 4,
|
||||
"type": "[Comfy3D] Stack Orbit Camera Poses",
|
||||
"pos": [
|
||||
339,
|
||||
476
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 586
|
||||
},
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "orbit_camposes",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"links": [
|
||||
2,
|
||||
4
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "orbit_radius_list",
|
||||
"type": "FLOAT",
|
||||
"links": null,
|
||||
"shape": 6
|
||||
},
|
||||
{
|
||||
"name": "elevation_list",
|
||||
"type": "FLOAT",
|
||||
"links": null,
|
||||
"shape": 6
|
||||
},
|
||||
{
|
||||
"name": "azimuth_list",
|
||||
"type": "FLOAT",
|
||||
"links": null,
|
||||
"shape": 6
|
||||
},
|
||||
{
|
||||
"name": "orbit_center_X_list",
|
||||
"type": "FLOAT",
|
||||
"links": null,
|
||||
"shape": 6
|
||||
},
|
||||
{
|
||||
"name": "orbit_center_Y_list",
|
||||
"type": "FLOAT",
|
||||
"links": null,
|
||||
"shape": 6
|
||||
},
|
||||
{
|
||||
"name": "orbit_center_Z_list",
|
||||
"type": "FLOAT",
|
||||
"links": null,
|
||||
"shape": 6
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Stack Orbit Camera Poses"
|
||||
},
|
||||
"widgets_values": [
|
||||
1.75,
|
||||
1.75,
|
||||
0.1,
|
||||
-45,
|
||||
45,
|
||||
45,
|
||||
0,
|
||||
-0.01,
|
||||
12,
|
||||
0,
|
||||
0,
|
||||
0.1,
|
||||
0,
|
||||
0,
|
||||
0.1,
|
||||
0,
|
||||
0,
|
||||
0.1
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 7,
|
||||
"type": "[Comfy3D] Gaussian Splatting Orbit Renderer",
|
||||
"pos": [
|
||||
730,
|
||||
720
|
||||
],
|
||||
"size": {
|
||||
"0": 556.800048828125,
|
||||
"1": 198
|
||||
},
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "gs_ply",
|
||||
"type": "GS_PLY",
|
||||
"link": 5
|
||||
},
|
||||
{
|
||||
"name": "render_orbit_camera_poses",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"link": 4
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "rendered_gs_images",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
6
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "rendered_gs_masks",
|
||||
"type": "MASK",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Gaussian Splatting Orbit Renderer"
|
||||
},
|
||||
"widgets_values": [
|
||||
1024,
|
||||
1024,
|
||||
49.1,
|
||||
0,
|
||||
0,
|
||||
0
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 9,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
1361,
|
||||
657
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 246
|
||||
},
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 6
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"type": "[Comfy3D] Mesh Orbit Renderer",
|
||||
"pos": [
|
||||
734,
|
||||
330
|
||||
],
|
||||
"size": {
|
||||
"0": 395,
|
||||
"1": 247
|
||||
},
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"link": 3
|
||||
},
|
||||
{
|
||||
"name": "render_orbit_camera_poses",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"link": 2
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "rendered_mesh_images",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
1
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "rendered_mesh_masks",
|
||||
"type": "MASK",
|
||||
"links": null,
|
||||
"shape": 3
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Mesh Orbit Renderer"
|
||||
},
|
||||
"widgets_values": [
|
||||
1024,
|
||||
1024,
|
||||
49.1,
|
||||
0,
|
||||
0,
|
||||
0,
|
||||
false
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
1197,
|
||||
230
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 246
|
||||
},
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "[Comfy3D] Load 3DGS",
|
||||
"pos": [
|
||||
329,
|
||||
1139
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "gs_ply",
|
||||
"type": "GS_PLY",
|
||||
"links": [
|
||||
5
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load 3DGS"
|
||||
},
|
||||
"widgets_values": [
|
||||
"C:\\Users\\reall\\Softwares\\ComfyUI_windows_portable\\ComfyUI\\input\\Meshes\\3DGS_TGS.ply"
|
||||
]
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
1,
|
||||
1,
|
||||
0,
|
||||
2,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
2,
|
||||
4,
|
||||
0,
|
||||
1,
|
||||
1,
|
||||
"ORBIT_CAMPOSES"
|
||||
],
|
||||
[
|
||||
3,
|
||||
6,
|
||||
0,
|
||||
1,
|
||||
0,
|
||||
"MESH"
|
||||
],
|
||||
[
|
||||
4,
|
||||
4,
|
||||
0,
|
||||
7,
|
||||
1,
|
||||
"ORBIT_CAMPOSES"
|
||||
],
|
||||
[
|
||||
5,
|
||||
8,
|
||||
0,
|
||||
7,
|
||||
0,
|
||||
"GS_PLY"
|
||||
],
|
||||
[
|
||||
6,
|
||||
7,
|
||||
0,
|
||||
9,
|
||||
0,
|
||||
"IMAGE"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.6830134553650709,
|
||||
"offset": [
|
||||
232.77218799488367,
|
||||
-20.214945769251532
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
@@ -0,0 +1,597 @@
|
||||
{
|
||||
"last_node_id": 12,
|
||||
"last_link_id": 14,
|
||||
"nodes": [
|
||||
{
|
||||
"id": 4,
|
||||
"type": "[Comfy3D] Set Diffusers Pipeline State Dict",
|
||||
"pos": [
|
||||
980,
|
||||
160
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 82
|
||||
},
|
||||
"flags": {},
|
||||
"order": 5,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"link": 2
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"links": [
|
||||
3
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Set Diffusers Pipeline State Dict"
|
||||
},
|
||||
"widgets_values": [
|
||||
"TencentARC/InstantMesh",
|
||||
"diffusion_pytorch_model.bin"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 6,
|
||||
"type": "LoadImage",
|
||||
"pos": [
|
||||
692,
|
||||
356
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 314
|
||||
},
|
||||
"flags": {},
|
||||
"order": 0,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "IMAGE",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
4
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
5
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "LoadImage"
|
||||
},
|
||||
"widgets_values": [
|
||||
"Mario_Cat.png",
|
||||
"image"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 7,
|
||||
"type": "InvertMask",
|
||||
"pos": [
|
||||
1060,
|
||||
423
|
||||
],
|
||||
"size": {
|
||||
"0": 210,
|
||||
"1": 26
|
||||
},
|
||||
"flags": {},
|
||||
"order": 3,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mask",
|
||||
"type": "MASK",
|
||||
"link": 5
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "MASK",
|
||||
"type": "MASK",
|
||||
"links": [
|
||||
6
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "InvertMask"
|
||||
}
|
||||
},
|
||||
{
|
||||
"id": 3,
|
||||
"type": "[Comfy3D] Set Diffusers Pipeline Scheduler",
|
||||
"pos": [
|
||||
539,
|
||||
159
|
||||
],
|
||||
"size": {
|
||||
"0": 412.3726501464844,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 4,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"link": 1
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"links": [
|
||||
2
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Set Diffusers Pipeline Scheduler"
|
||||
},
|
||||
"widgets_values": [
|
||||
"EulerAncestralDiscreteScheduler"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 1,
|
||||
"type": "[Comfy3D] Load Diffusers Pipeline",
|
||||
"pos": [
|
||||
116,
|
||||
160
|
||||
],
|
||||
"size": {
|
||||
"0": 383.9324645996094,
|
||||
"1": 109.8599853515625
|
||||
},
|
||||
"flags": {},
|
||||
"order": 1,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"links": [
|
||||
1
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load Diffusers Pipeline"
|
||||
},
|
||||
"widgets_values": [
|
||||
"Zero123PlusPipeline",
|
||||
"sudo-ai/zero123plus-v1.2",
|
||||
""
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 2,
|
||||
"type": "[Comfy3D] Load InstantMesh Reconstruction Model",
|
||||
"pos": [
|
||||
1428,
|
||||
159
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 2,
|
||||
"mode": 0,
|
||||
"outputs": [
|
||||
{
|
||||
"name": "lrm_model",
|
||||
"type": "LRM_MODEL",
|
||||
"links": [
|
||||
7
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Load InstantMesh Reconstruction Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
"instant_mesh_large.ckpt"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 5,
|
||||
"type": "[Comfy3D] Zero123Plus Diffusion Model",
|
||||
"pos": [
|
||||
1320,
|
||||
336
|
||||
],
|
||||
"size": {
|
||||
"0": 393,
|
||||
"1": 170
|
||||
},
|
||||
"flags": {},
|
||||
"order": 6,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "zero123plus_pipe",
|
||||
"type": "DIFFUSERS_PIPE",
|
||||
"link": 3
|
||||
},
|
||||
{
|
||||
"name": "reference_image",
|
||||
"type": "IMAGE",
|
||||
"link": 4
|
||||
},
|
||||
{
|
||||
"name": "reference_mask",
|
||||
"type": "MASK",
|
||||
"link": 6
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "multiviews",
|
||||
"type": "IMAGE",
|
||||
"links": [
|
||||
8,
|
||||
12
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
},
|
||||
{
|
||||
"name": "orbit_camposes",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"links": [
|
||||
9
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 1
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Zero123Plus Diffusion Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
42,
|
||||
"fixed",
|
||||
4,
|
||||
28
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 8,
|
||||
"type": "[Comfy3D] InstantMesh Reconstruction Model",
|
||||
"pos": [
|
||||
1807,
|
||||
159
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 122
|
||||
},
|
||||
"flags": {},
|
||||
"order": 8,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "lrm_model",
|
||||
"type": "LRM_MODEL",
|
||||
"link": 7
|
||||
},
|
||||
{
|
||||
"name": "multiview_images",
|
||||
"type": "IMAGE",
|
||||
"link": 8
|
||||
},
|
||||
{
|
||||
"name": "orbit_camera_poses",
|
||||
"type": "ORBIT_CAMPOSES",
|
||||
"link": 9
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"links": [
|
||||
13
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] InstantMesh Reconstruction Model"
|
||||
},
|
||||
"widgets_values": [
|
||||
30,
|
||||
1024
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 9,
|
||||
"type": "[Comfy3D] Save 3D Mesh",
|
||||
"pos": [
|
||||
2534,
|
||||
158
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 58
|
||||
},
|
||||
"flags": {},
|
||||
"order": 10,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"link": 14
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "save_path",
|
||||
"type": "STRING",
|
||||
"links": [
|
||||
11
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Save 3D Mesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"InstantMesh_test\\Mesh_Cat.obj"
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 10,
|
||||
"type": "[Comfy3D] Preview 3DMesh",
|
||||
"pos": [
|
||||
2919,
|
||||
158
|
||||
],
|
||||
"size": [
|
||||
600,
|
||||
500
|
||||
],
|
||||
"flags": {},
|
||||
"order": 11,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh_file_path",
|
||||
"type": "STRING",
|
||||
"link": 11,
|
||||
"widget": {
|
||||
"name": "mesh_file_path"
|
||||
}
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Preview 3DMesh"
|
||||
},
|
||||
"widgets_values": [
|
||||
"",
|
||||
null
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 12,
|
||||
"type": "[Comfy3D] Switch Mesh Axis",
|
||||
"pos": [
|
||||
2168,
|
||||
158
|
||||
],
|
||||
"size": {
|
||||
"0": 315,
|
||||
"1": 130
|
||||
},
|
||||
"flags": {},
|
||||
"order": 9,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "mesh",
|
||||
"type": "MESH",
|
||||
"link": 13
|
||||
}
|
||||
],
|
||||
"outputs": [
|
||||
{
|
||||
"name": "switched_mesh",
|
||||
"type": "MESH",
|
||||
"links": [
|
||||
14
|
||||
],
|
||||
"shape": 3,
|
||||
"slot_index": 0
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "[Comfy3D] Switch Mesh Axis"
|
||||
},
|
||||
"widgets_values": [
|
||||
"+y",
|
||||
"+z",
|
||||
"+x",
|
||||
false
|
||||
]
|
||||
},
|
||||
{
|
||||
"id": 11,
|
||||
"type": "PreviewImage",
|
||||
"pos": [
|
||||
1804,
|
||||
387
|
||||
],
|
||||
"size": {
|
||||
"0": 1043.55126953125,
|
||||
"1": 246
|
||||
},
|
||||
"flags": {},
|
||||
"order": 7,
|
||||
"mode": 0,
|
||||
"inputs": [
|
||||
{
|
||||
"name": "images",
|
||||
"type": "IMAGE",
|
||||
"link": 12
|
||||
}
|
||||
],
|
||||
"properties": {
|
||||
"Node name for S&R": "PreviewImage"
|
||||
}
|
||||
}
|
||||
],
|
||||
"links": [
|
||||
[
|
||||
1,
|
||||
1,
|
||||
0,
|
||||
3,
|
||||
0,
|
||||
"DIFFUSERS_PIPE"
|
||||
],
|
||||
[
|
||||
2,
|
||||
3,
|
||||
0,
|
||||
4,
|
||||
0,
|
||||
"DIFFUSERS_PIPE"
|
||||
],
|
||||
[
|
||||
3,
|
||||
4,
|
||||
0,
|
||||
5,
|
||||
0,
|
||||
"DIFFUSERS_PIPE"
|
||||
],
|
||||
[
|
||||
4,
|
||||
6,
|
||||
0,
|
||||
5,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
5,
|
||||
6,
|
||||
1,
|
||||
7,
|
||||
0,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
6,
|
||||
7,
|
||||
0,
|
||||
5,
|
||||
2,
|
||||
"MASK"
|
||||
],
|
||||
[
|
||||
7,
|
||||
2,
|
||||
0,
|
||||
8,
|
||||
0,
|
||||
"LRM_MODEL"
|
||||
],
|
||||
[
|
||||
8,
|
||||
5,
|
||||
0,
|
||||
8,
|
||||
1,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
9,
|
||||
5,
|
||||
1,
|
||||
8,
|
||||
2,
|
||||
"ORBIT_CAMPOSES"
|
||||
],
|
||||
[
|
||||
11,
|
||||
9,
|
||||
0,
|
||||
10,
|
||||
0,
|
||||
"STRING"
|
||||
],
|
||||
[
|
||||
12,
|
||||
5,
|
||||
0,
|
||||
11,
|
||||
0,
|
||||
"IMAGE"
|
||||
],
|
||||
[
|
||||
13,
|
||||
8,
|
||||
0,
|
||||
12,
|
||||
0,
|
||||
"MESH"
|
||||
],
|
||||
[
|
||||
14,
|
||||
12,
|
||||
0,
|
||||
9,
|
||||
0,
|
||||
"MESH"
|
||||
]
|
||||
],
|
||||
"groups": [],
|
||||
"config": {},
|
||||
"extra": {
|
||||
"ds": {
|
||||
"scale": 0.8264462809917354,
|
||||
"offset": [
|
||||
-788.9517631212848,
|
||||
67.47620178525084
|
||||
]
|
||||
}
|
||||
},
|
||||
"version": 0.4
|
||||
}
|
||||
Binary file not shown.
|
After Width: | Height: | Size: 554 KiB |
+13
-12
@@ -13,9 +13,14 @@ from pytorch_msssim import SSIM, MS_SSIM
|
||||
import comfy.utils
|
||||
|
||||
from .diff_mesh_renderer import DiffRastRenderer
|
||||
from ..shared_utils.camera_utils import OrbitCamera
|
||||
from ..shared_utils.camera_utils import BaseCameraController
|
||||
from ..shared_utils.image_utils import prepare_torch_img
|
||||
|
||||
class DiffMeshCameraController(BaseCameraController):
|
||||
def get_render_result(self, render_pose, bg_color):
|
||||
ref_cam = (render_pose, self.cam.perspective)
|
||||
return self.renderer.render(*ref_cam, self.cam.H, self.cam.W, ssaa=1, bg_color=bg_color, optional_render_types=[]) #ssaa = min(2.0, max(0.125, 2 * np.random.random()))
|
||||
|
||||
class DiffMesh:
|
||||
|
||||
def __init__(self, mesh, training_iterations, batch_size, texture_learning_rate, train_mesh_geometry, geometry_learning_rate, ms_ssim_loss_weight, remesh_after_n_iteration, force_cuda_rasterize):
|
||||
@@ -38,14 +43,16 @@ class DiffMesh:
|
||||
|
||||
def prepare_training(self, reference_images, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy):
|
||||
self.ref_imgs_num = len(reference_images)
|
||||
|
||||
self.all_ref_cam_poses = reference_orbit_camera_poses
|
||||
self.ref_cam_fovy = reference_orbit_camera_fovy
|
||||
|
||||
self.ref_size_H = reference_images[0].shape[0]
|
||||
self.ref_size_W = reference_images[0].shape[1]
|
||||
|
||||
self.cam = OrbitCamera(self.ref_size_W, self.ref_size_H, fovy=reference_orbit_camera_fovy)
|
||||
# default camera settings
|
||||
self.cam_controller = DiffMeshCameraController(
|
||||
self.renderer, self.ref_size_W, self.ref_size_H, reference_orbit_camera_fovy, self.gs_params.invert_bg_prob, None, self.device
|
||||
)
|
||||
|
||||
self.all_ref_cam_poses = reference_orbit_camera_poses
|
||||
|
||||
# prepare reference images and masks
|
||||
ref_imgs_torch_list = []
|
||||
@@ -80,13 +87,7 @@ class DiffMesh:
|
||||
|
||||
i = random.randint(0, ref_imgs_num_minus_1)
|
||||
|
||||
radius, elevation, azimuth, center_X, center_Y, center_Z = self.all_ref_cam_poses[i]
|
||||
|
||||
# render output
|
||||
orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
|
||||
ref_pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
|
||||
ref_cam = (ref_pose, self.cam.perspective)
|
||||
out = self.renderer.render(*ref_cam, self.ref_size_H, self.ref_size_W, ssaa=1) #ssaa = min(2.0, max(0.125, 2 * np.random.random()))
|
||||
out = self.cam_controller.render_at_pose(self.all_ref_cam_poses[i])
|
||||
|
||||
image = out["image"] # [H, W, 3] in [0, 1]
|
||||
image = image.permute(2, 0, 1).contiguous() # [3, H, W] in [0, 1]
|
||||
|
||||
@@ -51,6 +51,7 @@ class DiffRastRenderer(nn.Module):
|
||||
self.v_offsets = nn.Parameter(torch.zeros_like(self.mesh.v), requires_grad=True)
|
||||
self.raw_albedo = nn.Parameter(inverse_sigmoid(self.mesh.albedo), requires_grad=True)
|
||||
|
||||
self.train_geo = False
|
||||
|
||||
def get_params(self, texture_lr, train_geo, geom_lr):
|
||||
|
||||
@@ -68,7 +69,8 @@ class DiffRastRenderer(nn.Module):
|
||||
self.mesh.v = (self.mesh.v + self.v_offsets).detach()
|
||||
self.mesh.albedo = torch.sigmoid(self.raw_albedo.detach())
|
||||
|
||||
def render(self, pose, proj, h0, w0, ssaa=1, bg_color=1, texture_filter='linear-mipmap-linear'):
|
||||
def render(self, pose, proj, h0, w0, ssaa=1, bg_color=1, texture_filter='linear-mipmap-linear',
|
||||
optional_render_types=['depth', 'normal']):
|
||||
|
||||
# do super-sampling
|
||||
if ssaa != 1:
|
||||
@@ -94,36 +96,40 @@ class DiffRastRenderer(nn.Module):
|
||||
|
||||
rast, rast_db = dr.rasterize(self.glctx, v_clip, self.mesh.f, (h, w))
|
||||
|
||||
alpha = (rast[0, ..., 3:] > 0).float()
|
||||
depth, _ = dr.interpolate(-v_cam[..., [2]], rast, self.mesh.f) # [1, H, W, 1]
|
||||
depth = depth.squeeze(0) # [H, W, 1]
|
||||
|
||||
alpha = (rast[0, ..., 3:] > 0).float() # [H, W, 1]
|
||||
# render depth
|
||||
if 'depth' in optional_render_types:
|
||||
depth, _ = dr.interpolate(-v_cam[..., [2]], rast, self.mesh.f) # [1, H, W, 1]
|
||||
depth = depth.squeeze(0) # [H, W, 1]
|
||||
# render albedo
|
||||
texc, texc_db = dr.interpolate(self.mesh.vt.unsqueeze(0).contiguous(), rast, self.mesh.ft, rast_db=rast_db, diff_attrs='all')
|
||||
albedo = dr.texture(self.raw_albedo.unsqueeze(0), texc, uv_da=texc_db, filter_mode=texture_filter) # [1, H, W, 3]
|
||||
albedo = torch.sigmoid(albedo)
|
||||
|
||||
# get vn and render normal
|
||||
if self.train_geo:
|
||||
i0, i1, i2 = self.mesh.f[:, 0].long(), self.mesh.f[:, 1].long(), self.mesh.f[:, 2].long()
|
||||
v0, v1, v2 = v[i0, :], v[i1, :], v[i2, :]
|
||||
if 'normal' in optional_render_types:
|
||||
if self.train_geo:
|
||||
i0, i1, i2 = self.mesh.f[:, 0].long(), self.mesh.f[:, 1].long(), self.mesh.f[:, 2].long()
|
||||
v0, v1, v2 = v[i0, :], v[i1, :], v[i2, :]
|
||||
|
||||
face_normals = torch.cross(v1 - v0, v2 - v0)
|
||||
face_normals = safe_normalize(face_normals)
|
||||
face_normals = torch.cross(v1 - v0, v2 - v0)
|
||||
face_normals = safe_normalize(face_normals)
|
||||
|
||||
vn = torch.zeros_like(v)
|
||||
vn.scatter_add_(0, i0[:, None].repeat(1,3), face_normals)
|
||||
vn.scatter_add_(0, i1[:, None].repeat(1,3), face_normals)
|
||||
vn.scatter_add_(0, i2[:, None].repeat(1,3), face_normals)
|
||||
|
||||
vn = torch.where(torch.sum(vn * vn, -1, keepdim=True) > 1e-20, vn, torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device))
|
||||
else:
|
||||
vn = self.mesh.vn
|
||||
|
||||
vn = torch.zeros_like(v)
|
||||
vn.scatter_add_(0, i0[:, None].repeat(1,3), face_normals)
|
||||
vn.scatter_add_(0, i1[:, None].repeat(1,3), face_normals)
|
||||
vn.scatter_add_(0, i2[:, None].repeat(1,3), face_normals)
|
||||
normal, _ = dr.interpolate(vn.unsqueeze(0).contiguous(), rast, self.mesh.fn)
|
||||
normal = safe_normalize(normal[0])
|
||||
|
||||
vn = torch.where(torch.sum(vn * vn, -1, keepdim=True) > 1e-20, vn, torch.tensor([0.0, 0.0, 1.0], dtype=torch.float32, device=vn.device))
|
||||
else:
|
||||
vn = self.mesh.vn
|
||||
|
||||
normal, _ = dr.interpolate(vn.unsqueeze(0).contiguous(), rast, self.mesh.fn)
|
||||
normal = safe_normalize(normal[0])
|
||||
|
||||
# rotated normal (where [0, 0, 1] always faces camera)
|
||||
rot_normal = normal @ pose[:3, :3]
|
||||
viewcos = rot_normal[..., [2]]
|
||||
# rotated normal (where [0, 0, 1] always faces camera)
|
||||
rot_normal = normal @ pose[:3, :3]
|
||||
viewcos = rot_normal[..., [2]]
|
||||
|
||||
# antialias
|
||||
albedo = dr.antialias(albedo, rast, v_clip, self.mesh.f).squeeze(0) # [H, W, 3]
|
||||
@@ -133,14 +139,18 @@ class DiffRastRenderer(nn.Module):
|
||||
if ssaa != 1:
|
||||
albedo = scale_img_hwc(albedo, (h0, w0))
|
||||
alpha = scale_img_hwc(alpha, (h0, w0))
|
||||
depth = scale_img_hwc(depth, (h0, w0))
|
||||
normal = scale_img_hwc(normal, (h0, w0))
|
||||
viewcos = scale_img_hwc(viewcos, (h0, w0))
|
||||
if 'depth' in optional_render_types:
|
||||
depth = scale_img_hwc(depth, (h0, w0))
|
||||
if 'normal' in optional_render_types:
|
||||
normal = scale_img_hwc(normal, (h0, w0))
|
||||
viewcos = scale_img_hwc(viewcos, (h0, w0))
|
||||
|
||||
results['image'] = albedo.clamp(0, 1)
|
||||
results['alpha'] = alpha
|
||||
results['depth'] = depth
|
||||
results['normal'] = (normal + 1) / 2
|
||||
results['viewcos'] = viewcos
|
||||
if 'depth' in optional_render_types:
|
||||
results['depth'] = depth
|
||||
if 'normal' in optional_render_types:
|
||||
results['normal'] = (normal + 1) / 2
|
||||
results['viewcos'] = viewcos
|
||||
|
||||
return results
|
||||
+5
-39
@@ -4,13 +4,12 @@ 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 OrbitCamera, MiniCam, calculate_fovX, get_projection_matrix
|
||||
from ..shared_utils.camera_utils import BaseCameraController, MiniCam, calculate_fovX, get_projection_matrix
|
||||
from ..shared_utils.image_utils import prepare_torch_img
|
||||
|
||||
class GSParams:
|
||||
@@ -74,46 +73,13 @@ class GSParams:
|
||||
# other gaussian params
|
||||
self.sh_degree = sh_degree
|
||||
|
||||
class GaussianSplattingCameraController:
|
||||
def __init__(self, renderer, cam_size_W, cam_size_H, reference_orbit_camera_fovy, invert_bg_prob=1.0, static_bg=None, device='cuda'):
|
||||
self.device = torch.device(device)
|
||||
|
||||
self.renderer = renderer
|
||||
self.cam = OrbitCamera(cam_size_W, cam_size_H, fovy=reference_orbit_camera_fovy)
|
||||
self.projection_matrix = get_projection_matrix(self.cam.near, self.cam.far, self.cam.fovx, self.cam.fovy).transpose(0, 1).cuda()
|
||||
|
||||
self.invert_bg_prob = invert_bg_prob
|
||||
self.black_bg = torch.tensor([0, 0, 0], dtype=torch.float32, device=self.device)
|
||||
self.white_bg = torch.tensor([1, 1, 1], dtype=torch.float32, device=self.device)
|
||||
self.static_bg = None if static_bg is None else torch.tensor(static_bg, dtype=torch.float32, device=self.device)
|
||||
|
||||
def render_at_pose(self, cam_pose):
|
||||
radius, elevation, azimuth, center_X, center_Y, center_Z = cam_pose
|
||||
|
||||
orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
|
||||
render_pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
|
||||
class GaussianSplattingCameraController(BaseCameraController):
|
||||
def get_render_result(self, render_pose, bg_color):
|
||||
render_cam = MiniCam(render_pose, self.cam.W, self.cam.H, self.cam.fovy, self.cam.fovx, self.cam.near, self.cam.far, self.projection_matrix)
|
||||
|
||||
if self.static_bg is None:
|
||||
bg_color = self.white_bg if np.random.rand() > self.invert_bg_prob else self.black_bg
|
||||
else:
|
||||
bg_color = self.static_bg
|
||||
|
||||
return self.renderer.render(render_cam, bg_color=bg_color)
|
||||
|
||||
def render_all_pose(self, all_cam_poses):
|
||||
all_rendered_images, all_rendered_masks = [], []
|
||||
for cam_pose in all_cam_poses:
|
||||
out = self.render_at_pose(cam_pose)
|
||||
|
||||
image = out["image"] # [3, H, W] in [0, 1]
|
||||
mask = out["alpha"] # [1, H, W] in [0, 1]
|
||||
|
||||
all_rendered_images.append(image)
|
||||
all_rendered_masks.append(mask)
|
||||
|
||||
# [Number of Poses, 3, H, W], [Number of Poses, 1, H, W] both in [0, 1]
|
||||
return torch.stack(all_rendered_images, dim=0), torch.stack(all_rendered_masks, dim=0)
|
||||
def pose_init(self):
|
||||
self.projection_matrix = get_projection_matrix(self.cam.near, self.cam.far, self.cam.fovx, self.cam.fovy).transpose(0, 1).cuda()
|
||||
|
||||
class GaussianSplatting:
|
||||
|
||||
|
||||
@@ -0,0 +1,35 @@
|
||||
*.7z filter=lfs diff=lfs merge=lfs -text
|
||||
*.arrow filter=lfs diff=lfs merge=lfs -text
|
||||
*.bin filter=lfs diff=lfs merge=lfs -text
|
||||
*.bz2 filter=lfs diff=lfs merge=lfs -text
|
||||
*.ckpt filter=lfs diff=lfs merge=lfs -text
|
||||
*.ftz filter=lfs diff=lfs merge=lfs -text
|
||||
*.gz filter=lfs diff=lfs merge=lfs -text
|
||||
*.h5 filter=lfs diff=lfs merge=lfs -text
|
||||
*.joblib filter=lfs diff=lfs merge=lfs -text
|
||||
*.lfs.* filter=lfs diff=lfs merge=lfs -text
|
||||
*.mlmodel filter=lfs diff=lfs merge=lfs -text
|
||||
*.model filter=lfs diff=lfs merge=lfs -text
|
||||
*.msgpack filter=lfs diff=lfs merge=lfs -text
|
||||
*.npy filter=lfs diff=lfs merge=lfs -text
|
||||
*.npz filter=lfs diff=lfs merge=lfs -text
|
||||
*.onnx filter=lfs diff=lfs merge=lfs -text
|
||||
*.ot filter=lfs diff=lfs merge=lfs -text
|
||||
*.parquet filter=lfs diff=lfs merge=lfs -text
|
||||
*.pb filter=lfs diff=lfs merge=lfs -text
|
||||
*.pickle filter=lfs diff=lfs merge=lfs -text
|
||||
*.pkl filter=lfs diff=lfs merge=lfs -text
|
||||
*.pt filter=lfs diff=lfs merge=lfs -text
|
||||
*.pth filter=lfs diff=lfs merge=lfs -text
|
||||
*.rar filter=lfs diff=lfs merge=lfs -text
|
||||
*.safetensors filter=lfs diff=lfs merge=lfs -text
|
||||
saved_model/**/* filter=lfs diff=lfs merge=lfs -text
|
||||
*.tar.* filter=lfs diff=lfs merge=lfs -text
|
||||
*.tar filter=lfs diff=lfs merge=lfs -text
|
||||
*.tflite filter=lfs diff=lfs merge=lfs -text
|
||||
*.tgz filter=lfs diff=lfs merge=lfs -text
|
||||
*.wasm filter=lfs diff=lfs merge=lfs -text
|
||||
*.xz filter=lfs diff=lfs merge=lfs -text
|
||||
*.zip filter=lfs diff=lfs merge=lfs -text
|
||||
*.zst filter=lfs diff=lfs merge=lfs -text
|
||||
*tfevents* filter=lfs diff=lfs merge=lfs -text
|
||||
@@ -0,0 +1,27 @@
|
||||
{
|
||||
"crop_size": {
|
||||
"height": 224,
|
||||
"width": 224
|
||||
},
|
||||
"do_center_crop": true,
|
||||
"do_convert_rgb": true,
|
||||
"do_normalize": true,
|
||||
"do_rescale": true,
|
||||
"do_resize": true,
|
||||
"image_mean": [
|
||||
0.48145466,
|
||||
0.4578275,
|
||||
0.40821073
|
||||
],
|
||||
"image_processor_type": "CLIPImageProcessor",
|
||||
"image_std": [
|
||||
0.26862954,
|
||||
0.26130258,
|
||||
0.27577711
|
||||
],
|
||||
"resample": 3,
|
||||
"rescale_factor": 0.00392156862745098,
|
||||
"size": {
|
||||
"shortest_edge": 224
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,19 @@
|
||||
{
|
||||
"crop_size": {
|
||||
"height": 512,
|
||||
"width": 512
|
||||
},
|
||||
"do_center_crop": true,
|
||||
"do_convert_rgb": true,
|
||||
"do_normalize": true,
|
||||
"do_rescale": true,
|
||||
"do_resize": true,
|
||||
"image_mean": 0.5,
|
||||
"image_processor_type": "CLIPImageProcessor",
|
||||
"image_std": 0.8,
|
||||
"resample": 2,
|
||||
"rescale_factor": 0.00392156862745098,
|
||||
"size": {
|
||||
"shortest_edge": 512
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,120 @@
|
||||
{
|
||||
"_class_name": "Zero123PlusPipeline",
|
||||
"_diffusers_version": "0.24.0",
|
||||
"_name_or_path": "sudo-ai/zero123plus-v1.1",
|
||||
"feature_extractor_clip": [
|
||||
"transformers",
|
||||
"CLIPImageProcessor"
|
||||
],
|
||||
"feature_extractor_vae": [
|
||||
"transformers",
|
||||
"CLIPImageProcessor"
|
||||
],
|
||||
"ramping_coefficients": [
|
||||
0.00301829120144248,
|
||||
0.2204633206129074,
|
||||
0.21527841687202454,
|
||||
0.23498539626598358,
|
||||
0.1914631873369217,
|
||||
0.20188239216804504,
|
||||
0.19352824985980988,
|
||||
0.17249998450279236,
|
||||
0.15826298296451569,
|
||||
0.15236389636993408,
|
||||
0.13444548845291138,
|
||||
0.12044154852628708,
|
||||
0.12808501720428467,
|
||||
0.1271015852689743,
|
||||
0.13629068434238434,
|
||||
0.14516159892082214,
|
||||
0.15645112097263336,
|
||||
0.16885493695735931,
|
||||
0.18022602796554565,
|
||||
0.1958882212638855,
|
||||
0.21415705978870392,
|
||||
0.23056700825691223,
|
||||
0.2505834102630615,
|
||||
0.2574525773525238,
|
||||
0.275470107793808,
|
||||
0.2808215022087097,
|
||||
0.29953837394714355,
|
||||
0.2967497408390045,
|
||||
0.2883710563182831,
|
||||
0.3023308515548706,
|
||||
0.3054688572883606,
|
||||
0.32596179842948914,
|
||||
0.3225354254245758,
|
||||
0.3140765428543091,
|
||||
0.3288663625717163,
|
||||
0.3435625731945038,
|
||||
0.3342442810535431,
|
||||
0.32937031984329224,
|
||||
0.35734811425209045,
|
||||
0.3601177930831909,
|
||||
0.3517529368400574,
|
||||
0.3810708224773407,
|
||||
0.40007662773132324,
|
||||
0.4264647364616394,
|
||||
0.3977527916431427,
|
||||
0.4314143657684326,
|
||||
0.49558719992637634,
|
||||
0.4665665030479431,
|
||||
0.48960328102111816,
|
||||
0.5141982436180115,
|
||||
0.5230164527893066,
|
||||
0.5266074538230896,
|
||||
0.5456079840660095,
|
||||
0.5737904906272888,
|
||||
0.5882097482681274,
|
||||
0.6210350394248962,
|
||||
0.6530380845069885,
|
||||
0.6383244395256042,
|
||||
0.6792004704475403,
|
||||
0.6567418575286865,
|
||||
0.7517656683921814,
|
||||
0.736494243144989,
|
||||
0.7586457133293152,
|
||||
0.8130561709403992,
|
||||
0.9578766226768494,
|
||||
1.001284122467041,
|
||||
0.9404520988464355,
|
||||
1.004292368888855,
|
||||
0.9145274758338928,
|
||||
0.9771682620048523,
|
||||
1.0350638628005981,
|
||||
1.0265849828720093,
|
||||
1.0594775676727295,
|
||||
0.980824887752533,
|
||||
1.0715670585632324,
|
||||
1.0140161514282227,
|
||||
1.1983819007873535
|
||||
],
|
||||
"safety_checker": [
|
||||
null,
|
||||
null
|
||||
],
|
||||
"scheduler": [
|
||||
"diffusers",
|
||||
"EulerAncestralDiscreteScheduler"
|
||||
],
|
||||
"text_encoder": [
|
||||
"transformers",
|
||||
"CLIPTextModel"
|
||||
],
|
||||
"tokenizer": [
|
||||
"transformers",
|
||||
"CLIPTokenizer"
|
||||
],
|
||||
"unet": [
|
||||
"diffusers",
|
||||
"UNet2DConditionModel"
|
||||
],
|
||||
"vae": [
|
||||
"diffusers",
|
||||
"AutoencoderKL"
|
||||
],
|
||||
"vision_encoder": [
|
||||
"transformers",
|
||||
"CLIPVisionModelWithProjection"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
{
|
||||
"_class_name": "EulerAncestralDiscreteScheduler",
|
||||
"_diffusers_version": "0.24.0",
|
||||
"beta_end": 0.012,
|
||||
"beta_schedule": "linear",
|
||||
"beta_start": 0.00085,
|
||||
"clip_sample": false,
|
||||
"num_train_timesteps": 1000,
|
||||
"prediction_type": "v_prediction",
|
||||
"set_alpha_to_one": false,
|
||||
"skip_prk_steps": true,
|
||||
"steps_offset": 1,
|
||||
"timestep_spacing": "linspace",
|
||||
"trained_betas": null
|
||||
}
|
||||
@@ -0,0 +1,25 @@
|
||||
{
|
||||
"_name_or_path": "D:\\.cache\\huggingface\\hub\\models--sudo-ai--zero123plus-v1.1\\snapshots\\36df7de980afd15f80b2e1a4e9a920d7020e2654\\text_encoder",
|
||||
"architectures": [
|
||||
"CLIPTextModel"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"bos_token_id": 0,
|
||||
"dropout": 0.0,
|
||||
"eos_token_id": 2,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_size": 1024,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 4096,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"max_position_embeddings": 77,
|
||||
"model_type": "clip_text_model",
|
||||
"num_attention_heads": 16,
|
||||
"num_hidden_layers": 23,
|
||||
"pad_token_id": 1,
|
||||
"projection_dim": 512,
|
||||
"torch_dtype": "float16",
|
||||
"transformers_version": "4.36.0",
|
||||
"vocab_size": 49408
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,24 @@
|
||||
{
|
||||
"bos_token": {
|
||||
"content": "<|startoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"eos_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
},
|
||||
"pad_token": "!",
|
||||
"unk_token": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false
|
||||
}
|
||||
}
|
||||
@@ -0,0 +1,38 @@
|
||||
{
|
||||
"add_prefix_space": false,
|
||||
"added_tokens_decoder": {
|
||||
"0": {
|
||||
"content": "!",
|
||||
"lstrip": false,
|
||||
"normalized": false,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"49406": {
|
||||
"content": "<|startoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
},
|
||||
"49407": {
|
||||
"content": "<|endoftext|>",
|
||||
"lstrip": false,
|
||||
"normalized": true,
|
||||
"rstrip": false,
|
||||
"single_word": false,
|
||||
"special": true
|
||||
}
|
||||
},
|
||||
"bos_token": "<|startoftext|>",
|
||||
"clean_up_tokenization_spaces": true,
|
||||
"do_lower_case": true,
|
||||
"eos_token": "<|endoftext|>",
|
||||
"errors": "replace",
|
||||
"model_max_length": 77,
|
||||
"pad_token": "!",
|
||||
"tokenizer_class": "CLIPTokenizer",
|
||||
"unk_token": "<|endoftext|>"
|
||||
}
|
||||
File diff suppressed because it is too large
Load Diff
@@ -0,0 +1,73 @@
|
||||
{
|
||||
"_class_name": "UNet2DConditionModel",
|
||||
"_diffusers_version": "0.24.0",
|
||||
"_name_or_path": "D:\\.cache\\huggingface\\hub\\models--stabilityai--stable-diffusion-2\\snapshots\\1e128c8891e52218b74cde8f26dbfc701cb99d79\\unet",
|
||||
"act_fn": "silu",
|
||||
"addition_embed_type": null,
|
||||
"addition_embed_type_num_heads": 64,
|
||||
"addition_time_embed_dim": null,
|
||||
"attention_head_dim": [
|
||||
5,
|
||||
10,
|
||||
20,
|
||||
20
|
||||
],
|
||||
"attention_type": "default",
|
||||
"block_out_channels": [
|
||||
320,
|
||||
640,
|
||||
1280,
|
||||
1280
|
||||
],
|
||||
"center_input_sample": false,
|
||||
"class_embed_type": null,
|
||||
"class_embeddings_concat": false,
|
||||
"conv_in_kernel": 3,
|
||||
"conv_out_kernel": 3,
|
||||
"cross_attention_dim": 1024,
|
||||
"cross_attention_norm": null,
|
||||
"down_block_types": [
|
||||
"CrossAttnDownBlock2D",
|
||||
"CrossAttnDownBlock2D",
|
||||
"CrossAttnDownBlock2D",
|
||||
"DownBlock2D"
|
||||
],
|
||||
"downsample_padding": 1,
|
||||
"dropout": 0.0,
|
||||
"dual_cross_attention": false,
|
||||
"encoder_hid_dim": null,
|
||||
"encoder_hid_dim_type": null,
|
||||
"flip_sin_to_cos": true,
|
||||
"freq_shift": 0,
|
||||
"in_channels": 4,
|
||||
"layers_per_block": 2,
|
||||
"mid_block_only_cross_attention": null,
|
||||
"mid_block_scale_factor": 1,
|
||||
"mid_block_type": "UNetMidBlock2DCrossAttn",
|
||||
"norm_eps": 1e-05,
|
||||
"norm_num_groups": 32,
|
||||
"num_attention_heads": null,
|
||||
"num_class_embeds": null,
|
||||
"only_cross_attention": false,
|
||||
"out_channels": 4,
|
||||
"projection_class_embeddings_input_dim": null,
|
||||
"resnet_out_scale_factor": 1.0,
|
||||
"resnet_skip_time_act": false,
|
||||
"resnet_time_scale_shift": "default",
|
||||
"reverse_transformer_layers_per_block": null,
|
||||
"sample_size": 96,
|
||||
"time_cond_proj_dim": null,
|
||||
"time_embedding_act_fn": null,
|
||||
"time_embedding_dim": null,
|
||||
"time_embedding_type": "positional",
|
||||
"timestep_post_act": null,
|
||||
"transformer_layers_per_block": 1,
|
||||
"up_block_types": [
|
||||
"UpBlock2D",
|
||||
"CrossAttnUpBlock2D",
|
||||
"CrossAttnUpBlock2D",
|
||||
"CrossAttnUpBlock2D"
|
||||
],
|
||||
"upcast_attention": false,
|
||||
"use_linear_projection": true
|
||||
}
|
||||
@@ -0,0 +1,32 @@
|
||||
{
|
||||
"_class_name": "AutoencoderKL",
|
||||
"_diffusers_version": "0.24.0",
|
||||
"_name_or_path": "D:\\.cache\\huggingface\\hub\\models--sudo-ai--zero123plus-v1.1\\snapshots\\36df7de980afd15f80b2e1a4e9a920d7020e2654\\vae",
|
||||
"act_fn": "silu",
|
||||
"block_out_channels": [
|
||||
128,
|
||||
256,
|
||||
512,
|
||||
512
|
||||
],
|
||||
"down_block_types": [
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D",
|
||||
"DownEncoderBlock2D"
|
||||
],
|
||||
"force_upcast": true,
|
||||
"in_channels": 3,
|
||||
"latent_channels": 4,
|
||||
"layers_per_block": 2,
|
||||
"norm_num_groups": 32,
|
||||
"out_channels": 3,
|
||||
"sample_size": 256,
|
||||
"scaling_factor": 0.18215,
|
||||
"up_block_types": [
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D",
|
||||
"UpDecoderBlock2D"
|
||||
]
|
||||
}
|
||||
@@ -0,0 +1,23 @@
|
||||
{
|
||||
"_name_or_path": "D:\\.cache\\huggingface\\hub\\models--sudo-ai--zero123plus-v1.1\\snapshots\\36df7de980afd15f80b2e1a4e9a920d7020e2654\\vision_encoder",
|
||||
"architectures": [
|
||||
"CLIPVisionModelWithProjection"
|
||||
],
|
||||
"attention_dropout": 0.0,
|
||||
"dropout": 0.0,
|
||||
"hidden_act": "gelu",
|
||||
"hidden_size": 1280,
|
||||
"image_size": 224,
|
||||
"initializer_factor": 1.0,
|
||||
"initializer_range": 0.02,
|
||||
"intermediate_size": 5120,
|
||||
"layer_norm_eps": 1e-05,
|
||||
"model_type": "clip_vision_model",
|
||||
"num_attention_heads": 16,
|
||||
"num_channels": 3,
|
||||
"num_hidden_layers": 32,
|
||||
"patch_size": 14,
|
||||
"projection_dim": 1024,
|
||||
"torch_dtype": "float16",
|
||||
"transformers_version": "4.36.0"
|
||||
}
|
||||
@@ -0,0 +1,15 @@
|
||||
model_config:
|
||||
target: instant_mesh.models.lrm_mesh.InstantMesh
|
||||
params:
|
||||
encoder_feat_dim: 768
|
||||
encoder_freeze: false
|
||||
encoder_model_name: facebook/dino-vitb16
|
||||
transformer_dim: 1024
|
||||
transformer_layers: 12
|
||||
transformer_heads: 16
|
||||
triplane_low_res: 32
|
||||
triplane_high_res: 64
|
||||
triplane_dim: 40
|
||||
rendering_samples_per_ray: 96
|
||||
grid_res: 128
|
||||
grid_scale: 2.1
|
||||
@@ -0,0 +1,15 @@
|
||||
model_config:
|
||||
target: instant_mesh.models.lrm_mesh.InstantMesh
|
||||
params:
|
||||
encoder_feat_dim: 768
|
||||
encoder_freeze: false
|
||||
encoder_model_name: facebook/dino-vitb16
|
||||
transformer_dim: 1024
|
||||
transformer_layers: 16
|
||||
transformer_heads: 16
|
||||
triplane_low_res: 32
|
||||
triplane_high_res: 64
|
||||
triplane_dim: 80
|
||||
rendering_samples_per_ray: 128
|
||||
grid_res: 128
|
||||
grid_scale: 2.1
|
||||
@@ -0,0 +1,13 @@
|
||||
model_config:
|
||||
target: instant_mesh.models.lrm.InstantNeRF
|
||||
params:
|
||||
encoder_feat_dim: 768
|
||||
encoder_freeze: false
|
||||
encoder_model_name: facebook/dino-vitb16
|
||||
transformer_dim: 1024
|
||||
transformer_layers: 12
|
||||
transformer_heads: 16
|
||||
triplane_low_res: 32
|
||||
triplane_high_res: 64
|
||||
triplane_dim: 40
|
||||
rendering_samples_per_ray: 96
|
||||
@@ -0,0 +1,13 @@
|
||||
model_config:
|
||||
target: instant_mesh.models.lrm.InstantNeRF
|
||||
params:
|
||||
encoder_feat_dim: 768
|
||||
encoder_freeze: false
|
||||
encoder_model_name: facebook/dino-vitb16
|
||||
transformer_dim: 1024
|
||||
transformer_layers: 16
|
||||
transformer_heads: 16
|
||||
triplane_low_res: 32
|
||||
triplane_high_res: 64
|
||||
triplane_dim: 80
|
||||
rendering_samples_per_ray: 128
|
||||
@@ -0,0 +1,325 @@
|
||||
import os, sys
|
||||
import math
|
||||
import json
|
||||
import importlib
|
||||
from pathlib import Path
|
||||
|
||||
import cv2
|
||||
import random
|
||||
import numpy as np
|
||||
from PIL import Image
|
||||
import webdataset as wds
|
||||
import pytorch_lightning as pl
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torch.utils.data import Dataset
|
||||
from torch.utils.data import DataLoader
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
from torchvision import transforms
|
||||
|
||||
from instant_mesh.utils.train_util import instantiate_from_config
|
||||
from instant_mesh.utils.camera_util import (
|
||||
FOV_to_intrinsics,
|
||||
center_looking_at_camera_pose,
|
||||
get_circular_camera_poses,
|
||||
)
|
||||
|
||||
|
||||
class DataModuleFromConfig(pl.LightningDataModule):
|
||||
def __init__(
|
||||
self,
|
||||
batch_size=8,
|
||||
num_workers=4,
|
||||
train=None,
|
||||
validation=None,
|
||||
test=None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.batch_size = batch_size
|
||||
self.num_workers = num_workers
|
||||
|
||||
self.dataset_configs = dict()
|
||||
if train is not None:
|
||||
self.dataset_configs['train'] = train
|
||||
if validation is not None:
|
||||
self.dataset_configs['validation'] = validation
|
||||
if test is not None:
|
||||
self.dataset_configs['test'] = test
|
||||
|
||||
def setup(self, stage):
|
||||
|
||||
if stage in ['fit']:
|
||||
self.datasets = dict((k, instantiate_from_config(self.dataset_configs[k])) for k in self.dataset_configs)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
def train_dataloader(self):
|
||||
|
||||
sampler = DistributedSampler(self.datasets['train'])
|
||||
return wds.WebLoader(self.datasets['train'], batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, sampler=sampler)
|
||||
|
||||
def val_dataloader(self):
|
||||
|
||||
sampler = DistributedSampler(self.datasets['validation'])
|
||||
return wds.WebLoader(self.datasets['validation'], batch_size=1, num_workers=self.num_workers, shuffle=False, sampler=sampler)
|
||||
|
||||
def test_dataloader(self):
|
||||
|
||||
return wds.WebLoader(self.datasets['test'], batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
|
||||
|
||||
|
||||
class ObjaverseData(Dataset):
|
||||
def __init__(self,
|
||||
root_dir='objaverse/',
|
||||
meta_fname='valid_paths.json',
|
||||
input_image_dir='rendering_random_32views',
|
||||
target_image_dir='rendering_random_32views',
|
||||
input_view_num=6,
|
||||
target_view_num=4,
|
||||
total_view_n=32,
|
||||
fov=50,
|
||||
camera_rotation=True,
|
||||
validation=False,
|
||||
):
|
||||
self.root_dir = Path(root_dir)
|
||||
self.input_image_dir = input_image_dir
|
||||
self.target_image_dir = target_image_dir
|
||||
|
||||
self.input_view_num = input_view_num
|
||||
self.target_view_num = target_view_num
|
||||
self.total_view_n = total_view_n
|
||||
self.fov = fov
|
||||
self.camera_rotation = camera_rotation
|
||||
|
||||
with open(os.path.join(root_dir, meta_fname)) as f:
|
||||
filtered_dict = json.load(f)
|
||||
paths = filtered_dict['good_objs']
|
||||
self.paths = paths
|
||||
|
||||
self.depth_scale = 6.0
|
||||
|
||||
total_objects = len(self.paths)
|
||||
print('============= length of dataset %d =============' % len(self.paths))
|
||||
|
||||
def __len__(self):
|
||||
return len(self.paths)
|
||||
|
||||
def load_im(self, path, color):
|
||||
'''
|
||||
replace background pixel with random color in rendering
|
||||
'''
|
||||
pil_img = Image.open(path)
|
||||
|
||||
image = np.asarray(pil_img, dtype=np.float32) / 255.
|
||||
alpha = image[:, :, 3:]
|
||||
image = image[:, :, :3] * alpha + color * (1 - alpha)
|
||||
|
||||
image = torch.from_numpy(image).permute(2, 0, 1).contiguous().float()
|
||||
alpha = torch.from_numpy(alpha).permute(2, 0, 1).contiguous().float()
|
||||
return image, alpha
|
||||
|
||||
def __getitem__(self, index):
|
||||
while True:
|
||||
input_image_path = os.path.join(self.root_dir, self.input_image_dir, self.paths[index])
|
||||
target_image_path = os.path.join(self.root_dir, self.target_image_dir, self.paths[index])
|
||||
|
||||
indices = np.random.choice(range(self.total_view_n), self.input_view_num + self.target_view_num, replace=False)
|
||||
input_indices = indices[:self.input_view_num]
|
||||
target_indices = indices[self.input_view_num:]
|
||||
|
||||
'''background color, default: white'''
|
||||
bg_white = [1., 1., 1.]
|
||||
bg_black = [0., 0., 0.]
|
||||
|
||||
image_list = []
|
||||
alpha_list = []
|
||||
depth_list = []
|
||||
normal_list = []
|
||||
pose_list = []
|
||||
|
||||
try:
|
||||
input_cameras = np.load(os.path.join(input_image_path, 'cameras.npz'))['cam_poses']
|
||||
for idx in input_indices:
|
||||
image, alpha = self.load_im(os.path.join(input_image_path, '%03d.png' % idx), bg_white)
|
||||
normal, _ = self.load_im(os.path.join(input_image_path, '%03d_normal.png' % idx), bg_black)
|
||||
depth = cv2.imread(os.path.join(input_image_path, '%03d_depth.png' % idx), cv2.IMREAD_UNCHANGED) / 255.0 * self.depth_scale
|
||||
depth = torch.from_numpy(depth).unsqueeze(0)
|
||||
pose = input_cameras[idx]
|
||||
pose = np.concatenate([pose, np.array([[0, 0, 0, 1]])], axis=0)
|
||||
|
||||
image_list.append(image)
|
||||
alpha_list.append(alpha)
|
||||
depth_list.append(depth)
|
||||
normal_list.append(normal)
|
||||
pose_list.append(pose)
|
||||
|
||||
target_cameras = np.load(os.path.join(target_image_path, 'cameras.npz'))['cam_poses']
|
||||
for idx in target_indices:
|
||||
image, alpha = self.load_im(os.path.join(target_image_path, '%03d.png' % idx), bg_white)
|
||||
normal, _ = self.load_im(os.path.join(target_image_path, '%03d_normal.png' % idx), bg_black)
|
||||
depth = cv2.imread(os.path.join(target_image_path, '%03d_depth.png' % idx), cv2.IMREAD_UNCHANGED) / 255.0 * self.depth_scale
|
||||
depth = torch.from_numpy(depth).unsqueeze(0)
|
||||
pose = target_cameras[idx]
|
||||
pose = np.concatenate([pose, np.array([[0, 0, 0, 1]])], axis=0)
|
||||
|
||||
image_list.append(image)
|
||||
alpha_list.append(alpha)
|
||||
depth_list.append(depth)
|
||||
normal_list.append(normal)
|
||||
pose_list.append(pose)
|
||||
|
||||
except Exception as e:
|
||||
print(e)
|
||||
index = np.random.randint(0, len(self.paths))
|
||||
continue
|
||||
|
||||
break
|
||||
|
||||
images = torch.stack(image_list, dim=0).float() # (6+V, 3, H, W)
|
||||
alphas = torch.stack(alpha_list, dim=0).float() # (6+V, 1, H, W)
|
||||
depths = torch.stack(depth_list, dim=0).float() # (6+V, 1, H, W)
|
||||
normals = torch.stack(normal_list, dim=0).float() # (6+V, 3, H, W)
|
||||
w2cs = torch.from_numpy(np.stack(pose_list, axis=0)).float() # (6+V, 4, 4)
|
||||
c2ws = torch.linalg.inv(w2cs).float()
|
||||
|
||||
normals = normals * 2.0 - 1.0
|
||||
normals = F.normalize(normals, dim=1)
|
||||
normals = (normals + 1.0) / 2.0
|
||||
normals = torch.lerp(torch.zeros_like(normals), normals, alphas)
|
||||
|
||||
# random rotation along z axis
|
||||
if self.camera_rotation:
|
||||
degree = np.random.uniform(0, math.pi * 2)
|
||||
rot = torch.tensor([
|
||||
[np.cos(degree), -np.sin(degree), 0, 0],
|
||||
[np.sin(degree), np.cos(degree), 0, 0],
|
||||
[0, 0, 1, 0],
|
||||
[0, 0, 0, 1],
|
||||
]).unsqueeze(0).float()
|
||||
c2ws = torch.matmul(rot, c2ws)
|
||||
|
||||
# rotate normals
|
||||
N, _, H, W = normals.shape
|
||||
normals = normals * 2.0 - 1.0
|
||||
normals = torch.matmul(rot[:, :3, :3], normals.view(N, 3, -1)).view(N, 3, H, W)
|
||||
normals = F.normalize(normals, dim=1)
|
||||
normals = (normals + 1.0) / 2.0
|
||||
normals = torch.lerp(torch.zeros_like(normals), normals, alphas)
|
||||
|
||||
# random scaling
|
||||
if np.random.rand() < 0.5:
|
||||
scale = np.random.uniform(0.7, 1.1)
|
||||
c2ws[:, :3, 3] *= scale
|
||||
depths *= scale
|
||||
|
||||
# instrinsics of perspective cameras
|
||||
K = FOV_to_intrinsics(self.fov)
|
||||
Ks = K.unsqueeze(0).repeat(self.input_view_num + self.target_view_num, 1, 1).float()
|
||||
|
||||
data = {
|
||||
'input_images': images[:self.input_view_num], # (6, 3, H, W)
|
||||
'input_alphas': alphas[:self.input_view_num], # (6, 1, H, W)
|
||||
'input_depths': depths[:self.input_view_num], # (6, 1, H, W)
|
||||
'input_normals': normals[:self.input_view_num], # (6, 3, H, W)
|
||||
'input_c2ws': c2ws[:self.input_view_num], # (6, 4, 4)
|
||||
'input_Ks': Ks[:self.input_view_num], # (6, 3, 3)
|
||||
|
||||
# lrm generator input and supervision
|
||||
'target_images': images[self.input_view_num:], # (V, 3, H, W)
|
||||
'target_alphas': alphas[self.input_view_num:], # (V, 1, H, W)
|
||||
'target_depths': depths[self.input_view_num:], # (V, 1, H, W)
|
||||
'target_normals': normals[self.input_view_num:], # (V, 3, H, W)
|
||||
'target_c2ws': c2ws[self.input_view_num:], # (V, 4, 4)
|
||||
'target_Ks': Ks[self.input_view_num:], # (V, 3, 3)
|
||||
}
|
||||
return data
|
||||
|
||||
|
||||
class ValidationData(Dataset):
|
||||
def __init__(self,
|
||||
root_dir='objaverse/',
|
||||
input_view_num=6,
|
||||
input_image_size=320,
|
||||
fov=30,
|
||||
):
|
||||
self.root_dir = Path(root_dir)
|
||||
self.input_view_num = input_view_num
|
||||
self.input_image_size = input_image_size
|
||||
self.fov = fov
|
||||
|
||||
self.paths = sorted(os.listdir(self.root_dir))
|
||||
print('============= length of dataset %d =============' % len(self.paths))
|
||||
|
||||
cam_distance = 4.0
|
||||
azimuths = np.array([30, 90, 150, 210, 270, 330])
|
||||
elevations = np.array([20, -10, 20, -10, 20, -10])
|
||||
azimuths = np.deg2rad(azimuths)
|
||||
elevations = np.deg2rad(elevations)
|
||||
|
||||
x = cam_distance * np.cos(elevations) * np.cos(azimuths)
|
||||
y = cam_distance * np.cos(elevations) * np.sin(azimuths)
|
||||
z = cam_distance * np.sin(elevations)
|
||||
|
||||
cam_locations = np.stack([x, y, z], axis=-1)
|
||||
cam_locations = torch.from_numpy(cam_locations).float()
|
||||
c2ws = center_looking_at_camera_pose(cam_locations)
|
||||
self.c2ws = c2ws.float()
|
||||
self.Ks = FOV_to_intrinsics(self.fov).unsqueeze(0).repeat(6, 1, 1).float()
|
||||
|
||||
render_c2ws = get_circular_camera_poses(M=8, radius=cam_distance, elevation=20.0)
|
||||
render_Ks = FOV_to_intrinsics(self.fov).unsqueeze(0).repeat(render_c2ws.shape[0], 1, 1)
|
||||
self.render_c2ws = render_c2ws.float()
|
||||
self.render_Ks = render_Ks.float()
|
||||
|
||||
def __len__(self):
|
||||
return len(self.paths)
|
||||
|
||||
def load_im(self, path, color):
|
||||
'''
|
||||
replace background pixel with random color in rendering
|
||||
'''
|
||||
pil_img = Image.open(path)
|
||||
pil_img = pil_img.resize((self.input_image_size, self.input_image_size), resample=Image.BICUBIC)
|
||||
|
||||
image = np.asarray(pil_img, dtype=np.float32) / 255.
|
||||
if image.shape[-1] == 4:
|
||||
alpha = image[:, :, 3:]
|
||||
image = image[:, :, :3] * alpha + color * (1 - alpha)
|
||||
else:
|
||||
alpha = np.ones_like(image[:, :, :1])
|
||||
|
||||
image = torch.from_numpy(image).permute(2, 0, 1).contiguous().float()
|
||||
alpha = torch.from_numpy(alpha).permute(2, 0, 1).contiguous().float()
|
||||
return image, alpha
|
||||
|
||||
def __getitem__(self, index):
|
||||
# load data
|
||||
input_image_path = os.path.join(self.root_dir, self.paths[index])
|
||||
|
||||
'''background color, default: white'''
|
||||
bkg_color = [1.0, 1.0, 1.0]
|
||||
|
||||
image_list = []
|
||||
alpha_list = []
|
||||
|
||||
for idx in range(self.input_view_num):
|
||||
image, alpha = self.load_im(os.path.join(input_image_path, f'{idx:03d}.png'), bkg_color)
|
||||
image_list.append(image)
|
||||
alpha_list.append(alpha)
|
||||
|
||||
images = torch.stack(image_list, dim=0).float()
|
||||
alphas = torch.stack(alpha_list, dim=0).float()
|
||||
|
||||
data = {
|
||||
'input_images': images,
|
||||
'input_alphas': alphas,
|
||||
'input_c2ws': self.c2ws,
|
||||
'input_Ks': self.Ks,
|
||||
|
||||
'render_c2ws': self.render_c2ws,
|
||||
'render_Ks': self.render_Ks,
|
||||
}
|
||||
return data
|
||||
@@ -0,0 +1,124 @@
|
||||
import os
|
||||
import json
|
||||
import numpy as np
|
||||
import webdataset as wds
|
||||
import pytorch_lightning as pl
|
||||
import torch
|
||||
from torch.utils.data import Dataset
|
||||
from torch.utils.data.distributed import DistributedSampler
|
||||
from PIL import Image
|
||||
from pathlib import Path
|
||||
|
||||
from instant_mesh.utils.train_util import instantiate_from_config
|
||||
|
||||
|
||||
class DataModuleFromConfig(pl.LightningDataModule):
|
||||
def __init__(
|
||||
self,
|
||||
batch_size=8,
|
||||
num_workers=4,
|
||||
train=None,
|
||||
validation=None,
|
||||
test=None,
|
||||
**kwargs,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.batch_size = batch_size
|
||||
self.num_workers = num_workers
|
||||
|
||||
self.dataset_configs = dict()
|
||||
if train is not None:
|
||||
self.dataset_configs['train'] = train
|
||||
if validation is not None:
|
||||
self.dataset_configs['validation'] = validation
|
||||
if test is not None:
|
||||
self.dataset_configs['test'] = test
|
||||
|
||||
def setup(self, stage):
|
||||
|
||||
if stage in ['fit']:
|
||||
self.datasets = dict((k, instantiate_from_config(self.dataset_configs[k])) for k in self.dataset_configs)
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
def train_dataloader(self):
|
||||
|
||||
sampler = DistributedSampler(self.datasets['train'])
|
||||
return wds.WebLoader(self.datasets['train'], batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False, sampler=sampler)
|
||||
|
||||
def val_dataloader(self):
|
||||
|
||||
sampler = DistributedSampler(self.datasets['validation'])
|
||||
return wds.WebLoader(self.datasets['validation'], batch_size=4, num_workers=self.num_workers, shuffle=False, sampler=sampler)
|
||||
|
||||
def test_dataloader(self):
|
||||
|
||||
return wds.WebLoader(self.datasets['test'], batch_size=self.batch_size, num_workers=self.num_workers, shuffle=False)
|
||||
|
||||
|
||||
class ObjaverseData(Dataset):
|
||||
def __init__(self,
|
||||
root_dir='objaverse/',
|
||||
meta_fname='valid_paths.json',
|
||||
image_dir='rendering_zero123plus',
|
||||
validation=False,
|
||||
):
|
||||
self.root_dir = Path(root_dir)
|
||||
self.image_dir = image_dir
|
||||
|
||||
with open(os.path.join(root_dir, meta_fname)) as f:
|
||||
lvis_dict = json.load(f)
|
||||
paths = []
|
||||
for k in lvis_dict.keys():
|
||||
paths.extend(lvis_dict[k])
|
||||
self.paths = paths
|
||||
|
||||
total_objects = len(self.paths)
|
||||
if validation:
|
||||
self.paths = self.paths[-16:] # used last 16 as validation
|
||||
else:
|
||||
self.paths = self.paths[:-16]
|
||||
print('============= length of dataset %d =============' % len(self.paths))
|
||||
|
||||
def __len__(self):
|
||||
return len(self.paths)
|
||||
|
||||
def load_im(self, path, color):
|
||||
pil_img = Image.open(path)
|
||||
|
||||
image = np.asarray(pil_img, dtype=np.float32) / 255.
|
||||
alpha = image[:, :, 3:]
|
||||
image = image[:, :, :3] * alpha + color * (1 - alpha)
|
||||
|
||||
image = torch.from_numpy(image).permute(2, 0, 1).contiguous().float()
|
||||
alpha = torch.from_numpy(alpha).permute(2, 0, 1).contiguous().float()
|
||||
return image, alpha
|
||||
|
||||
def __getitem__(self, index):
|
||||
while True:
|
||||
image_path = os.path.join(self.root_dir, self.image_dir, self.paths[index])
|
||||
|
||||
'''background color, default: white'''
|
||||
bkg_color = [1., 1., 1.]
|
||||
|
||||
img_list = []
|
||||
try:
|
||||
for idx in range(7):
|
||||
img, alpha = self.load_im(os.path.join(image_path, '%03d.png' % idx), bkg_color)
|
||||
img_list.append(img)
|
||||
|
||||
except Exception as e:
|
||||
print(e)
|
||||
index = np.random.randint(0, len(self.paths))
|
||||
continue
|
||||
|
||||
break
|
||||
|
||||
imgs = torch.stack(img_list, dim=0).float()
|
||||
|
||||
data = {
|
||||
'cond_imgs': imgs[0], # (3, H, W)
|
||||
'target_imgs': imgs[1:], # (6, 3, H, W)
|
||||
}
|
||||
return data
|
||||
@@ -0,0 +1,303 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torchvision.transforms import v2
|
||||
from torchvision.utils import make_grid, save_image
|
||||
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
|
||||
import pytorch_lightning as pl
|
||||
from einops import rearrange, repeat
|
||||
|
||||
from instant_mesh.utils.train_util import instantiate_from_config
|
||||
|
||||
|
||||
class MVRecon(pl.LightningModule):
|
||||
def __init__(
|
||||
self,
|
||||
lrm_generator_config,
|
||||
lrm_path=None,
|
||||
input_size=256,
|
||||
render_size=192,
|
||||
):
|
||||
super(MVRecon, self).__init__()
|
||||
|
||||
self.input_size = input_size
|
||||
self.render_size = render_size
|
||||
|
||||
# init modules
|
||||
self.lrm_generator = instantiate_from_config(lrm_generator_config)
|
||||
if lrm_path is not None:
|
||||
lrm_ckpt = torch.load(lrm_path)
|
||||
self.lrm_generator.load_state_dict(lrm_ckpt['weights'], strict=False)
|
||||
|
||||
self.lpips = LearnedPerceptualImagePatchSimilarity(net_type='vgg')
|
||||
|
||||
self.validation_step_outputs = []
|
||||
|
||||
def on_fit_start(self):
|
||||
if self.global_rank == 0:
|
||||
os.makedirs(os.path.join(self.logdir, 'images'), exist_ok=True)
|
||||
os.makedirs(os.path.join(self.logdir, 'images_val'), exist_ok=True)
|
||||
|
||||
def prepare_batch_data(self, batch):
|
||||
lrm_generator_input = {}
|
||||
render_gt = {} # for supervision
|
||||
|
||||
# input images
|
||||
images = batch['input_images']
|
||||
images = v2.functional.resize(
|
||||
images, self.input_size, interpolation=3, antialias=True).clamp(0, 1)
|
||||
|
||||
lrm_generator_input['images'] = images.to(self.device)
|
||||
|
||||
# input cameras and render cameras
|
||||
input_c2ws = batch['input_c2ws'].flatten(-2)
|
||||
input_Ks = batch['input_Ks'].flatten(-2)
|
||||
target_c2ws = batch['target_c2ws'].flatten(-2)
|
||||
target_Ks = batch['target_Ks'].flatten(-2)
|
||||
render_cameras_input = torch.cat([input_c2ws, input_Ks], dim=-1)
|
||||
render_cameras_target = torch.cat([target_c2ws, target_Ks], dim=-1)
|
||||
render_cameras = torch.cat([render_cameras_input, render_cameras_target], dim=1)
|
||||
|
||||
input_extrinsics = input_c2ws[:, :, :12]
|
||||
input_intrinsics = torch.stack([
|
||||
input_Ks[:, :, 0], input_Ks[:, :, 4],
|
||||
input_Ks[:, :, 2], input_Ks[:, :, 5],
|
||||
], dim=-1)
|
||||
cameras = torch.cat([input_extrinsics, input_intrinsics], dim=-1)
|
||||
|
||||
# add noise to input cameras
|
||||
cameras = cameras + torch.rand_like(cameras) * 0.04 - 0.02
|
||||
|
||||
lrm_generator_input['cameras'] = cameras.to(self.device)
|
||||
lrm_generator_input['render_cameras'] = render_cameras.to(self.device)
|
||||
|
||||
# target images
|
||||
target_images = torch.cat([batch['input_images'], batch['target_images']], dim=1)
|
||||
target_depths = torch.cat([batch['input_depths'], batch['target_depths']], dim=1)
|
||||
target_alphas = torch.cat([batch['input_alphas'], batch['target_alphas']], dim=1)
|
||||
|
||||
# random crop
|
||||
render_size = np.random.randint(self.render_size, 513)
|
||||
target_images = v2.functional.resize(
|
||||
target_images, render_size, interpolation=3, antialias=True).clamp(0, 1)
|
||||
target_depths = v2.functional.resize(
|
||||
target_depths, render_size, interpolation=0, antialias=True)
|
||||
target_alphas = v2.functional.resize(
|
||||
target_alphas, render_size, interpolation=0, antialias=True)
|
||||
|
||||
crop_params = v2.RandomCrop.get_params(
|
||||
target_images, output_size=(self.render_size, self.render_size))
|
||||
target_images = v2.functional.crop(target_images, *crop_params)
|
||||
target_depths = v2.functional.crop(target_depths, *crop_params)[:, :, 0:1]
|
||||
target_alphas = v2.functional.crop(target_alphas, *crop_params)[:, :, 0:1]
|
||||
|
||||
lrm_generator_input['render_size'] = render_size
|
||||
lrm_generator_input['crop_params'] = crop_params
|
||||
|
||||
render_gt['target_images'] = target_images.to(self.device)
|
||||
render_gt['target_depths'] = target_depths.to(self.device)
|
||||
render_gt['target_alphas'] = target_alphas.to(self.device)
|
||||
|
||||
return lrm_generator_input, render_gt
|
||||
|
||||
def prepare_validation_batch_data(self, batch):
|
||||
lrm_generator_input = {}
|
||||
|
||||
# input images
|
||||
images = batch['input_images']
|
||||
images = v2.functional.resize(
|
||||
images, self.input_size, interpolation=3, antialias=True).clamp(0, 1)
|
||||
|
||||
lrm_generator_input['images'] = images.to(self.device)
|
||||
|
||||
input_c2ws = batch['input_c2ws'].flatten(-2)
|
||||
input_Ks = batch['input_Ks'].flatten(-2)
|
||||
|
||||
input_extrinsics = input_c2ws[:, :, :12]
|
||||
input_intrinsics = torch.stack([
|
||||
input_Ks[:, :, 0], input_Ks[:, :, 4],
|
||||
input_Ks[:, :, 2], input_Ks[:, :, 5],
|
||||
], dim=-1)
|
||||
cameras = torch.cat([input_extrinsics, input_intrinsics], dim=-1)
|
||||
|
||||
lrm_generator_input['cameras'] = cameras.to(self.device)
|
||||
|
||||
render_c2ws = batch['render_c2ws'].flatten(-2)
|
||||
render_Ks = batch['render_Ks'].flatten(-2)
|
||||
render_cameras = torch.cat([render_c2ws, render_Ks], dim=-1)
|
||||
|
||||
lrm_generator_input['render_cameras'] = render_cameras.to(self.device)
|
||||
lrm_generator_input['render_size'] = 384
|
||||
lrm_generator_input['crop_params'] = None
|
||||
|
||||
return lrm_generator_input
|
||||
|
||||
def forward_lrm_generator(
|
||||
self,
|
||||
images,
|
||||
cameras,
|
||||
render_cameras,
|
||||
render_size=192,
|
||||
crop_params=None,
|
||||
chunk_size=1,
|
||||
):
|
||||
planes = torch.utils.checkpoint.checkpoint(
|
||||
self.lrm_generator.forward_planes,
|
||||
images,
|
||||
cameras,
|
||||
use_reentrant=False,
|
||||
)
|
||||
frames = []
|
||||
for i in range(0, render_cameras.shape[1], chunk_size):
|
||||
frames.append(
|
||||
torch.utils.checkpoint.checkpoint(
|
||||
self.lrm_generator.synthesizer,
|
||||
planes,
|
||||
cameras=render_cameras[:, i:i+chunk_size],
|
||||
render_size=render_size,
|
||||
crop_params=crop_params,
|
||||
use_reentrant=False
|
||||
)
|
||||
)
|
||||
frames = {
|
||||
k: torch.cat([r[k] for r in frames], dim=1)
|
||||
for k in frames[0].keys()
|
||||
}
|
||||
return frames
|
||||
|
||||
def forward(self, lrm_generator_input):
|
||||
images = lrm_generator_input['images']
|
||||
cameras = lrm_generator_input['cameras']
|
||||
render_cameras = lrm_generator_input['render_cameras']
|
||||
render_size = lrm_generator_input['render_size']
|
||||
crop_params = lrm_generator_input['crop_params']
|
||||
|
||||
out = self.forward_lrm_generator(
|
||||
images,
|
||||
cameras,
|
||||
render_cameras,
|
||||
render_size=render_size,
|
||||
crop_params=crop_params,
|
||||
chunk_size=1,
|
||||
)
|
||||
render_images = torch.clamp(out['images_rgb'], 0.0, 1.0)
|
||||
render_depths = out['images_depth']
|
||||
render_alphas = torch.clamp(out['images_weight'], 0.0, 1.0)
|
||||
|
||||
out = {
|
||||
'render_images': render_images,
|
||||
'render_depths': render_depths,
|
||||
'render_alphas': render_alphas,
|
||||
}
|
||||
return out
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
lrm_generator_input, render_gt = self.prepare_batch_data(batch)
|
||||
|
||||
render_out = self.forward(lrm_generator_input)
|
||||
|
||||
loss, loss_dict = self.compute_loss(render_out, render_gt)
|
||||
|
||||
self.log_dict(loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=True)
|
||||
|
||||
if self.global_step % 1000 == 0 and self.global_rank == 0:
|
||||
B, N, C, H, W = render_gt['target_images'].shape
|
||||
N_in = lrm_generator_input['images'].shape[1]
|
||||
|
||||
input_images = v2.functional.resize(
|
||||
lrm_generator_input['images'], (H, W), interpolation=3, antialias=True).clamp(0, 1)
|
||||
input_images = torch.cat(
|
||||
[input_images, torch.ones(B, N-N_in, C, H, W).to(input_images)], dim=1)
|
||||
|
||||
input_images = rearrange(
|
||||
input_images, 'b n c h w -> b c h (n w)')
|
||||
target_images = rearrange(
|
||||
render_gt['target_images'], 'b n c h w -> b c h (n w)')
|
||||
render_images = rearrange(
|
||||
render_out['render_images'], 'b n c h w -> b c h (n w)')
|
||||
target_alphas = rearrange(
|
||||
repeat(render_gt['target_alphas'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)')
|
||||
render_alphas = rearrange(
|
||||
repeat(render_out['render_alphas'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)')
|
||||
target_depths = rearrange(
|
||||
repeat(render_gt['target_depths'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)')
|
||||
render_depths = rearrange(
|
||||
repeat(render_out['render_depths'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)')
|
||||
MAX_DEPTH = torch.max(target_depths)
|
||||
target_depths = target_depths / MAX_DEPTH * target_alphas
|
||||
render_depths = render_depths / MAX_DEPTH
|
||||
|
||||
grid = torch.cat([
|
||||
input_images,
|
||||
target_images, render_images,
|
||||
target_alphas, render_alphas,
|
||||
target_depths, render_depths,
|
||||
], dim=-2)
|
||||
grid = make_grid(grid, nrow=target_images.shape[0], normalize=True, value_range=(0, 1))
|
||||
|
||||
save_image(grid, os.path.join(self.logdir, 'images', f'train_{self.global_step:07d}.png'))
|
||||
|
||||
return loss
|
||||
|
||||
def compute_loss(self, render_out, render_gt):
|
||||
# NOTE: the rgb value range of OpenLRM is [0, 1]
|
||||
render_images = render_out['render_images']
|
||||
target_images = render_gt['target_images'].to(render_images)
|
||||
render_images = rearrange(render_images, 'b n ... -> (b n) ...') * 2.0 - 1.0
|
||||
target_images = rearrange(target_images, 'b n ... -> (b n) ...') * 2.0 - 1.0
|
||||
|
||||
loss_mse = F.mse_loss(render_images, target_images)
|
||||
loss_lpips = 2.0 * self.lpips(render_images, target_images)
|
||||
|
||||
render_alphas = render_out['render_alphas']
|
||||
target_alphas = render_gt['target_alphas']
|
||||
loss_mask = F.mse_loss(render_alphas, target_alphas)
|
||||
|
||||
loss = loss_mse + loss_lpips + loss_mask
|
||||
|
||||
prefix = 'train'
|
||||
loss_dict = {}
|
||||
loss_dict.update({f'{prefix}/loss_mse': loss_mse})
|
||||
loss_dict.update({f'{prefix}/loss_lpips': loss_lpips})
|
||||
loss_dict.update({f'{prefix}/loss_mask': loss_mask})
|
||||
loss_dict.update({f'{prefix}/loss': loss})
|
||||
|
||||
return loss, loss_dict
|
||||
|
||||
@torch.no_grad()
|
||||
def validation_step(self, batch, batch_idx):
|
||||
lrm_generator_input = self.prepare_validation_batch_data(batch)
|
||||
|
||||
render_out = self.forward(lrm_generator_input)
|
||||
render_images = render_out['render_images']
|
||||
render_images = rearrange(render_images, 'b n c h w -> b c h (n w)')
|
||||
|
||||
self.validation_step_outputs.append(render_images)
|
||||
|
||||
def on_validation_epoch_end(self):
|
||||
images = torch.cat(self.validation_step_outputs, dim=-1)
|
||||
|
||||
all_images = self.all_gather(images)
|
||||
all_images = rearrange(all_images, 'r b c h w -> (r b) c h w')
|
||||
|
||||
if self.global_rank == 0:
|
||||
image_path = os.path.join(self.logdir, 'images_val', f'val_{self.global_step:07d}.png')
|
||||
|
||||
grid = make_grid(all_images, nrow=1, normalize=True, value_range=(0, 1))
|
||||
save_image(grid, image_path)
|
||||
print(f"Saved image to {image_path}")
|
||||
|
||||
self.validation_step_outputs.clear()
|
||||
|
||||
def configure_optimizers(self):
|
||||
lr = self.learning_rate
|
||||
|
||||
params = []
|
||||
|
||||
params.append({"params": self.lrm_generator.parameters(), "lr": lr, "weight_decay": 0.01 })
|
||||
|
||||
optimizer = torch.optim.AdamW(params, lr=lr, betas=(0.90, 0.95))
|
||||
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 3000, eta_min=lr/10)
|
||||
|
||||
return {'optimizer': optimizer, 'lr_scheduler': scheduler}
|
||||
@@ -0,0 +1,325 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
from torchvision.transforms import v2
|
||||
from torchvision.utils import make_grid, save_image
|
||||
from torchmetrics.image.lpip import LearnedPerceptualImagePatchSimilarity
|
||||
import pytorch_lightning as pl
|
||||
from einops import rearrange, repeat
|
||||
|
||||
from instant_mesh.utils.train_util import instantiate_from_config
|
||||
|
||||
|
||||
# Regulrarization loss for FlexiCubes
|
||||
def sdf_reg_loss_batch(sdf, all_edges):
|
||||
sdf_f1x6x2 = sdf[:, all_edges.reshape(-1)].reshape(sdf.shape[0], -1, 2)
|
||||
mask = torch.sign(sdf_f1x6x2[..., 0]) != torch.sign(sdf_f1x6x2[..., 1])
|
||||
sdf_f1x6x2 = sdf_f1x6x2[mask]
|
||||
sdf_diff = F.binary_cross_entropy_with_logits(
|
||||
sdf_f1x6x2[..., 0], (sdf_f1x6x2[..., 1] > 0).float()) + \
|
||||
F.binary_cross_entropy_with_logits(
|
||||
sdf_f1x6x2[..., 1], (sdf_f1x6x2[..., 0] > 0).float())
|
||||
return sdf_diff
|
||||
|
||||
|
||||
class MVRecon(pl.LightningModule):
|
||||
def __init__(
|
||||
self,
|
||||
lrm_generator_config,
|
||||
input_size=256,
|
||||
render_size=512,
|
||||
init_ckpt=None,
|
||||
):
|
||||
super(MVRecon, self).__init__()
|
||||
|
||||
self.input_size = input_size
|
||||
self.render_size = render_size
|
||||
|
||||
# init modules
|
||||
self.lrm_generator = instantiate_from_config(lrm_generator_config)
|
||||
|
||||
self.lpips = LearnedPerceptualImagePatchSimilarity(net_type='vgg')
|
||||
|
||||
# Load weights from pretrained MVRecon model, and use the mlp
|
||||
# weights to initialize the weights of sdf and rgb mlps.
|
||||
if init_ckpt is not None:
|
||||
sd = torch.load(init_ckpt, map_location='cpu')['state_dict']
|
||||
sd = {k: v for k, v in sd.items() if k.startswith('lrm_generator')}
|
||||
sd_fc = {}
|
||||
for k, v in sd.items():
|
||||
if k.startswith('lrm_generator.synthesizer.decoder.net.'):
|
||||
if k.startswith('lrm_generator.synthesizer.decoder.net.6.'): # last layer
|
||||
# Here we assume the density filed's isosurface threshold is t,
|
||||
# we reverse the sign of density filed to initialize SDF field.
|
||||
# -(w*x + b - t) = (-w)*x + (t - b)
|
||||
if 'weight' in k:
|
||||
sd_fc[k.replace('net.', 'net_sdf.')] = -v[0:1]
|
||||
else:
|
||||
sd_fc[k.replace('net.', 'net_sdf.')] = 10.0 - v[0:1]
|
||||
sd_fc[k.replace('net.', 'net_rgb.')] = v[1:4]
|
||||
else:
|
||||
sd_fc[k.replace('net.', 'net_sdf.')] = v
|
||||
sd_fc[k.replace('net.', 'net_rgb.')] = v
|
||||
else:
|
||||
sd_fc[k] = v
|
||||
sd_fc = {k.replace('lrm_generator.', ''): v for k, v in sd_fc.items()}
|
||||
# missing `net_deformation` and `net_weight` parameters
|
||||
self.lrm_generator.load_state_dict(sd_fc, strict=False)
|
||||
print(f'Loaded weights from {init_ckpt}')
|
||||
|
||||
self.validation_step_outputs = []
|
||||
|
||||
def on_fit_start(self):
|
||||
device = torch.device(f'cuda:{self.global_rank}')
|
||||
self.lrm_generator.init_flexicubes_geometry(device)
|
||||
if self.global_rank == 0:
|
||||
os.makedirs(os.path.join(self.logdir, 'images'), exist_ok=True)
|
||||
os.makedirs(os.path.join(self.logdir, 'images_val'), exist_ok=True)
|
||||
|
||||
def prepare_batch_data(self, batch):
|
||||
lrm_generator_input = {}
|
||||
render_gt = {}
|
||||
|
||||
# input images
|
||||
images = batch['input_images']
|
||||
images = v2.functional.resize(
|
||||
images, self.input_size, interpolation=3, antialias=True).clamp(0, 1)
|
||||
|
||||
lrm_generator_input['images'] = images.to(self.device)
|
||||
|
||||
# input cameras and render cameras
|
||||
input_c2ws = batch['input_c2ws']
|
||||
input_Ks = batch['input_Ks']
|
||||
target_c2ws = batch['target_c2ws']
|
||||
|
||||
render_c2ws = torch.cat([input_c2ws, target_c2ws], dim=1)
|
||||
render_w2cs = torch.linalg.inv(render_c2ws)
|
||||
|
||||
input_extrinsics = input_c2ws.flatten(-2)
|
||||
input_extrinsics = input_extrinsics[:, :, :12]
|
||||
input_intrinsics = input_Ks.flatten(-2)
|
||||
input_intrinsics = torch.stack([
|
||||
input_intrinsics[:, :, 0], input_intrinsics[:, :, 4],
|
||||
input_intrinsics[:, :, 2], input_intrinsics[:, :, 5],
|
||||
], dim=-1)
|
||||
cameras = torch.cat([input_extrinsics, input_intrinsics], dim=-1)
|
||||
|
||||
# add noise to input_cameras
|
||||
cameras = cameras + torch.rand_like(cameras) * 0.04 - 0.02
|
||||
|
||||
lrm_generator_input['cameras'] = cameras.to(self.device)
|
||||
lrm_generator_input['render_cameras'] = render_w2cs.to(self.device)
|
||||
|
||||
# target images
|
||||
target_images = torch.cat([batch['input_images'], batch['target_images']], dim=1)
|
||||
target_depths = torch.cat([batch['input_depths'], batch['target_depths']], dim=1)
|
||||
target_alphas = torch.cat([batch['input_alphas'], batch['target_alphas']], dim=1)
|
||||
target_normals = torch.cat([batch['input_normals'], batch['target_normals']], dim=1)
|
||||
|
||||
render_size = self.render_size
|
||||
target_images = v2.functional.resize(
|
||||
target_images, render_size, interpolation=3, antialias=True).clamp(0, 1)
|
||||
target_depths = v2.functional.resize(
|
||||
target_depths, render_size, interpolation=0, antialias=True)
|
||||
target_alphas = v2.functional.resize(
|
||||
target_alphas, render_size, interpolation=0, antialias=True)
|
||||
target_normals = v2.functional.resize(
|
||||
target_normals, render_size, interpolation=3, antialias=True)
|
||||
|
||||
lrm_generator_input['render_size'] = render_size
|
||||
|
||||
render_gt['target_images'] = target_images.to(self.device)
|
||||
render_gt['target_depths'] = target_depths.to(self.device)
|
||||
render_gt['target_alphas'] = target_alphas.to(self.device)
|
||||
render_gt['target_normals'] = target_normals.to(self.device)
|
||||
|
||||
return lrm_generator_input, render_gt
|
||||
|
||||
def prepare_validation_batch_data(self, batch):
|
||||
lrm_generator_input = {}
|
||||
|
||||
# input images
|
||||
images = batch['input_images']
|
||||
images = v2.functional.resize(
|
||||
images, self.input_size, interpolation=3, antialias=True).clamp(0, 1)
|
||||
|
||||
lrm_generator_input['images'] = images.to(self.device)
|
||||
|
||||
# input cameras
|
||||
input_c2ws = batch['input_c2ws'].flatten(-2)
|
||||
input_Ks = batch['input_Ks'].flatten(-2)
|
||||
|
||||
input_extrinsics = input_c2ws[:, :, :12]
|
||||
input_intrinsics = torch.stack([
|
||||
input_Ks[:, :, 0], input_Ks[:, :, 4],
|
||||
input_Ks[:, :, 2], input_Ks[:, :, 5],
|
||||
], dim=-1)
|
||||
cameras = torch.cat([input_extrinsics, input_intrinsics], dim=-1)
|
||||
|
||||
lrm_generator_input['cameras'] = cameras.to(self.device)
|
||||
|
||||
# render cameras
|
||||
render_c2ws = batch['render_c2ws']
|
||||
render_w2cs = torch.linalg.inv(render_c2ws)
|
||||
|
||||
lrm_generator_input['render_cameras'] = render_w2cs.to(self.device)
|
||||
lrm_generator_input['render_size'] = 384
|
||||
|
||||
return lrm_generator_input
|
||||
|
||||
def forward_lrm_generator(self, images, cameras, render_cameras, render_size=512):
|
||||
planes = torch.utils.checkpoint.checkpoint(
|
||||
self.lrm_generator.forward_planes,
|
||||
images,
|
||||
cameras,
|
||||
use_reentrant=False,
|
||||
)
|
||||
out = self.lrm_generator.forward_geometry(
|
||||
planes,
|
||||
render_cameras,
|
||||
render_size,
|
||||
)
|
||||
return out
|
||||
|
||||
def forward(self, lrm_generator_input):
|
||||
images = lrm_generator_input['images']
|
||||
cameras = lrm_generator_input['cameras']
|
||||
render_cameras = lrm_generator_input['render_cameras']
|
||||
render_size = lrm_generator_input['render_size']
|
||||
|
||||
out = self.forward_lrm_generator(
|
||||
images, cameras, render_cameras, render_size=render_size)
|
||||
|
||||
return out
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
lrm_generator_input, render_gt = self.prepare_batch_data(batch)
|
||||
|
||||
render_out = self.forward(lrm_generator_input)
|
||||
|
||||
loss, loss_dict = self.compute_loss(render_out, render_gt)
|
||||
|
||||
self.log_dict(loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=True)
|
||||
|
||||
if self.global_step % 1000 == 0 and self.global_rank == 0:
|
||||
B, N, C, H, W = render_gt['target_images'].shape
|
||||
N_in = lrm_generator_input['images'].shape[1]
|
||||
|
||||
target_images = rearrange(
|
||||
render_gt['target_images'], 'b n c h w -> b c h (n w)')
|
||||
render_images = rearrange(
|
||||
render_out['img'], 'b n c h w -> b c h (n w)')
|
||||
target_alphas = rearrange(
|
||||
repeat(render_gt['target_alphas'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)')
|
||||
render_alphas = rearrange(
|
||||
repeat(render_out['mask'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)')
|
||||
target_depths = rearrange(
|
||||
repeat(render_gt['target_depths'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)')
|
||||
render_depths = rearrange(
|
||||
repeat(render_out['depth'], 'b n 1 h w -> b n 3 h w'), 'b n c h w -> b c h (n w)')
|
||||
target_normals = rearrange(
|
||||
render_gt['target_normals'], 'b n c h w -> b c h (n w)')
|
||||
render_normals = rearrange(
|
||||
render_out['normal'], 'b n c h w -> b c h (n w)')
|
||||
MAX_DEPTH = torch.max(target_depths)
|
||||
target_depths = target_depths / MAX_DEPTH * target_alphas
|
||||
render_depths = render_depths / MAX_DEPTH
|
||||
|
||||
grid = torch.cat([
|
||||
target_images, render_images,
|
||||
target_alphas, render_alphas,
|
||||
target_depths, render_depths,
|
||||
target_normals, render_normals,
|
||||
], dim=-2)
|
||||
grid = make_grid(grid, nrow=target_images.shape[0], normalize=True, value_range=(0, 1))
|
||||
|
||||
image_path = os.path.join(self.logdir, 'images', f'train_{self.global_step:07d}.png')
|
||||
save_image(grid, image_path)
|
||||
print(f"Saved image to {image_path}")
|
||||
|
||||
return loss
|
||||
|
||||
def compute_loss(self, render_out, render_gt):
|
||||
# NOTE: the rgb value range of OpenLRM is [0, 1]
|
||||
render_images = render_out['img']
|
||||
target_images = render_gt['target_images'].to(render_images)
|
||||
render_images = rearrange(render_images, 'b n ... -> (b n) ...') * 2.0 - 1.0
|
||||
target_images = rearrange(target_images, 'b n ... -> (b n) ...') * 2.0 - 1.0
|
||||
loss_mse = F.mse_loss(render_images, target_images)
|
||||
loss_lpips = 2.0 * self.lpips(render_images, target_images)
|
||||
|
||||
render_alphas = render_out['mask']
|
||||
target_alphas = render_gt['target_alphas']
|
||||
loss_mask = F.mse_loss(render_alphas, target_alphas)
|
||||
|
||||
render_depths = render_out['depth']
|
||||
target_depths = render_gt['target_depths']
|
||||
loss_depth = 0.5 * F.l1_loss(render_depths[target_alphas>0], target_depths[target_alphas>0])
|
||||
|
||||
render_normals = render_out['normal'] * 2.0 - 1.0
|
||||
target_normals = render_gt['target_normals'] * 2.0 - 1.0
|
||||
similarity = (render_normals * target_normals).sum(dim=-3).abs()
|
||||
normal_mask = target_alphas.squeeze(-3)
|
||||
loss_normal = 1 - similarity[normal_mask>0].mean()
|
||||
loss_normal = 0.2 * loss_normal
|
||||
|
||||
# flexicubes regularization loss
|
||||
sdf = render_out['sdf']
|
||||
sdf_reg_loss = render_out['sdf_reg_loss']
|
||||
sdf_reg_loss_entropy = sdf_reg_loss_batch(sdf, self.lrm_generator.geometry.all_edges).mean() * 0.01
|
||||
_, flexicubes_surface_reg, flexicubes_weights_reg = sdf_reg_loss
|
||||
flexicubes_surface_reg = flexicubes_surface_reg.mean() * 0.5
|
||||
flexicubes_weights_reg = flexicubes_weights_reg.mean() * 0.1
|
||||
|
||||
loss_reg = sdf_reg_loss_entropy + flexicubes_surface_reg + flexicubes_weights_reg
|
||||
|
||||
loss = loss_mse + loss_lpips + loss_mask + loss_depth + loss_normal + loss_reg
|
||||
|
||||
prefix = 'train'
|
||||
loss_dict = {}
|
||||
loss_dict.update({f'{prefix}/loss_mse': loss_mse})
|
||||
loss_dict.update({f'{prefix}/loss_lpips': loss_lpips})
|
||||
loss_dict.update({f'{prefix}/loss_mask': loss_mask})
|
||||
loss_dict.update({f'{prefix}/loss_normal': loss_normal})
|
||||
loss_dict.update({f'{prefix}/loss_depth': loss_depth})
|
||||
loss_dict.update({f'{prefix}/loss_reg_sdf': sdf_reg_loss_entropy})
|
||||
loss_dict.update({f'{prefix}/loss_reg_surface': flexicubes_surface_reg})
|
||||
loss_dict.update({f'{prefix}/loss_reg_weights': flexicubes_weights_reg})
|
||||
loss_dict.update({f'{prefix}/loss': loss})
|
||||
|
||||
return loss, loss_dict
|
||||
|
||||
@torch.no_grad()
|
||||
def validation_step(self, batch, batch_idx):
|
||||
lrm_generator_input = self.prepare_validation_batch_data(batch)
|
||||
|
||||
render_out = self.forward(lrm_generator_input)
|
||||
render_images = render_out['img']
|
||||
render_images = rearrange(render_images, 'b n c h w -> b c h (n w)')
|
||||
|
||||
self.validation_step_outputs.append(render_images)
|
||||
|
||||
def on_validation_epoch_end(self):
|
||||
images = torch.cat(self.validation_step_outputs, dim=-1)
|
||||
|
||||
all_images = self.all_gather(images)
|
||||
all_images = rearrange(all_images, 'r b c h w -> (r b) c h w')
|
||||
|
||||
if self.global_rank == 0:
|
||||
image_path = os.path.join(self.logdir, 'images_val', f'val_{self.global_step:07d}.png')
|
||||
|
||||
grid = make_grid(all_images, nrow=1, normalize=True, value_range=(0, 1))
|
||||
save_image(grid, image_path)
|
||||
print(f"Saved image to {image_path}")
|
||||
|
||||
self.validation_step_outputs.clear()
|
||||
|
||||
def configure_optimizers(self):
|
||||
lr = self.learning_rate
|
||||
|
||||
optimizer = torch.optim.AdamW(
|
||||
self.lrm_generator.parameters(), lr=lr, betas=(0.90, 0.95), weight_decay=0.01)
|
||||
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 100000, eta_min=0)
|
||||
|
||||
return {'optimizer': optimizer, 'lr_scheduler': scheduler}
|
||||
@@ -0,0 +1,123 @@
|
||||
# Copyright (c) 2023, Zexin He
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# https://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
|
||||
class BasicTransformerBlock(nn.Module):
|
||||
"""
|
||||
Transformer block that takes in a cross-attention condition and another modulation vector applied to sub-blocks.
|
||||
"""
|
||||
# use attention from torch.nn.MultiHeadAttention
|
||||
# Block contains a cross-attention layer, a self-attention layer, and a MLP
|
||||
def __init__(
|
||||
self,
|
||||
inner_dim: int,
|
||||
cond_dim: int,
|
||||
num_heads: int,
|
||||
eps: float,
|
||||
attn_drop: float = 0.,
|
||||
attn_bias: bool = False,
|
||||
mlp_ratio: float = 4.,
|
||||
mlp_drop: float = 0.,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.norm1 = nn.LayerNorm(inner_dim)
|
||||
self.cross_attn = nn.MultiheadAttention(
|
||||
embed_dim=inner_dim, num_heads=num_heads, kdim=cond_dim, vdim=cond_dim,
|
||||
dropout=attn_drop, bias=attn_bias, batch_first=True)
|
||||
self.norm2 = nn.LayerNorm(inner_dim)
|
||||
self.self_attn = nn.MultiheadAttention(
|
||||
embed_dim=inner_dim, num_heads=num_heads,
|
||||
dropout=attn_drop, bias=attn_bias, batch_first=True)
|
||||
self.norm3 = nn.LayerNorm(inner_dim)
|
||||
self.mlp = nn.Sequential(
|
||||
nn.Linear(inner_dim, int(inner_dim * mlp_ratio)),
|
||||
nn.GELU(),
|
||||
nn.Dropout(mlp_drop),
|
||||
nn.Linear(int(inner_dim * mlp_ratio), inner_dim),
|
||||
nn.Dropout(mlp_drop),
|
||||
)
|
||||
|
||||
def forward(self, x, cond):
|
||||
# x: [N, L, D]
|
||||
# cond: [N, L_cond, D_cond]
|
||||
x = x + self.cross_attn(self.norm1(x), cond, cond)[0]
|
||||
before_sa = self.norm2(x)
|
||||
x = x + self.self_attn(before_sa, before_sa, before_sa)[0]
|
||||
x = x + self.mlp(self.norm3(x))
|
||||
return x
|
||||
|
||||
|
||||
class TriplaneTransformer(nn.Module):
|
||||
"""
|
||||
Transformer with condition that generates a triplane representation.
|
||||
|
||||
Reference:
|
||||
Timm: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L486
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
inner_dim: int,
|
||||
image_feat_dim: int,
|
||||
triplane_low_res: int,
|
||||
triplane_high_res: int,
|
||||
triplane_dim: int,
|
||||
num_layers: int,
|
||||
num_heads: int,
|
||||
eps: float = 1e-6,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# attributes
|
||||
self.triplane_low_res = triplane_low_res
|
||||
self.triplane_high_res = triplane_high_res
|
||||
self.triplane_dim = triplane_dim
|
||||
|
||||
# modules
|
||||
# initialize pos_embed with 1/sqrt(dim) * N(0, 1)
|
||||
self.pos_embed = nn.Parameter(torch.randn(1, 3*triplane_low_res**2, inner_dim) * (1. / inner_dim) ** 0.5)
|
||||
self.layers = nn.ModuleList([
|
||||
BasicTransformerBlock(
|
||||
inner_dim=inner_dim, cond_dim=image_feat_dim, num_heads=num_heads, eps=eps)
|
||||
for _ in range(num_layers)
|
||||
])
|
||||
self.norm = nn.LayerNorm(inner_dim, eps=eps)
|
||||
self.deconv = nn.ConvTranspose2d(inner_dim, triplane_dim, kernel_size=2, stride=2, padding=0)
|
||||
|
||||
def forward(self, image_feats):
|
||||
# image_feats: [N, L_cond, D_cond]
|
||||
|
||||
N = image_feats.shape[0]
|
||||
H = W = self.triplane_low_res
|
||||
L = 3 * H * W
|
||||
|
||||
x = self.pos_embed.repeat(N, 1, 1) # [N, L, D]
|
||||
for layer in self.layers:
|
||||
x = layer(x, image_feats)
|
||||
x = self.norm(x)
|
||||
|
||||
# separate each plane and apply deconv
|
||||
x = x.view(N, 3, H, W, -1)
|
||||
x = torch.einsum('nihwd->indhw', x) # [3, N, D, H, W]
|
||||
x = x.contiguous().view(3*N, -1, H, W) # [3*N, D, H, W]
|
||||
x = self.deconv(x) # [3*N, D', H', W']
|
||||
x = x.view(3, N, *x.shape[-3:]) # [3, N, D', H', W']
|
||||
x = torch.einsum('indhw->nidhw', x) # [N, 3, D', H', W']
|
||||
x = x.contiguous()
|
||||
|
||||
return x
|
||||
@@ -0,0 +1,550 @@
|
||||
# coding=utf-8
|
||||
# Copyright 2021 Google AI, Ross Wightman, The HuggingFace Inc. team. All rights reserved.
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# http://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
""" PyTorch ViT model."""
|
||||
|
||||
|
||||
import collections.abc
|
||||
import math
|
||||
from typing import Dict, List, Optional, Set, Tuple, Union
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
from transformers.activations import ACT2FN
|
||||
from transformers.modeling_outputs import (
|
||||
BaseModelOutput,
|
||||
BaseModelOutputWithPooling,
|
||||
)
|
||||
from transformers import PreTrainedModel, ViTConfig
|
||||
from transformers.pytorch_utils import find_pruneable_heads_and_indices, prune_linear_layer
|
||||
|
||||
|
||||
class ViTEmbeddings(nn.Module):
|
||||
"""
|
||||
Construct the CLS token, position and patch embeddings. Optionally, also the mask token.
|
||||
"""
|
||||
|
||||
def __init__(self, config: ViTConfig, use_mask_token: bool = False) -> None:
|
||||
super().__init__()
|
||||
|
||||
self.cls_token = nn.Parameter(torch.randn(1, 1, config.hidden_size))
|
||||
self.mask_token = nn.Parameter(torch.zeros(1, 1, config.hidden_size)) if use_mask_token else None
|
||||
self.patch_embeddings = ViTPatchEmbeddings(config)
|
||||
num_patches = self.patch_embeddings.num_patches
|
||||
self.position_embeddings = nn.Parameter(torch.randn(1, num_patches + 1, config.hidden_size))
|
||||
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
||||
self.config = config
|
||||
|
||||
def interpolate_pos_encoding(self, embeddings: torch.Tensor, height: int, width: int) -> torch.Tensor:
|
||||
"""
|
||||
This method allows to interpolate the pre-trained position encodings, to be able to use the model on higher
|
||||
resolution images.
|
||||
|
||||
Source:
|
||||
https://github.com/facebookresearch/dino/blob/de9ee3df6cf39fac952ab558447af1fa1365362a/vision_transformer.py#L174
|
||||
"""
|
||||
|
||||
num_patches = embeddings.shape[1] - 1
|
||||
num_positions = self.position_embeddings.shape[1] - 1
|
||||
if num_patches == num_positions and height == width:
|
||||
return self.position_embeddings
|
||||
class_pos_embed = self.position_embeddings[:, 0]
|
||||
patch_pos_embed = self.position_embeddings[:, 1:]
|
||||
dim = embeddings.shape[-1]
|
||||
h0 = height // self.config.patch_size
|
||||
w0 = width // self.config.patch_size
|
||||
# we add a small number to avoid floating point error in the interpolation
|
||||
# see discussion at https://github.com/facebookresearch/dino/issues/8
|
||||
h0, w0 = h0 + 0.1, w0 + 0.1
|
||||
patch_pos_embed = patch_pos_embed.reshape(1, int(math.sqrt(num_positions)), int(math.sqrt(num_positions)), dim)
|
||||
patch_pos_embed = patch_pos_embed.permute(0, 3, 1, 2)
|
||||
patch_pos_embed = nn.functional.interpolate(
|
||||
patch_pos_embed,
|
||||
scale_factor=(h0 / math.sqrt(num_positions), w0 / math.sqrt(num_positions)),
|
||||
mode="bicubic",
|
||||
align_corners=False,
|
||||
)
|
||||
assert int(h0) == patch_pos_embed.shape[-2] and int(w0) == patch_pos_embed.shape[-1]
|
||||
patch_pos_embed = patch_pos_embed.permute(0, 2, 3, 1).view(1, -1, dim)
|
||||
return torch.cat((class_pos_embed.unsqueeze(0), patch_pos_embed), dim=1)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pixel_values: torch.Tensor,
|
||||
bool_masked_pos: Optional[torch.BoolTensor] = None,
|
||||
interpolate_pos_encoding: bool = False,
|
||||
) -> torch.Tensor:
|
||||
batch_size, num_channels, height, width = pixel_values.shape
|
||||
embeddings = self.patch_embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)
|
||||
|
||||
if bool_masked_pos is not None:
|
||||
seq_length = embeddings.shape[1]
|
||||
mask_tokens = self.mask_token.expand(batch_size, seq_length, -1)
|
||||
# replace the masked visual tokens by mask_tokens
|
||||
mask = bool_masked_pos.unsqueeze(-1).type_as(mask_tokens)
|
||||
embeddings = embeddings * (1.0 - mask) + mask_tokens * mask
|
||||
|
||||
# add the [CLS] token to the embedded patch tokens
|
||||
cls_tokens = self.cls_token.expand(batch_size, -1, -1)
|
||||
embeddings = torch.cat((cls_tokens, embeddings), dim=1)
|
||||
|
||||
# add positional encoding to each token
|
||||
if interpolate_pos_encoding:
|
||||
embeddings = embeddings + self.interpolate_pos_encoding(embeddings, height, width)
|
||||
else:
|
||||
embeddings = embeddings + self.position_embeddings
|
||||
|
||||
embeddings = self.dropout(embeddings)
|
||||
|
||||
return embeddings
|
||||
|
||||
|
||||
class ViTPatchEmbeddings(nn.Module):
|
||||
"""
|
||||
This class turns `pixel_values` of shape `(batch_size, num_channels, height, width)` into the initial
|
||||
`hidden_states` (patch embeddings) of shape `(batch_size, seq_length, hidden_size)` to be consumed by a
|
||||
Transformer.
|
||||
"""
|
||||
|
||||
def __init__(self, config):
|
||||
super().__init__()
|
||||
image_size, patch_size = config.image_size, config.patch_size
|
||||
num_channels, hidden_size = config.num_channels, config.hidden_size
|
||||
|
||||
image_size = image_size if isinstance(image_size, collections.abc.Iterable) else (image_size, image_size)
|
||||
patch_size = patch_size if isinstance(patch_size, collections.abc.Iterable) else (patch_size, patch_size)
|
||||
num_patches = (image_size[1] // patch_size[1]) * (image_size[0] // patch_size[0])
|
||||
self.image_size = image_size
|
||||
self.patch_size = patch_size
|
||||
self.num_channels = num_channels
|
||||
self.num_patches = num_patches
|
||||
|
||||
self.projection = nn.Conv2d(num_channels, hidden_size, kernel_size=patch_size, stride=patch_size)
|
||||
|
||||
def forward(self, pixel_values: torch.Tensor, interpolate_pos_encoding: bool = False) -> torch.Tensor:
|
||||
batch_size, num_channels, height, width = pixel_values.shape
|
||||
if num_channels != self.num_channels:
|
||||
raise ValueError(
|
||||
"Make sure that the channel dimension of the pixel values match with the one set in the configuration."
|
||||
f" Expected {self.num_channels} but got {num_channels}."
|
||||
)
|
||||
if not interpolate_pos_encoding:
|
||||
if height != self.image_size[0] or width != self.image_size[1]:
|
||||
raise ValueError(
|
||||
f"Input image size ({height}*{width}) doesn't match model"
|
||||
f" ({self.image_size[0]}*{self.image_size[1]})."
|
||||
)
|
||||
embeddings = self.projection(pixel_values).flatten(2).transpose(1, 2)
|
||||
return embeddings
|
||||
|
||||
|
||||
class ViTSelfAttention(nn.Module):
|
||||
def __init__(self, config: ViTConfig) -> None:
|
||||
super().__init__()
|
||||
if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
|
||||
raise ValueError(
|
||||
f"The hidden size {config.hidden_size,} is not a multiple of the number of attention "
|
||||
f"heads {config.num_attention_heads}."
|
||||
)
|
||||
|
||||
self.num_attention_heads = config.num_attention_heads
|
||||
self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
|
||||
self.all_head_size = self.num_attention_heads * self.attention_head_size
|
||||
|
||||
self.query = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
|
||||
self.key = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
|
||||
self.value = nn.Linear(config.hidden_size, self.all_head_size, bias=config.qkv_bias)
|
||||
|
||||
self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
|
||||
|
||||
def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
|
||||
new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
|
||||
x = x.view(new_x_shape)
|
||||
return x.permute(0, 2, 1, 3)
|
||||
|
||||
def forward(
|
||||
self, hidden_states, head_mask: Optional[torch.Tensor] = None, output_attentions: bool = False
|
||||
) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
|
||||
mixed_query_layer = self.query(hidden_states)
|
||||
|
||||
key_layer = self.transpose_for_scores(self.key(hidden_states))
|
||||
value_layer = self.transpose_for_scores(self.value(hidden_states))
|
||||
query_layer = self.transpose_for_scores(mixed_query_layer)
|
||||
|
||||
# Take the dot product between "query" and "key" to get the raw attention scores.
|
||||
attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
|
||||
|
||||
attention_scores = attention_scores / math.sqrt(self.attention_head_size)
|
||||
|
||||
# Normalize the attention scores to probabilities.
|
||||
attention_probs = nn.functional.softmax(attention_scores, dim=-1)
|
||||
|
||||
# This is actually dropping out entire tokens to attend to, which might
|
||||
# seem a bit unusual, but is taken from the original Transformer paper.
|
||||
attention_probs = self.dropout(attention_probs)
|
||||
|
||||
# Mask heads if we want to
|
||||
if head_mask is not None:
|
||||
attention_probs = attention_probs * head_mask
|
||||
|
||||
context_layer = torch.matmul(attention_probs, value_layer)
|
||||
|
||||
context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
|
||||
new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
|
||||
context_layer = context_layer.view(new_context_layer_shape)
|
||||
|
||||
outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
class ViTSelfOutput(nn.Module):
|
||||
"""
|
||||
The residual connection is defined in ViTLayer instead of here (as is the case with other models), due to the
|
||||
layernorm applied before each block.
|
||||
"""
|
||||
|
||||
def __init__(self, config: ViTConfig) -> None:
|
||||
super().__init__()
|
||||
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
|
||||
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
|
||||
hidden_states = self.dense(hidden_states)
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class ViTAttention(nn.Module):
|
||||
def __init__(self, config: ViTConfig) -> None:
|
||||
super().__init__()
|
||||
self.attention = ViTSelfAttention(config)
|
||||
self.output = ViTSelfOutput(config)
|
||||
self.pruned_heads = set()
|
||||
|
||||
def prune_heads(self, heads: Set[int]) -> None:
|
||||
if len(heads) == 0:
|
||||
return
|
||||
heads, index = find_pruneable_heads_and_indices(
|
||||
heads, self.attention.num_attention_heads, self.attention.attention_head_size, self.pruned_heads
|
||||
)
|
||||
|
||||
# Prune linear layers
|
||||
self.attention.query = prune_linear_layer(self.attention.query, index)
|
||||
self.attention.key = prune_linear_layer(self.attention.key, index)
|
||||
self.attention.value = prune_linear_layer(self.attention.value, index)
|
||||
self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
|
||||
|
||||
# Update hyper params and store pruned heads
|
||||
self.attention.num_attention_heads = self.attention.num_attention_heads - len(heads)
|
||||
self.attention.all_head_size = self.attention.attention_head_size * self.attention.num_attention_heads
|
||||
self.pruned_heads = self.pruned_heads.union(heads)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
output_attentions: bool = False,
|
||||
) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
|
||||
self_outputs = self.attention(hidden_states, head_mask, output_attentions)
|
||||
|
||||
attention_output = self.output(self_outputs[0], hidden_states)
|
||||
|
||||
outputs = (attention_output,) + self_outputs[1:] # add attentions if we output them
|
||||
return outputs
|
||||
|
||||
|
||||
class ViTIntermediate(nn.Module):
|
||||
def __init__(self, config: ViTConfig) -> None:
|
||||
super().__init__()
|
||||
self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
|
||||
if isinstance(config.hidden_act, str):
|
||||
self.intermediate_act_fn = ACT2FN[config.hidden_act]
|
||||
else:
|
||||
self.intermediate_act_fn = config.hidden_act
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
|
||||
hidden_states = self.dense(hidden_states)
|
||||
hidden_states = self.intermediate_act_fn(hidden_states)
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
class ViTOutput(nn.Module):
|
||||
def __init__(self, config: ViTConfig) -> None:
|
||||
super().__init__()
|
||||
self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
|
||||
self.dropout = nn.Dropout(config.hidden_dropout_prob)
|
||||
|
||||
def forward(self, hidden_states: torch.Tensor, input_tensor: torch.Tensor) -> torch.Tensor:
|
||||
hidden_states = self.dense(hidden_states)
|
||||
hidden_states = self.dropout(hidden_states)
|
||||
|
||||
hidden_states = hidden_states + input_tensor
|
||||
|
||||
return hidden_states
|
||||
|
||||
|
||||
def modulate(x, shift, scale):
|
||||
return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
|
||||
|
||||
|
||||
class ViTLayer(nn.Module):
|
||||
"""This corresponds to the Block class in the timm implementation."""
|
||||
|
||||
def __init__(self, config: ViTConfig) -> None:
|
||||
super().__init__()
|
||||
self.chunk_size_feed_forward = config.chunk_size_feed_forward
|
||||
self.seq_len_dim = 1
|
||||
self.attention = ViTAttention(config)
|
||||
self.intermediate = ViTIntermediate(config)
|
||||
self.output = ViTOutput(config)
|
||||
self.layernorm_before = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
||||
self.layernorm_after = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
||||
|
||||
self.adaLN_modulation = nn.Sequential(
|
||||
nn.SiLU(),
|
||||
nn.Linear(config.hidden_size, 4 * config.hidden_size, bias=True)
|
||||
)
|
||||
nn.init.constant_(self.adaLN_modulation[-1].weight, 0)
|
||||
nn.init.constant_(self.adaLN_modulation[-1].bias, 0)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
adaln_input: torch.Tensor = None,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
output_attentions: bool = False,
|
||||
) -> Union[Tuple[torch.Tensor, torch.Tensor], Tuple[torch.Tensor]]:
|
||||
shift_msa, scale_msa, shift_mlp, scale_mlp = self.adaLN_modulation(adaln_input).chunk(4, dim=1)
|
||||
|
||||
self_attention_outputs = self.attention(
|
||||
modulate(self.layernorm_before(hidden_states), shift_msa, scale_msa), # in ViT, layernorm is applied before self-attention
|
||||
head_mask,
|
||||
output_attentions=output_attentions,
|
||||
)
|
||||
attention_output = self_attention_outputs[0]
|
||||
outputs = self_attention_outputs[1:] # add self attentions if we output attention weights
|
||||
|
||||
# first residual connection
|
||||
hidden_states = attention_output + hidden_states
|
||||
|
||||
# in ViT, layernorm is also applied after self-attention
|
||||
layer_output = modulate(self.layernorm_after(hidden_states), shift_mlp, scale_mlp)
|
||||
layer_output = self.intermediate(layer_output)
|
||||
|
||||
# second residual connection is done here
|
||||
layer_output = self.output(layer_output, hidden_states)
|
||||
|
||||
outputs = (layer_output,) + outputs
|
||||
|
||||
return outputs
|
||||
|
||||
|
||||
class ViTEncoder(nn.Module):
|
||||
def __init__(self, config: ViTConfig) -> None:
|
||||
super().__init__()
|
||||
self.config = config
|
||||
self.layer = nn.ModuleList([ViTLayer(config) for _ in range(config.num_hidden_layers)])
|
||||
self.gradient_checkpointing = False
|
||||
|
||||
def forward(
|
||||
self,
|
||||
hidden_states: torch.Tensor,
|
||||
adaln_input: torch.Tensor = None,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
output_attentions: bool = False,
|
||||
output_hidden_states: bool = False,
|
||||
return_dict: bool = True,
|
||||
) -> Union[tuple, BaseModelOutput]:
|
||||
all_hidden_states = () if output_hidden_states else None
|
||||
all_self_attentions = () if output_attentions else None
|
||||
|
||||
for i, layer_module in enumerate(self.layer):
|
||||
if output_hidden_states:
|
||||
all_hidden_states = all_hidden_states + (hidden_states,)
|
||||
|
||||
layer_head_mask = head_mask[i] if head_mask is not None else None
|
||||
|
||||
if self.gradient_checkpointing and self.training:
|
||||
layer_outputs = self._gradient_checkpointing_func(
|
||||
layer_module.__call__,
|
||||
hidden_states,
|
||||
adaln_input,
|
||||
layer_head_mask,
|
||||
output_attentions,
|
||||
)
|
||||
else:
|
||||
layer_outputs = layer_module(hidden_states, adaln_input, layer_head_mask, output_attentions)
|
||||
|
||||
hidden_states = layer_outputs[0]
|
||||
|
||||
if output_attentions:
|
||||
all_self_attentions = all_self_attentions + (layer_outputs[1],)
|
||||
|
||||
if output_hidden_states:
|
||||
all_hidden_states = all_hidden_states + (hidden_states,)
|
||||
|
||||
if not return_dict:
|
||||
return tuple(v for v in [hidden_states, all_hidden_states, all_self_attentions] if v is not None)
|
||||
return BaseModelOutput(
|
||||
last_hidden_state=hidden_states,
|
||||
hidden_states=all_hidden_states,
|
||||
attentions=all_self_attentions,
|
||||
)
|
||||
|
||||
|
||||
class ViTPreTrainedModel(PreTrainedModel):
|
||||
"""
|
||||
An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
|
||||
models.
|
||||
"""
|
||||
|
||||
config_class = ViTConfig
|
||||
base_model_prefix = "vit"
|
||||
main_input_name = "pixel_values"
|
||||
supports_gradient_checkpointing = True
|
||||
_no_split_modules = ["ViTEmbeddings", "ViTLayer"]
|
||||
|
||||
def _init_weights(self, module: Union[nn.Linear, nn.Conv2d, nn.LayerNorm]) -> None:
|
||||
"""Initialize the weights"""
|
||||
if isinstance(module, (nn.Linear, nn.Conv2d)):
|
||||
# Upcast the input in `fp32` and cast it back to desired `dtype` to avoid
|
||||
# `trunc_normal_cpu` not implemented in `half` issues
|
||||
module.weight.data = nn.init.trunc_normal_(
|
||||
module.weight.data.to(torch.float32), mean=0.0, std=self.config.initializer_range
|
||||
).to(module.weight.dtype)
|
||||
if module.bias is not None:
|
||||
module.bias.data.zero_()
|
||||
elif isinstance(module, nn.LayerNorm):
|
||||
module.bias.data.zero_()
|
||||
module.weight.data.fill_(1.0)
|
||||
elif isinstance(module, ViTEmbeddings):
|
||||
module.position_embeddings.data = nn.init.trunc_normal_(
|
||||
module.position_embeddings.data.to(torch.float32),
|
||||
mean=0.0,
|
||||
std=self.config.initializer_range,
|
||||
).to(module.position_embeddings.dtype)
|
||||
|
||||
module.cls_token.data = nn.init.trunc_normal_(
|
||||
module.cls_token.data.to(torch.float32),
|
||||
mean=0.0,
|
||||
std=self.config.initializer_range,
|
||||
).to(module.cls_token.dtype)
|
||||
|
||||
|
||||
class ViTModel(ViTPreTrainedModel):
|
||||
def __init__(self, config: ViTConfig, add_pooling_layer: bool = True, use_mask_token: bool = False):
|
||||
super().__init__(config)
|
||||
self.config = config
|
||||
|
||||
self.embeddings = ViTEmbeddings(config, use_mask_token=use_mask_token)
|
||||
self.encoder = ViTEncoder(config)
|
||||
|
||||
self.layernorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
|
||||
self.pooler = ViTPooler(config) if add_pooling_layer else None
|
||||
|
||||
# Initialize weights and apply final processing
|
||||
self.post_init()
|
||||
|
||||
def get_input_embeddings(self) -> ViTPatchEmbeddings:
|
||||
return self.embeddings.patch_embeddings
|
||||
|
||||
def _prune_heads(self, heads_to_prune: Dict[int, List[int]]) -> None:
|
||||
"""
|
||||
Prunes heads of the model. heads_to_prune: dict of {layer_num: list of heads to prune in this layer} See base
|
||||
class PreTrainedModel
|
||||
"""
|
||||
for layer, heads in heads_to_prune.items():
|
||||
self.encoder.layer[layer].attention.prune_heads(heads)
|
||||
|
||||
def forward(
|
||||
self,
|
||||
pixel_values: Optional[torch.Tensor] = None,
|
||||
adaln_input: Optional[torch.Tensor] = None,
|
||||
bool_masked_pos: Optional[torch.BoolTensor] = None,
|
||||
head_mask: Optional[torch.Tensor] = None,
|
||||
output_attentions: Optional[bool] = None,
|
||||
output_hidden_states: Optional[bool] = None,
|
||||
interpolate_pos_encoding: Optional[bool] = None,
|
||||
return_dict: Optional[bool] = None,
|
||||
) -> Union[Tuple, BaseModelOutputWithPooling]:
|
||||
r"""
|
||||
bool_masked_pos (`torch.BoolTensor` of shape `(batch_size, num_patches)`, *optional*):
|
||||
Boolean masked positions. Indicates which patches are masked (1) and which aren't (0).
|
||||
"""
|
||||
output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
|
||||
output_hidden_states = (
|
||||
output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
|
||||
)
|
||||
return_dict = return_dict if return_dict is not None else self.config.use_return_dict
|
||||
|
||||
if pixel_values is None:
|
||||
raise ValueError("You have to specify pixel_values")
|
||||
|
||||
# Prepare head mask if needed
|
||||
# 1.0 in head_mask indicate we keep the head
|
||||
# attention_probs has shape bsz x n_heads x N x N
|
||||
# input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
|
||||
# and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
|
||||
head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
|
||||
|
||||
# TODO: maybe have a cleaner way to cast the input (from `ImageProcessor` side?)
|
||||
expected_dtype = self.embeddings.patch_embeddings.projection.weight.dtype
|
||||
if pixel_values.dtype != expected_dtype:
|
||||
pixel_values = pixel_values.to(expected_dtype)
|
||||
|
||||
embedding_output = self.embeddings(
|
||||
pixel_values, bool_masked_pos=bool_masked_pos, interpolate_pos_encoding=interpolate_pos_encoding
|
||||
)
|
||||
|
||||
encoder_outputs = self.encoder(
|
||||
embedding_output,
|
||||
adaln_input=adaln_input,
|
||||
head_mask=head_mask,
|
||||
output_attentions=output_attentions,
|
||||
output_hidden_states=output_hidden_states,
|
||||
return_dict=return_dict,
|
||||
)
|
||||
sequence_output = encoder_outputs[0]
|
||||
sequence_output = self.layernorm(sequence_output)
|
||||
pooled_output = self.pooler(sequence_output) if self.pooler is not None else None
|
||||
|
||||
if not return_dict:
|
||||
head_outputs = (sequence_output, pooled_output) if pooled_output is not None else (sequence_output,)
|
||||
return head_outputs + encoder_outputs[1:]
|
||||
|
||||
return BaseModelOutputWithPooling(
|
||||
last_hidden_state=sequence_output,
|
||||
pooler_output=pooled_output,
|
||||
hidden_states=encoder_outputs.hidden_states,
|
||||
attentions=encoder_outputs.attentions,
|
||||
)
|
||||
|
||||
|
||||
class ViTPooler(nn.Module):
|
||||
def __init__(self, config: ViTConfig):
|
||||
super().__init__()
|
||||
self.dense = nn.Linear(config.hidden_size, config.hidden_size)
|
||||
self.activation = nn.Tanh()
|
||||
|
||||
def forward(self, hidden_states):
|
||||
# We "pool" the model by simply taking the hidden state corresponding
|
||||
# to the first token.
|
||||
first_token_tensor = hidden_states[:, 0]
|
||||
pooled_output = self.dense(first_token_tensor)
|
||||
pooled_output = self.activation(pooled_output)
|
||||
return pooled_output
|
||||
@@ -0,0 +1,80 @@
|
||||
# Copyright (c) 2023, Zexin He
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# https://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
|
||||
import torch.nn as nn
|
||||
from transformers import ViTImageProcessor
|
||||
from einops import rearrange, repeat
|
||||
from .dino import ViTModel
|
||||
|
||||
|
||||
class DinoWrapper(nn.Module):
|
||||
"""
|
||||
Dino v1 wrapper using huggingface transformer implementation.
|
||||
"""
|
||||
def __init__(self, model_name: str, freeze: bool = True):
|
||||
super().__init__()
|
||||
self.model, self.processor = self._build_dino(model_name)
|
||||
self.camera_embedder = nn.Sequential(
|
||||
nn.Linear(16, self.model.config.hidden_size, bias=True),
|
||||
nn.SiLU(),
|
||||
nn.Linear(self.model.config.hidden_size, self.model.config.hidden_size, bias=True)
|
||||
)
|
||||
if freeze:
|
||||
self._freeze()
|
||||
|
||||
def forward(self, image, camera):
|
||||
# image: [B, N, C, H, W]
|
||||
# camera: [B, N, D]
|
||||
# RGB image with [0,1] scale and properly sized
|
||||
if image.ndim == 5:
|
||||
image = rearrange(image, 'b n c h w -> (b n) c h w')
|
||||
dtype = image.dtype
|
||||
inputs = self.processor(
|
||||
images=image.float(),
|
||||
return_tensors="pt",
|
||||
do_rescale=False,
|
||||
do_resize=False,
|
||||
).to(self.model.device).to(dtype)
|
||||
# embed camera
|
||||
N = camera.shape[1]
|
||||
camera_embeddings = self.camera_embedder(camera)
|
||||
camera_embeddings = rearrange(camera_embeddings, 'b n d -> (b n) d')
|
||||
embeddings = camera_embeddings
|
||||
# This resampling of positional embedding uses bicubic interpolation
|
||||
outputs = self.model(**inputs, adaln_input=embeddings, interpolate_pos_encoding=True)
|
||||
last_hidden_states = outputs.last_hidden_state
|
||||
return last_hidden_states
|
||||
|
||||
def _freeze(self):
|
||||
print(f"======== Freezing DinoWrapper ========")
|
||||
self.model.eval()
|
||||
for name, param in self.model.named_parameters():
|
||||
param.requires_grad = False
|
||||
|
||||
@staticmethod
|
||||
def _build_dino(model_name: str, proxy_error_retries: int = 3, proxy_error_cooldown: int = 5):
|
||||
import requests
|
||||
try:
|
||||
model = ViTModel.from_pretrained(model_name, add_pooling_layer=False)
|
||||
processor = ViTImageProcessor.from_pretrained(model_name)
|
||||
return model, processor
|
||||
except requests.exceptions.ProxyError as err:
|
||||
if proxy_error_retries > 0:
|
||||
print(f"Huggingface ProxyError: Retrying in {proxy_error_cooldown} seconds...")
|
||||
import time
|
||||
time.sleep(proxy_error_cooldown)
|
||||
return DinoWrapper._build_dino(model_name, proxy_error_retries - 1, proxy_error_cooldown)
|
||||
else:
|
||||
raise err
|
||||
@@ -0,0 +1,7 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
#
|
||||
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
|
||||
# and proprietary rights in and to this software, related documentation
|
||||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
|
||||
@@ -0,0 +1,16 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
#
|
||||
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
|
||||
# and proprietary rights in and to this software, related documentation
|
||||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
|
||||
|
||||
import torch
|
||||
from torch import nn
|
||||
|
||||
|
||||
class Camera(nn.Module):
|
||||
def __init__(self):
|
||||
super(Camera, self).__init__()
|
||||
pass
|
||||
@@ -0,0 +1,35 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
#
|
||||
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
|
||||
# and proprietary rights in and to this software, related documentation
|
||||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
|
||||
|
||||
import torch
|
||||
from . import Camera
|
||||
import numpy as np
|
||||
|
||||
|
||||
def projection(x=0.1, n=1.0, f=50.0, near_plane=None):
|
||||
if near_plane is None:
|
||||
near_plane = n
|
||||
return np.array(
|
||||
[[n / x, 0, 0, 0],
|
||||
[0, n / -x, 0, 0],
|
||||
[0, 0, -(f + near_plane) / (f - near_plane), -(2 * f * near_plane) / (f - near_plane)],
|
||||
[0, 0, -1, 0]]).astype(np.float32)
|
||||
|
||||
|
||||
class PerspectiveCamera(Camera):
|
||||
def __init__(self, fovy=49.0, device='cuda'):
|
||||
super(PerspectiveCamera, self).__init__()
|
||||
self.device = device
|
||||
focal = np.tan(fovy / 180.0 * np.pi * 0.5)
|
||||
self.proj_mtx = torch.from_numpy(projection(x=focal, f=1000.0, n=1.0, near_plane=0.1)).to(self.device).unsqueeze(dim=0)
|
||||
|
||||
def project(self, points_bxnx4):
|
||||
out = torch.matmul(
|
||||
points_bxnx4,
|
||||
torch.transpose(self.proj_mtx, 1, 2))
|
||||
return out
|
||||
@@ -0,0 +1,8 @@
|
||||
import torch
|
||||
|
||||
class Renderer():
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def forward(self):
|
||||
pass
|
||||
@@ -0,0 +1,121 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
#
|
||||
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
|
||||
# and proprietary rights in and to this software, related documentation
|
||||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
|
||||
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import nvdiffrast.torch as dr
|
||||
from . import Renderer
|
||||
|
||||
_FG_LUT = None
|
||||
|
||||
|
||||
def interpolate(attr, rast, attr_idx, rast_db=None):
|
||||
return dr.interpolate(
|
||||
attr.contiguous(), rast, attr_idx, rast_db=rast_db,
|
||||
diff_attrs=None if rast_db is None else 'all')
|
||||
|
||||
|
||||
def xfm_points(points, matrix, use_python=True):
|
||||
'''Transform points.
|
||||
Args:
|
||||
points: Tensor containing 3D points with shape [minibatch_size, num_vertices, 3] or [1, num_vertices, 3]
|
||||
matrix: A 4x4 transform matrix with shape [minibatch_size, 4, 4]
|
||||
use_python: Use PyTorch's torch.matmul (for validation)
|
||||
Returns:
|
||||
Transformed points in homogeneous 4D with shape [minibatch_size, num_vertices, 4].
|
||||
'''
|
||||
out = torch.matmul(torch.nn.functional.pad(points, pad=(0, 1), mode='constant', value=1.0), torch.transpose(matrix, 1, 2))
|
||||
if torch.is_anomaly_enabled():
|
||||
assert torch.all(torch.isfinite(out)), "Output of xfm_points contains inf or NaN"
|
||||
return out
|
||||
|
||||
|
||||
def dot(x, y):
|
||||
return torch.sum(x * y, -1, keepdim=True)
|
||||
|
||||
|
||||
def compute_vertex_normal(v_pos, t_pos_idx):
|
||||
i0 = t_pos_idx[:, 0]
|
||||
i1 = t_pos_idx[:, 1]
|
||||
i2 = t_pos_idx[:, 2]
|
||||
|
||||
v0 = v_pos[i0, :]
|
||||
v1 = v_pos[i1, :]
|
||||
v2 = v_pos[i2, :]
|
||||
|
||||
face_normals = torch.cross(v1 - v0, v2 - v0)
|
||||
|
||||
# Splat face normals to vertices
|
||||
v_nrm = torch.zeros_like(v_pos)
|
||||
v_nrm.scatter_add_(0, i0[:, None].repeat(1, 3), face_normals)
|
||||
v_nrm.scatter_add_(0, i1[:, None].repeat(1, 3), face_normals)
|
||||
v_nrm.scatter_add_(0, i2[:, None].repeat(1, 3), face_normals)
|
||||
|
||||
# Normalize, replace zero (degenerated) normals with some default value
|
||||
v_nrm = torch.where(
|
||||
dot(v_nrm, v_nrm) > 1e-20, v_nrm, torch.as_tensor([0.0, 0.0, 1.0]).to(v_nrm)
|
||||
)
|
||||
v_nrm = F.normalize(v_nrm, dim=1)
|
||||
assert torch.all(torch.isfinite(v_nrm))
|
||||
|
||||
return v_nrm
|
||||
|
||||
|
||||
class NeuralRender(Renderer):
|
||||
def __init__(self, device='cuda', camera_model=None):
|
||||
super(NeuralRender, self).__init__()
|
||||
self.device = device
|
||||
self.ctx = dr.RasterizeCudaContext(device=device)
|
||||
self.projection_mtx = None
|
||||
self.camera = camera_model
|
||||
|
||||
def render_mesh(
|
||||
self,
|
||||
mesh_v_pos_bxnx3,
|
||||
mesh_t_pos_idx_fx3,
|
||||
camera_mv_bx4x4,
|
||||
mesh_v_feat_bxnxd,
|
||||
resolution=256,
|
||||
spp=1,
|
||||
device='cuda',
|
||||
hierarchical_mask=False
|
||||
):
|
||||
assert not hierarchical_mask
|
||||
|
||||
mtx_in = torch.tensor(camera_mv_bx4x4, dtype=torch.float32, device=device) if not torch.is_tensor(camera_mv_bx4x4) else camera_mv_bx4x4
|
||||
v_pos = xfm_points(mesh_v_pos_bxnx3, mtx_in) # Rotate it to camera coordinates
|
||||
v_pos_clip = self.camera.project(v_pos) # Projection in the camera
|
||||
|
||||
v_nrm = compute_vertex_normal(mesh_v_pos_bxnx3[0], mesh_t_pos_idx_fx3.long()) # vertex normals in world coordinates
|
||||
|
||||
# Render the image,
|
||||
# Here we only return the feature (3D location) at each pixel, which will be used as the input for neural render
|
||||
num_layers = 1
|
||||
mask_pyramid = None
|
||||
assert mesh_t_pos_idx_fx3.shape[0] > 0 # Make sure we have shapes
|
||||
mesh_v_feat_bxnxd = torch.cat([mesh_v_feat_bxnxd.repeat(v_pos.shape[0], 1, 1), v_pos], dim=-1) # Concatenate the pos
|
||||
|
||||
with dr.DepthPeeler(self.ctx, v_pos_clip, mesh_t_pos_idx_fx3, [resolution * spp, resolution * spp]) as peeler:
|
||||
for _ in range(num_layers):
|
||||
rast, db = peeler.rasterize_next_layer()
|
||||
gb_feat, _ = interpolate(mesh_v_feat_bxnxd, rast, mesh_t_pos_idx_fx3)
|
||||
|
||||
hard_mask = torch.clamp(rast[..., -1:], 0, 1)
|
||||
antialias_mask = dr.antialias(
|
||||
hard_mask.clone().contiguous(), rast, v_pos_clip,
|
||||
mesh_t_pos_idx_fx3)
|
||||
|
||||
depth = gb_feat[..., -2:-1]
|
||||
ori_mesh_feature = gb_feat[..., :-4]
|
||||
|
||||
normal, _ = interpolate(v_nrm[None, ...], rast, mesh_t_pos_idx_fx3)
|
||||
normal = dr.antialias(normal.clone().contiguous(), rast, v_pos_clip, mesh_t_pos_idx_fx3)
|
||||
normal = F.normalize(normal, dim=-1)
|
||||
normal = torch.lerp(torch.zeros_like(normal), (normal + 1.0) / 2.0, hard_mask.float()) # black background
|
||||
|
||||
return ori_mesh_feature, antialias_mask, hard_mask, rast, v_pos_clip, mask_pyramid, depth, normal
|
||||
@@ -0,0 +1,18 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
#
|
||||
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
|
||||
# and proprietary rights in and to this software, related documentation
|
||||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
|
||||
|
||||
class Geometry():
|
||||
def __init__(self):
|
||||
pass
|
||||
|
||||
def forward(self):
|
||||
pass
|
||||
@@ -0,0 +1,504 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
#
|
||||
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
|
||||
# and proprietary rights in and to this software, related documentation
|
||||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
import os
|
||||
from . import Geometry
|
||||
from .dmtet_utils import get_center_boundary_index
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
###############################################################################
|
||||
# DMTet utility functions
|
||||
###############################################################################
|
||||
def create_mt_variable(device):
|
||||
triangle_table = torch.tensor(
|
||||
[
|
||||
[-1, -1, -1, -1, -1, -1],
|
||||
[1, 0, 2, -1, -1, -1],
|
||||
[4, 0, 3, -1, -1, -1],
|
||||
[1, 4, 2, 1, 3, 4],
|
||||
[3, 1, 5, -1, -1, -1],
|
||||
[2, 3, 0, 2, 5, 3],
|
||||
[1, 4, 0, 1, 5, 4],
|
||||
[4, 2, 5, -1, -1, -1],
|
||||
[4, 5, 2, -1, -1, -1],
|
||||
[4, 1, 0, 4, 5, 1],
|
||||
[3, 2, 0, 3, 5, 2],
|
||||
[1, 3, 5, -1, -1, -1],
|
||||
[4, 1, 2, 4, 3, 1],
|
||||
[3, 0, 4, -1, -1, -1],
|
||||
[2, 0, 1, -1, -1, -1],
|
||||
[-1, -1, -1, -1, -1, -1]
|
||||
], dtype=torch.long, device=device)
|
||||
|
||||
num_triangles_table = torch.tensor([0, 1, 1, 2, 1, 2, 2, 1, 1, 2, 2, 1, 2, 1, 1, 0], dtype=torch.long, device=device)
|
||||
base_tet_edges = torch.tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long, device=device)
|
||||
v_id = torch.pow(2, torch.arange(4, dtype=torch.long, device=device))
|
||||
return triangle_table, num_triangles_table, base_tet_edges, v_id
|
||||
|
||||
|
||||
def sort_edges(edges_ex2):
|
||||
with torch.no_grad():
|
||||
order = (edges_ex2[:, 0] > edges_ex2[:, 1]).long()
|
||||
order = order.unsqueeze(dim=1)
|
||||
a = torch.gather(input=edges_ex2, index=order, dim=1)
|
||||
b = torch.gather(input=edges_ex2, index=1 - order, dim=1)
|
||||
return torch.stack([a, b], -1)
|
||||
|
||||
|
||||
###############################################################################
|
||||
# marching tetrahedrons (differentiable)
|
||||
###############################################################################
|
||||
|
||||
def marching_tets(pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id):
|
||||
with torch.no_grad():
|
||||
occ_n = sdf_n > 0
|
||||
occ_fx4 = occ_n[tet_fx4.reshape(-1)].reshape(-1, 4)
|
||||
occ_sum = torch.sum(occ_fx4, -1)
|
||||
valid_tets = (occ_sum > 0) & (occ_sum < 4)
|
||||
occ_sum = occ_sum[valid_tets]
|
||||
|
||||
# find all vertices
|
||||
all_edges = tet_fx4[valid_tets][:, base_tet_edges].reshape(-1, 2)
|
||||
all_edges = sort_edges(all_edges)
|
||||
unique_edges, idx_map = torch.unique(all_edges, dim=0, return_inverse=True)
|
||||
|
||||
unique_edges = unique_edges.long()
|
||||
mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 1
|
||||
mapping = torch.ones((unique_edges.shape[0]), dtype=torch.long, device=sdf_n.device) * -1
|
||||
mapping[mask_edges] = torch.arange(mask_edges.sum(), dtype=torch.long, device=sdf_n.device)
|
||||
idx_map = mapping[idx_map] # map edges to verts
|
||||
|
||||
interp_v = unique_edges[mask_edges] # .long()
|
||||
edges_to_interp = pos_nx3[interp_v.reshape(-1)].reshape(-1, 2, 3)
|
||||
edges_to_interp_sdf = sdf_n[interp_v.reshape(-1)].reshape(-1, 2, 1)
|
||||
edges_to_interp_sdf[:, -1] *= -1
|
||||
|
||||
denominator = edges_to_interp_sdf.sum(1, keepdim=True)
|
||||
|
||||
edges_to_interp_sdf = torch.flip(edges_to_interp_sdf, [1]) / denominator
|
||||
verts = (edges_to_interp * edges_to_interp_sdf).sum(1)
|
||||
|
||||
idx_map = idx_map.reshape(-1, 6)
|
||||
|
||||
tetindex = (occ_fx4[valid_tets] * v_id.unsqueeze(0)).sum(-1)
|
||||
num_triangles = num_triangles_table[tetindex]
|
||||
|
||||
# Generate triangle indices
|
||||
faces = torch.cat(
|
||||
(
|
||||
torch.gather(
|
||||
input=idx_map[num_triangles == 1], dim=1,
|
||||
index=triangle_table[tetindex[num_triangles == 1]][:, :3]).reshape(-1, 3),
|
||||
torch.gather(
|
||||
input=idx_map[num_triangles == 2], dim=1,
|
||||
index=triangle_table[tetindex[num_triangles == 2]][:, :6]).reshape(-1, 3),
|
||||
), dim=0)
|
||||
return verts, faces
|
||||
|
||||
|
||||
def create_tetmesh_variables(device='cuda'):
|
||||
tet_table = torch.tensor(
|
||||
[[-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1],
|
||||
[0, 4, 5, 6, -1, -1, -1, -1, -1, -1, -1, -1],
|
||||
[1, 4, 7, 8, -1, -1, -1, -1, -1, -1, -1, -1],
|
||||
[1, 0, 8, 7, 0, 5, 8, 7, 0, 5, 6, 8],
|
||||
[2, 5, 7, 9, -1, -1, -1, -1, -1, -1, -1, -1],
|
||||
[2, 0, 9, 7, 0, 4, 9, 7, 0, 4, 6, 9],
|
||||
[2, 1, 9, 5, 1, 4, 9, 5, 1, 4, 8, 9],
|
||||
[6, 0, 1, 2, 6, 1, 2, 8, 6, 8, 2, 9],
|
||||
[3, 6, 8, 9, -1, -1, -1, -1, -1, -1, -1, -1],
|
||||
[3, 0, 9, 8, 0, 4, 9, 8, 0, 4, 5, 9],
|
||||
[3, 1, 9, 6, 1, 4, 9, 6, 1, 4, 7, 9],
|
||||
[5, 0, 1, 3, 5, 1, 3, 7, 5, 7, 3, 9],
|
||||
[3, 2, 8, 6, 2, 5, 8, 6, 2, 5, 7, 8],
|
||||
[4, 0, 2, 3, 4, 2, 3, 7, 4, 7, 3, 8],
|
||||
[4, 1, 2, 3, 4, 2, 3, 5, 4, 5, 3, 6],
|
||||
[-1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1, -1]], dtype=torch.long, device=device)
|
||||
num_tets_table = torch.tensor([0, 1, 1, 3, 1, 3, 3, 3, 1, 3, 3, 3, 3, 3, 3, 0], dtype=torch.long, device=device)
|
||||
return tet_table, num_tets_table
|
||||
|
||||
|
||||
def marching_tets_tetmesh(
|
||||
pos_nx3, sdf_n, tet_fx4, triangle_table, num_triangles_table, base_tet_edges, v_id,
|
||||
return_tet_mesh=False, ori_v=None, num_tets_table=None, tet_table=None):
|
||||
with torch.no_grad():
|
||||
occ_n = sdf_n > 0
|
||||
occ_fx4 = occ_n[tet_fx4.reshape(-1)].reshape(-1, 4)
|
||||
occ_sum = torch.sum(occ_fx4, -1)
|
||||
valid_tets = (occ_sum > 0) & (occ_sum < 4)
|
||||
occ_sum = occ_sum[valid_tets]
|
||||
|
||||
# find all vertices
|
||||
all_edges = tet_fx4[valid_tets][:, base_tet_edges].reshape(-1, 2)
|
||||
all_edges = sort_edges(all_edges)
|
||||
unique_edges, idx_map = torch.unique(all_edges, dim=0, return_inverse=True)
|
||||
|
||||
unique_edges = unique_edges.long()
|
||||
mask_edges = occ_n[unique_edges.reshape(-1)].reshape(-1, 2).sum(-1) == 1
|
||||
mapping = torch.ones((unique_edges.shape[0]), dtype=torch.long, device=sdf_n.device) * -1
|
||||
mapping[mask_edges] = torch.arange(mask_edges.sum(), dtype=torch.long, device=sdf_n.device)
|
||||
idx_map = mapping[idx_map] # map edges to verts
|
||||
|
||||
interp_v = unique_edges[mask_edges] # .long()
|
||||
edges_to_interp = pos_nx3[interp_v.reshape(-1)].reshape(-1, 2, 3)
|
||||
edges_to_interp_sdf = sdf_n[interp_v.reshape(-1)].reshape(-1, 2, 1)
|
||||
edges_to_interp_sdf[:, -1] *= -1
|
||||
|
||||
denominator = edges_to_interp_sdf.sum(1, keepdim=True)
|
||||
|
||||
edges_to_interp_sdf = torch.flip(edges_to_interp_sdf, [1]) / denominator
|
||||
verts = (edges_to_interp * edges_to_interp_sdf).sum(1)
|
||||
|
||||
idx_map = idx_map.reshape(-1, 6)
|
||||
|
||||
tetindex = (occ_fx4[valid_tets] * v_id.unsqueeze(0)).sum(-1)
|
||||
num_triangles = num_triangles_table[tetindex]
|
||||
|
||||
# Generate triangle indices
|
||||
faces = torch.cat(
|
||||
(
|
||||
torch.gather(
|
||||
input=idx_map[num_triangles == 1], dim=1,
|
||||
index=triangle_table[tetindex[num_triangles == 1]][:, :3]).reshape(-1, 3),
|
||||
torch.gather(
|
||||
input=idx_map[num_triangles == 2], dim=1,
|
||||
index=triangle_table[tetindex[num_triangles == 2]][:, :6]).reshape(-1, 3),
|
||||
), dim=0)
|
||||
if not return_tet_mesh:
|
||||
return verts, faces
|
||||
occupied_verts = ori_v[occ_n]
|
||||
mapping = torch.ones((pos_nx3.shape[0]), dtype=torch.long, device="cuda") * -1
|
||||
mapping[occ_n] = torch.arange(occupied_verts.shape[0], device="cuda")
|
||||
tet_fx4 = mapping[tet_fx4.reshape(-1)].reshape((-1, 4))
|
||||
|
||||
idx_map = torch.cat([tet_fx4[valid_tets] + verts.shape[0], idx_map], -1) # t x 10
|
||||
tet_verts = torch.cat([verts, occupied_verts], 0)
|
||||
num_tets = num_tets_table[tetindex]
|
||||
|
||||
tets = torch.cat(
|
||||
(
|
||||
torch.gather(input=idx_map[num_tets == 1], dim=1, index=tet_table[tetindex[num_tets == 1]][:, :4]).reshape(
|
||||
-1,
|
||||
4),
|
||||
torch.gather(input=idx_map[num_tets == 3], dim=1, index=tet_table[tetindex[num_tets == 3]][:, :12]).reshape(
|
||||
-1,
|
||||
4),
|
||||
), dim=0)
|
||||
# add fully occupied tets
|
||||
fully_occupied = occ_fx4.sum(-1) == 4
|
||||
tet_fully_occupied = tet_fx4[fully_occupied] + verts.shape[0]
|
||||
tets = torch.cat([tets, tet_fully_occupied])
|
||||
|
||||
return verts, faces, tet_verts, tets
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Compact tet grid
|
||||
###############################################################################
|
||||
|
||||
def compact_tets(pos_nx3, sdf_n, tet_fx4):
|
||||
with torch.no_grad():
|
||||
# Find surface tets
|
||||
occ_n = sdf_n > 0
|
||||
occ_fx4 = occ_n[tet_fx4.reshape(-1)].reshape(-1, 4)
|
||||
occ_sum = torch.sum(occ_fx4, -1)
|
||||
valid_tets = (occ_sum > 0) & (occ_sum < 4) # one value per tet, these are the surface tets
|
||||
|
||||
valid_vtx = tet_fx4[valid_tets].reshape(-1)
|
||||
unique_vtx, idx_map = torch.unique(valid_vtx, dim=0, return_inverse=True)
|
||||
new_pos = pos_nx3[unique_vtx]
|
||||
new_sdf = sdf_n[unique_vtx]
|
||||
new_tets = idx_map.reshape(-1, 4)
|
||||
return new_pos, new_sdf, new_tets
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Subdivide volume
|
||||
###############################################################################
|
||||
|
||||
def batch_subdivide_volume(tet_pos_bxnx3, tet_bxfx4, grid_sdf):
|
||||
device = tet_pos_bxnx3.device
|
||||
# get new verts
|
||||
tet_fx4 = tet_bxfx4[0]
|
||||
edges = [0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3]
|
||||
all_edges = tet_fx4[:, edges].reshape(-1, 2)
|
||||
all_edges = sort_edges(all_edges)
|
||||
unique_edges, idx_map = torch.unique(all_edges, dim=0, return_inverse=True)
|
||||
idx_map = idx_map + tet_pos_bxnx3.shape[1]
|
||||
all_values = torch.cat([tet_pos_bxnx3, grid_sdf], -1)
|
||||
mid_points_pos = all_values[:, unique_edges.reshape(-1)].reshape(
|
||||
all_values.shape[0], -1, 2,
|
||||
all_values.shape[-1]).mean(2)
|
||||
new_v = torch.cat([all_values, mid_points_pos], 1)
|
||||
new_v, new_sdf = new_v[..., :3], new_v[..., 3]
|
||||
|
||||
# get new tets
|
||||
|
||||
idx_a, idx_b, idx_c, idx_d = tet_fx4[:, 0], tet_fx4[:, 1], tet_fx4[:, 2], tet_fx4[:, 3]
|
||||
idx_ab = idx_map[0::6]
|
||||
idx_ac = idx_map[1::6]
|
||||
idx_ad = idx_map[2::6]
|
||||
idx_bc = idx_map[3::6]
|
||||
idx_bd = idx_map[4::6]
|
||||
idx_cd = idx_map[5::6]
|
||||
|
||||
tet_1 = torch.stack([idx_a, idx_ab, idx_ac, idx_ad], dim=1)
|
||||
tet_2 = torch.stack([idx_b, idx_bc, idx_ab, idx_bd], dim=1)
|
||||
tet_3 = torch.stack([idx_c, idx_ac, idx_bc, idx_cd], dim=1)
|
||||
tet_4 = torch.stack([idx_d, idx_ad, idx_cd, idx_bd], dim=1)
|
||||
tet_5 = torch.stack([idx_ab, idx_ac, idx_ad, idx_bd], dim=1)
|
||||
tet_6 = torch.stack([idx_ab, idx_ac, idx_bd, idx_bc], dim=1)
|
||||
tet_7 = torch.stack([idx_cd, idx_ac, idx_bd, idx_ad], dim=1)
|
||||
tet_8 = torch.stack([idx_cd, idx_ac, idx_bc, idx_bd], dim=1)
|
||||
|
||||
tet_np = torch.cat([tet_1, tet_2, tet_3, tet_4, tet_5, tet_6, tet_7, tet_8], dim=0)
|
||||
tet_np = tet_np.reshape(1, -1, 4).expand(tet_pos_bxnx3.shape[0], -1, -1)
|
||||
tet = tet_np.long().to(device)
|
||||
|
||||
return new_v, tet, new_sdf
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Adjacency
|
||||
###############################################################################
|
||||
def tet_to_tet_adj_sparse(tet_tx4):
|
||||
# include self connection!!!!!!!!!!!!!!!!!!!
|
||||
with torch.no_grad():
|
||||
t = tet_tx4.shape[0]
|
||||
device = tet_tx4.device
|
||||
idx_array = torch.LongTensor(
|
||||
[0, 1, 2,
|
||||
1, 0, 3,
|
||||
2, 3, 0,
|
||||
3, 2, 1]).to(device).reshape(4, 3).unsqueeze(0).expand(t, -1, -1) # (t, 4, 3)
|
||||
|
||||
# get all faces
|
||||
all_faces = torch.gather(input=tet_tx4.unsqueeze(1).expand(-1, 4, -1), index=idx_array, dim=-1).reshape(
|
||||
-1,
|
||||
3) # (tx4, 3)
|
||||
all_faces_tet_idx = torch.arange(t, device=device).unsqueeze(-1).expand(-1, 4).reshape(-1)
|
||||
# sort and group
|
||||
all_faces_sorted, _ = torch.sort(all_faces, dim=1)
|
||||
|
||||
all_faces_unique, inverse_indices, counts = torch.unique(
|
||||
all_faces_sorted, dim=0, return_counts=True,
|
||||
return_inverse=True)
|
||||
tet_face_fx3 = all_faces_unique[counts == 2]
|
||||
counts = counts[inverse_indices] # tx4
|
||||
valid = (counts == 2)
|
||||
|
||||
group = inverse_indices[valid]
|
||||
# print (inverse_indices.shape, group.shape, all_faces_tet_idx.shape)
|
||||
_, indices = torch.sort(group)
|
||||
all_faces_tet_idx_grouped = all_faces_tet_idx[valid][indices]
|
||||
tet_face_tetidx_fx2 = torch.stack([all_faces_tet_idx_grouped[::2], all_faces_tet_idx_grouped[1::2]], dim=-1)
|
||||
|
||||
tet_adj_idx = torch.cat([tet_face_tetidx_fx2, torch.flip(tet_face_tetidx_fx2, [1])])
|
||||
adj_self = torch.arange(t, device=tet_tx4.device)
|
||||
adj_self = torch.stack([adj_self, adj_self], -1)
|
||||
tet_adj_idx = torch.cat([tet_adj_idx, adj_self])
|
||||
|
||||
tet_adj_idx = torch.unique(tet_adj_idx, dim=0)
|
||||
values = torch.ones(
|
||||
tet_adj_idx.shape[0], device=tet_tx4.device).float()
|
||||
adj_sparse = torch.sparse.FloatTensor(
|
||||
tet_adj_idx.t(), values, torch.Size([t, t]))
|
||||
|
||||
# normalization
|
||||
neighbor_num = 1.0 / torch.sparse.sum(
|
||||
adj_sparse, dim=1).to_dense()
|
||||
values = torch.index_select(neighbor_num, 0, tet_adj_idx[:, 0])
|
||||
adj_sparse = torch.sparse.FloatTensor(
|
||||
tet_adj_idx.t(), values, torch.Size([t, t]))
|
||||
return adj_sparse
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Compact grid
|
||||
###############################################################################
|
||||
|
||||
def get_tet_bxfx4x3(bxnxz, bxfx4):
|
||||
n_batch, z = bxnxz.shape[0], bxnxz.shape[2]
|
||||
gather_input = bxnxz.unsqueeze(2).expand(
|
||||
n_batch, bxnxz.shape[1], 4, z)
|
||||
gather_index = bxfx4.unsqueeze(-1).expand(
|
||||
n_batch, bxfx4.shape[1], 4, z).long()
|
||||
tet_bxfx4xz = torch.gather(
|
||||
input=gather_input, dim=1, index=gather_index)
|
||||
|
||||
return tet_bxfx4xz
|
||||
|
||||
|
||||
def shrink_grid(tet_pos_bxnx3, tet_bxfx4, grid_sdf):
|
||||
with torch.no_grad():
|
||||
assert tet_pos_bxnx3.shape[0] == 1
|
||||
|
||||
occ = grid_sdf[0] > 0
|
||||
occ_sum = get_tet_bxfx4x3(occ.unsqueeze(0).unsqueeze(-1), tet_bxfx4).reshape(-1, 4).sum(-1)
|
||||
mask = (occ_sum > 0) & (occ_sum < 4)
|
||||
|
||||
# build connectivity graph
|
||||
adj_matrix = tet_to_tet_adj_sparse(tet_bxfx4[0])
|
||||
mask = mask.float().unsqueeze(-1)
|
||||
|
||||
# Include a one ring of neighbors
|
||||
for i in range(1):
|
||||
mask = torch.sparse.mm(adj_matrix, mask)
|
||||
mask = mask.squeeze(-1) > 0
|
||||
|
||||
mapping = torch.zeros((tet_pos_bxnx3.shape[1]), device=tet_pos_bxnx3.device, dtype=torch.long)
|
||||
new_tet_bxfx4 = tet_bxfx4[:, mask].long()
|
||||
selected_verts_idx = torch.unique(new_tet_bxfx4)
|
||||
new_tet_pos_bxnx3 = tet_pos_bxnx3[:, selected_verts_idx]
|
||||
mapping[selected_verts_idx] = torch.arange(selected_verts_idx.shape[0], device=tet_pos_bxnx3.device)
|
||||
new_tet_bxfx4 = mapping[new_tet_bxfx4.reshape(-1)].reshape(new_tet_bxfx4.shape)
|
||||
new_grid_sdf = grid_sdf[:, selected_verts_idx]
|
||||
return new_tet_pos_bxnx3, new_tet_bxfx4, new_grid_sdf
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Regularizer
|
||||
###############################################################################
|
||||
|
||||
def sdf_reg_loss(sdf, all_edges):
|
||||
sdf_f1x6x2 = sdf[all_edges.reshape(-1)].reshape(-1, 2)
|
||||
mask = torch.sign(sdf_f1x6x2[..., 0]) != torch.sign(sdf_f1x6x2[..., 1])
|
||||
sdf_f1x6x2 = sdf_f1x6x2[mask]
|
||||
sdf_diff = torch.nn.functional.binary_cross_entropy_with_logits(
|
||||
sdf_f1x6x2[..., 0],
|
||||
(sdf_f1x6x2[..., 1] > 0).float()) + \
|
||||
torch.nn.functional.binary_cross_entropy_with_logits(
|
||||
sdf_f1x6x2[..., 1],
|
||||
(sdf_f1x6x2[..., 0] > 0).float())
|
||||
return sdf_diff
|
||||
|
||||
|
||||
def sdf_reg_loss_batch(sdf, all_edges):
|
||||
sdf_f1x6x2 = sdf[:, all_edges.reshape(-1)].reshape(sdf.shape[0], -1, 2)
|
||||
mask = torch.sign(sdf_f1x6x2[..., 0]) != torch.sign(sdf_f1x6x2[..., 1])
|
||||
sdf_f1x6x2 = sdf_f1x6x2[mask]
|
||||
sdf_diff = torch.nn.functional.binary_cross_entropy_with_logits(sdf_f1x6x2[..., 0], (sdf_f1x6x2[..., 1] > 0).float()) + \
|
||||
torch.nn.functional.binary_cross_entropy_with_logits(sdf_f1x6x2[..., 1], (sdf_f1x6x2[..., 0] > 0).float())
|
||||
return sdf_diff
|
||||
|
||||
|
||||
###############################################################################
|
||||
# Geometry interface
|
||||
###############################################################################
|
||||
class DMTetGeometry(Geometry):
|
||||
def __init__(
|
||||
self, grid_res=64, scale=2.0, device='cuda', renderer=None,
|
||||
render_type='neural_render', args=None):
|
||||
super(DMTetGeometry, self).__init__()
|
||||
self.grid_res = grid_res
|
||||
self.device = device
|
||||
self.args = args
|
||||
tets = np.load('data/tets/%d_compress.npz' % (grid_res))
|
||||
self.verts = torch.from_numpy(tets['vertices']).float().to(self.device)
|
||||
# Make sure the tet is zero-centered and length is equal to 1
|
||||
length = self.verts.max(dim=0)[0] - self.verts.min(dim=0)[0]
|
||||
length = length.max()
|
||||
mid = (self.verts.max(dim=0)[0] + self.verts.min(dim=0)[0]) / 2.0
|
||||
self.verts = (self.verts - mid.unsqueeze(dim=0)) / length
|
||||
if isinstance(scale, list):
|
||||
self.verts[:, 0] = self.verts[:, 0] * scale[0]
|
||||
self.verts[:, 1] = self.verts[:, 1] * scale[1]
|
||||
self.verts[:, 2] = self.verts[:, 2] * scale[1]
|
||||
else:
|
||||
self.verts = self.verts * scale
|
||||
self.indices = torch.from_numpy(tets['tets']).long().to(self.device)
|
||||
self.triangle_table, self.num_triangles_table, self.base_tet_edges, self.v_id = create_mt_variable(self.device)
|
||||
self.tet_table, self.num_tets_table = create_tetmesh_variables(self.device)
|
||||
# Parameters for regularization computation
|
||||
edges = torch.tensor([0, 1, 0, 2, 0, 3, 1, 2, 1, 3, 2, 3], dtype=torch.long, device=self.device)
|
||||
all_edges = self.indices[:, edges].reshape(-1, 2)
|
||||
all_edges_sorted = torch.sort(all_edges, dim=1)[0]
|
||||
self.all_edges = torch.unique(all_edges_sorted, dim=0)
|
||||
|
||||
# Parameters used for fix boundary sdf
|
||||
self.center_indices, self.boundary_indices = get_center_boundary_index(self.verts)
|
||||
self.renderer = renderer
|
||||
self.render_type = render_type
|
||||
|
||||
def getAABB(self):
|
||||
return torch.min(self.verts, dim=0).values, torch.max(self.verts, dim=0).values
|
||||
|
||||
def get_mesh(self, v_deformed_nx3, sdf_n, with_uv=False, indices=None):
|
||||
if indices is None:
|
||||
indices = self.indices
|
||||
verts, faces = marching_tets(
|
||||
v_deformed_nx3, sdf_n, indices, self.triangle_table,
|
||||
self.num_triangles_table, self.base_tet_edges, self.v_id)
|
||||
faces = torch.cat(
|
||||
[faces[:, 0:1],
|
||||
faces[:, 2:3],
|
||||
faces[:, 1:2], ], dim=-1)
|
||||
return verts, faces
|
||||
|
||||
def get_tet_mesh(self, v_deformed_nx3, sdf_n, with_uv=False, indices=None):
|
||||
if indices is None:
|
||||
indices = self.indices
|
||||
verts, faces, tet_verts, tets = marching_tets_tetmesh(
|
||||
v_deformed_nx3, sdf_n, indices, self.triangle_table,
|
||||
self.num_triangles_table, self.base_tet_edges, self.v_id, return_tet_mesh=True,
|
||||
num_tets_table=self.num_tets_table, tet_table=self.tet_table, ori_v=v_deformed_nx3)
|
||||
faces = torch.cat(
|
||||
[faces[:, 0:1],
|
||||
faces[:, 2:3],
|
||||
faces[:, 1:2], ], dim=-1)
|
||||
return verts, faces, tet_verts, tets
|
||||
|
||||
def render_mesh(self, mesh_v_nx3, mesh_f_fx3, camera_mv_bx4x4, resolution=256, hierarchical_mask=False):
|
||||
return_value = dict()
|
||||
if self.render_type == 'neural_render':
|
||||
tex_pos, mask, hard_mask, rast, v_pos_clip, mask_pyramid, depth = self.renderer.render_mesh(
|
||||
mesh_v_nx3.unsqueeze(dim=0),
|
||||
mesh_f_fx3.int(),
|
||||
camera_mv_bx4x4,
|
||||
mesh_v_nx3.unsqueeze(dim=0),
|
||||
resolution=resolution,
|
||||
device=self.device,
|
||||
hierarchical_mask=hierarchical_mask
|
||||
)
|
||||
|
||||
return_value['tex_pos'] = tex_pos
|
||||
return_value['mask'] = mask
|
||||
return_value['hard_mask'] = hard_mask
|
||||
return_value['rast'] = rast
|
||||
return_value['v_pos_clip'] = v_pos_clip
|
||||
return_value['mask_pyramid'] = mask_pyramid
|
||||
return_value['depth'] = depth
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return return_value
|
||||
|
||||
def render(self, v_deformed_bxnx3=None, sdf_bxn=None, camera_mv_bxnviewx4x4=None, resolution=256):
|
||||
# Here I assume a batch of meshes (can be different mesh and geometry), for the other shapes, the batch is 1
|
||||
v_list = []
|
||||
f_list = []
|
||||
n_batch = v_deformed_bxnx3.shape[0]
|
||||
all_render_output = []
|
||||
for i_batch in range(n_batch):
|
||||
verts_nx3, faces_fx3 = self.get_mesh(v_deformed_bxnx3[i_batch], sdf_bxn[i_batch])
|
||||
v_list.append(verts_nx3)
|
||||
f_list.append(faces_fx3)
|
||||
render_output = self.render_mesh(verts_nx3, faces_fx3, camera_mv_bxnviewx4x4[i_batch], resolution)
|
||||
all_render_output.append(render_output)
|
||||
|
||||
# Concatenate all render output
|
||||
return_keys = all_render_output[0].keys()
|
||||
return_value = dict()
|
||||
for k in return_keys:
|
||||
value = [v[k] for v in all_render_output]
|
||||
return_value[k] = value
|
||||
# We can do concatenation outside of the render
|
||||
return return_value
|
||||
@@ -0,0 +1,20 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
#
|
||||
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
|
||||
# and proprietary rights in and to this software, related documentation
|
||||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
|
||||
|
||||
import torch
|
||||
|
||||
|
||||
def get_center_boundary_index(verts):
|
||||
length_ = torch.sum(verts ** 2, dim=-1)
|
||||
center_idx = torch.argmin(length_)
|
||||
boundary_neg = verts == verts.max()
|
||||
boundary_pos = verts == verts.min()
|
||||
boundary = torch.bitwise_or(boundary_pos, boundary_neg)
|
||||
boundary = torch.sum(boundary.float(), dim=-1)
|
||||
boundary_idx = torch.nonzero(boundary)
|
||||
return center_idx, boundary_idx.squeeze(dim=-1)
|
||||
@@ -0,0 +1,40 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
#
|
||||
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
|
||||
# and proprietary rights in and to this software, related documentation
|
||||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
|
||||
|
||||
import torch
|
||||
import xatlas
|
||||
import numpy as np
|
||||
import nvdiffrast.torch as dr
|
||||
|
||||
|
||||
# ==============================================================================================
|
||||
def interpolate(attr, rast, attr_idx, rast_db=None):
|
||||
return dr.interpolate(attr.contiguous(), rast, attr_idx, rast_db=rast_db, diff_attrs=None if rast_db is None else 'all')
|
||||
|
||||
|
||||
def xatlas_uvmap(ctx, mesh_v, mesh_pos_idx, resolution):
|
||||
vmapping, indices, uvs = xatlas.parametrize(mesh_v.detach().cpu().numpy(), mesh_pos_idx.detach().cpu().numpy())
|
||||
|
||||
# Convert to tensors
|
||||
indices_int64 = indices.astype(np.uint64, casting='same_kind').view(np.int64)
|
||||
|
||||
uvs = torch.tensor(uvs, dtype=torch.float32, device=mesh_v.device)
|
||||
mesh_tex_idx = torch.tensor(indices_int64, dtype=torch.int64, device=mesh_v.device)
|
||||
# mesh_v_tex. ture
|
||||
uv_clip = uvs[None, ...] * 2.0 - 1.0
|
||||
|
||||
# pad to four component coordinate
|
||||
uv_clip4 = torch.cat((uv_clip, torch.zeros_like(uv_clip[..., 0:1]), torch.ones_like(uv_clip[..., 0:1])), dim=-1)
|
||||
|
||||
# rasterize
|
||||
rast, _ = dr.rasterize(ctx, uv_clip4, mesh_tex_idx.int(), (resolution, resolution))
|
||||
|
||||
# Interpolate world space position
|
||||
gb_pos, _ = interpolate(mesh_v[None, ...], rast, mesh_pos_idx.int())
|
||||
mask = rast[..., 3:4] > 0
|
||||
return uvs, mesh_tex_idx, gb_pos, mask
|
||||
@@ -0,0 +1,120 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
#
|
||||
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
|
||||
# and proprietary rights in and to this software, related documentation
|
||||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
|
||||
|
||||
import torch
|
||||
import numpy as np
|
||||
import os
|
||||
from . import Geometry
|
||||
from flexicubes.flexicubes import FlexiCubes
|
||||
from .dmtet import sdf_reg_loss_batch
|
||||
import torch.nn.functional as F
|
||||
|
||||
def get_center_boundary_index(grid_res, device):
|
||||
v = torch.zeros((grid_res + 1, grid_res + 1, grid_res + 1), dtype=torch.bool, device=device)
|
||||
v[grid_res // 2 + 1, grid_res // 2 + 1, grid_res // 2 + 1] = True
|
||||
center_indices = torch.nonzero(v.reshape(-1))
|
||||
|
||||
v[grid_res // 2 + 1, grid_res // 2 + 1, grid_res // 2 + 1] = False
|
||||
v[:2, ...] = True
|
||||
v[-2:, ...] = True
|
||||
v[:, :2, ...] = True
|
||||
v[:, -2:, ...] = True
|
||||
v[:, :, :2] = True
|
||||
v[:, :, -2:] = True
|
||||
boundary_indices = torch.nonzero(v.reshape(-1))
|
||||
return center_indices, boundary_indices
|
||||
|
||||
###############################################################################
|
||||
# Geometry interface
|
||||
###############################################################################
|
||||
class FlexiCubesGeometry(Geometry):
|
||||
def __init__(
|
||||
self, grid_res=64, scale=2.0, device='cuda', renderer=None,
|
||||
render_type='neural_render', args=None):
|
||||
super(FlexiCubesGeometry, self).__init__()
|
||||
self.grid_res = grid_res
|
||||
self.device = device
|
||||
self.args = args
|
||||
self.fc = FlexiCubes(device, weight_scale=0.5)
|
||||
self.verts, self.indices = self.fc.construct_voxel_grid(grid_res)
|
||||
if isinstance(scale, list):
|
||||
self.verts[:, 0] = self.verts[:, 0] * scale[0]
|
||||
self.verts[:, 1] = self.verts[:, 1] * scale[1]
|
||||
self.verts[:, 2] = self.verts[:, 2] * scale[1]
|
||||
else:
|
||||
self.verts = self.verts * scale
|
||||
|
||||
all_edges = self.indices[:, self.fc.cube_edges].reshape(-1, 2)
|
||||
self.all_edges = torch.unique(all_edges, dim=0)
|
||||
|
||||
# Parameters used for fix boundary sdf
|
||||
self.center_indices, self.boundary_indices = get_center_boundary_index(self.grid_res, device)
|
||||
self.renderer = renderer
|
||||
self.render_type = render_type
|
||||
|
||||
def getAABB(self):
|
||||
return torch.min(self.verts, dim=0).values, torch.max(self.verts, dim=0).values
|
||||
|
||||
def get_mesh(self, v_deformed_nx3, sdf_n, weight_n=None, with_uv=False, indices=None, is_training=False):
|
||||
if indices is None:
|
||||
indices = self.indices
|
||||
|
||||
verts, faces, v_reg_loss = self.fc(v_deformed_nx3, sdf_n, indices, self.grid_res,
|
||||
beta_fx12=weight_n[:, :12], alpha_fx8=weight_n[:, 12:20],
|
||||
gamma_f=weight_n[:, 20], training=is_training
|
||||
)
|
||||
return verts, faces, v_reg_loss
|
||||
|
||||
|
||||
def render_mesh(self, mesh_v_nx3, mesh_f_fx3, camera_mv_bx4x4, resolution=256, hierarchical_mask=False):
|
||||
return_value = dict()
|
||||
if self.render_type == 'neural_render':
|
||||
tex_pos, mask, hard_mask, rast, v_pos_clip, mask_pyramid, depth, normal = self.renderer.render_mesh(
|
||||
mesh_v_nx3.unsqueeze(dim=0),
|
||||
mesh_f_fx3.int(),
|
||||
camera_mv_bx4x4,
|
||||
mesh_v_nx3.unsqueeze(dim=0),
|
||||
resolution=resolution,
|
||||
device=self.device,
|
||||
hierarchical_mask=hierarchical_mask
|
||||
)
|
||||
|
||||
return_value['tex_pos'] = tex_pos
|
||||
return_value['mask'] = mask
|
||||
return_value['hard_mask'] = hard_mask
|
||||
return_value['rast'] = rast
|
||||
return_value['v_pos_clip'] = v_pos_clip
|
||||
return_value['mask_pyramid'] = mask_pyramid
|
||||
return_value['depth'] = depth
|
||||
return_value['normal'] = normal
|
||||
else:
|
||||
raise NotImplementedError
|
||||
|
||||
return return_value
|
||||
|
||||
def render(self, v_deformed_bxnx3=None, sdf_bxn=None, camera_mv_bxnviewx4x4=None, resolution=256):
|
||||
# Here I assume a batch of meshes (can be different mesh and geometry), for the other shapes, the batch is 1
|
||||
v_list = []
|
||||
f_list = []
|
||||
n_batch = v_deformed_bxnx3.shape[0]
|
||||
all_render_output = []
|
||||
for i_batch in range(n_batch):
|
||||
verts_nx3, faces_fx3 = self.get_mesh(v_deformed_bxnx3[i_batch], sdf_bxn[i_batch])
|
||||
v_list.append(verts_nx3)
|
||||
f_list.append(faces_fx3)
|
||||
render_output = self.render_mesh(verts_nx3, faces_fx3, camera_mv_bxnviewx4x4[i_batch], resolution)
|
||||
all_render_output.append(render_output)
|
||||
|
||||
# Concatenate all render output
|
||||
return_keys = all_render_output[0].keys()
|
||||
return_value = dict()
|
||||
for k in return_keys:
|
||||
value = [v[k] for v in all_render_output]
|
||||
return_value[k] = value
|
||||
# We can do concatenation outside of the render
|
||||
return return_value
|
||||
@@ -0,0 +1,209 @@
|
||||
# Copyright (c) 2023, Zexin He
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# https://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import mcubes
|
||||
import nvdiffrast.torch as dr
|
||||
from einops import rearrange, repeat
|
||||
|
||||
from .encoder.dino_wrapper import DinoWrapper
|
||||
from .decoder.transformer import TriplaneTransformer
|
||||
from .renderer.synthesizer import TriplaneSynthesizer
|
||||
from ..utils.mesh_util import xatlas_uvmap
|
||||
|
||||
|
||||
class InstantNeRF(nn.Module):
|
||||
"""
|
||||
Full model of the large reconstruction model.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
encoder_freeze: bool = False,
|
||||
encoder_model_name: str = 'facebook/dino-vitb16',
|
||||
encoder_feat_dim: int = 768,
|
||||
transformer_dim: int = 1024,
|
||||
transformer_layers: int = 16,
|
||||
transformer_heads: int = 16,
|
||||
triplane_low_res: int = 32,
|
||||
triplane_high_res: int = 64,
|
||||
triplane_dim: int = 80,
|
||||
rendering_samples_per_ray: int = 128,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# modules
|
||||
self.encoder = DinoWrapper(
|
||||
model_name=encoder_model_name,
|
||||
freeze=encoder_freeze,
|
||||
)
|
||||
|
||||
self.transformer = TriplaneTransformer(
|
||||
inner_dim=transformer_dim,
|
||||
num_layers=transformer_layers,
|
||||
num_heads=transformer_heads,
|
||||
image_feat_dim=encoder_feat_dim,
|
||||
triplane_low_res=triplane_low_res,
|
||||
triplane_high_res=triplane_high_res,
|
||||
triplane_dim=triplane_dim,
|
||||
)
|
||||
|
||||
self.synthesizer = TriplaneSynthesizer(
|
||||
triplane_dim=triplane_dim,
|
||||
samples_per_ray=rendering_samples_per_ray,
|
||||
)
|
||||
|
||||
def forward_planes(self, images, cameras):
|
||||
# images: [B, V, C_img, H_img, W_img]
|
||||
# cameras: [B, V, 16]
|
||||
B = images.shape[0]
|
||||
|
||||
# encode images
|
||||
image_feats = self.encoder(images, cameras)
|
||||
image_feats = rearrange(image_feats, '(b v) l d -> b (v l) d', b=B)
|
||||
|
||||
# transformer generating planes
|
||||
planes = self.transformer(image_feats)
|
||||
|
||||
return planes
|
||||
|
||||
def forward_synthesizer(self, planes, render_cameras, render_size: int):
|
||||
render_results = self.synthesizer(
|
||||
planes,
|
||||
render_cameras,
|
||||
render_size,
|
||||
)
|
||||
return render_results
|
||||
|
||||
def forward(self, images, cameras, render_cameras, render_size: int):
|
||||
# images: [B, V, C_img, H_img, W_img]
|
||||
# cameras: [B, V, 16]
|
||||
# render_cameras: [B, M, D_cam_render]
|
||||
# render_size: int
|
||||
B, M = render_cameras.shape[:2]
|
||||
|
||||
planes = self.forward_planes(images, cameras)
|
||||
|
||||
# render target views
|
||||
render_results = self.synthesizer(planes, render_cameras, render_size)
|
||||
|
||||
return {
|
||||
'planes': planes,
|
||||
**render_results,
|
||||
}
|
||||
|
||||
def get_texture_prediction(self, planes, tex_pos, hard_mask=None):
|
||||
'''
|
||||
Predict Texture given triplanes
|
||||
:param planes: the triplane feature map
|
||||
:param tex_pos: Position we want to query the texture field
|
||||
:param hard_mask: 2D silhoueete of the rendered image
|
||||
'''
|
||||
tex_pos = torch.cat(tex_pos, dim=0)
|
||||
if not hard_mask is None:
|
||||
tex_pos = tex_pos * hard_mask.float()
|
||||
batch_size = tex_pos.shape[0]
|
||||
tex_pos = tex_pos.reshape(batch_size, -1, 3)
|
||||
###################
|
||||
# We use mask to get the texture location (to save the memory)
|
||||
if hard_mask is not None:
|
||||
n_point_list = torch.sum(hard_mask.long().reshape(hard_mask.shape[0], -1), dim=-1)
|
||||
sample_tex_pose_list = []
|
||||
max_point = n_point_list.max()
|
||||
expanded_hard_mask = hard_mask.reshape(batch_size, -1, 1).expand(-1, -1, 3) > 0.5
|
||||
for i in range(tex_pos.shape[0]):
|
||||
tex_pos_one_shape = tex_pos[i][expanded_hard_mask[i]].reshape(1, -1, 3)
|
||||
if tex_pos_one_shape.shape[1] < max_point:
|
||||
tex_pos_one_shape = torch.cat(
|
||||
[tex_pos_one_shape, torch.zeros(
|
||||
1, max_point - tex_pos_one_shape.shape[1], 3,
|
||||
device=tex_pos_one_shape.device, dtype=torch.float32)], dim=1)
|
||||
sample_tex_pose_list.append(tex_pos_one_shape)
|
||||
tex_pos = torch.cat(sample_tex_pose_list, dim=0)
|
||||
|
||||
tex_feat = torch.utils.checkpoint.checkpoint(
|
||||
self.synthesizer.forward_points,
|
||||
planes,
|
||||
tex_pos,
|
||||
use_reentrant=False,
|
||||
)['rgb']
|
||||
|
||||
if hard_mask is not None:
|
||||
final_tex_feat = torch.zeros(
|
||||
planes.shape[0], hard_mask.shape[1] * hard_mask.shape[2], tex_feat.shape[-1], device=tex_feat.device)
|
||||
expanded_hard_mask = hard_mask.reshape(hard_mask.shape[0], -1, 1).expand(-1, -1, final_tex_feat.shape[-1]) > 0.5
|
||||
for i in range(planes.shape[0]):
|
||||
final_tex_feat[i][expanded_hard_mask[i]] = tex_feat[i][:n_point_list[i]].reshape(-1)
|
||||
tex_feat = final_tex_feat
|
||||
|
||||
return tex_feat.reshape(planes.shape[0], hard_mask.shape[1], hard_mask.shape[2], tex_feat.shape[-1])
|
||||
|
||||
def extract_mesh(
|
||||
self,
|
||||
planes: torch.Tensor,
|
||||
mesh_resolution: int = 256,
|
||||
mesh_threshold: int = 10.0,
|
||||
use_texture_map: bool = False,
|
||||
texture_resolution: int = 1024,
|
||||
**kwargs,
|
||||
):
|
||||
'''
|
||||
Extract a 3D mesh from triplane nerf. Only support batch_size 1.
|
||||
:param planes: triplane features
|
||||
:param mesh_resolution: marching cubes resolution
|
||||
:param mesh_threshold: iso-surface threshold
|
||||
:param use_texture_map: use texture map or vertex color
|
||||
:param texture_resolution: the resolution of texture map
|
||||
'''
|
||||
assert planes.shape[0] == 1
|
||||
device = planes.device
|
||||
|
||||
grid_out = self.synthesizer.forward_grid(
|
||||
planes=planes,
|
||||
grid_size=mesh_resolution,
|
||||
)
|
||||
|
||||
vertices, faces = mcubes.marching_cubes(
|
||||
grid_out['sigma'].squeeze(0).squeeze(-1).cpu().numpy(),
|
||||
mesh_threshold,
|
||||
)
|
||||
vertices = vertices / (mesh_resolution - 1) * 2 - 1
|
||||
|
||||
if not use_texture_map:
|
||||
# query vertex colors
|
||||
vertices_tensor = torch.tensor(vertices, dtype=torch.float32, device=device).unsqueeze(0)
|
||||
vertices_colors = self.synthesizer.forward_points(
|
||||
planes, vertices_tensor)['rgb'].squeeze(0).cpu().numpy()
|
||||
vertices_colors = (vertices_colors * 255).astype(np.uint8)
|
||||
|
||||
return vertices, faces, vertices_colors
|
||||
|
||||
# use x-atlas to get uv mapping for the mesh
|
||||
vertices = torch.tensor(vertices, dtype=torch.float32, device=device)
|
||||
faces = torch.tensor(faces.astype(int), dtype=torch.long, device=device)
|
||||
|
||||
ctx = dr.RasterizeCudaContext(device=device)
|
||||
uvs, mesh_tex_idx, gb_pos, tex_hard_mask = xatlas_uvmap(
|
||||
ctx, vertices, faces, resolution=texture_resolution)
|
||||
tex_hard_mask = tex_hard_mask.float()
|
||||
|
||||
# query the texture field to get the RGB color for texture map
|
||||
tex_feat = self.get_texture_prediction(
|
||||
planes, [gb_pos], tex_hard_mask)
|
||||
background_feature = torch.zeros_like(tex_feat)
|
||||
img_feat = torch.lerp(background_feature, tex_feat, tex_hard_mask)
|
||||
texture_map = img_feat.permute(0, 3, 1, 2).squeeze(0)
|
||||
|
||||
return vertices, faces, uvs, mesh_tex_idx, texture_map
|
||||
@@ -0,0 +1,382 @@
|
||||
# Copyright (c) 2023, Tencent Inc
|
||||
#
|
||||
# Licensed under the Apache License, Version 2.0 (the "License");
|
||||
# you may not use this file except in compliance with the License.
|
||||
# You may obtain a copy of the License at
|
||||
#
|
||||
# https://www.apache.org/licenses/LICENSE-2.0
|
||||
#
|
||||
# Unless required by applicable law or agreed to in writing, software
|
||||
# distributed under the License is distributed on an "AS IS" BASIS,
|
||||
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
||||
# See the License for the specific language governing permissions and
|
||||
# limitations under the License.
|
||||
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import nvdiffrast.torch as dr
|
||||
from einops import rearrange, repeat
|
||||
|
||||
from .encoder.dino_wrapper import DinoWrapper
|
||||
from .decoder.transformer import TriplaneTransformer
|
||||
from .renderer.synthesizer_mesh import TriplaneSynthesizer
|
||||
from .geometry.camera.perspective_camera import PerspectiveCamera
|
||||
from .geometry.render.neural_render import NeuralRender
|
||||
from .geometry.rep_3d.flexicubes_geometry import FlexiCubesGeometry
|
||||
from ..utils.mesh_util import xatlas_uvmap
|
||||
|
||||
|
||||
class InstantMesh(nn.Module):
|
||||
"""
|
||||
Full model of the large reconstruction model.
|
||||
"""
|
||||
def __init__(
|
||||
self,
|
||||
encoder_freeze: bool = False,
|
||||
encoder_model_name: str = 'facebook/dino-vitb16',
|
||||
encoder_feat_dim: int = 768,
|
||||
transformer_dim: int = 1024,
|
||||
transformer_layers: int = 16,
|
||||
transformer_heads: int = 16,
|
||||
triplane_low_res: int = 32,
|
||||
triplane_high_res: int = 64,
|
||||
triplane_dim: int = 80,
|
||||
rendering_samples_per_ray: int = 128,
|
||||
grid_res: int = 128,
|
||||
grid_scale: float = 2.0,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
# attributes
|
||||
self.grid_res = grid_res
|
||||
self.grid_scale = grid_scale
|
||||
self.deformation_multiplier = 4.0
|
||||
|
||||
# modules
|
||||
self.encoder = DinoWrapper(
|
||||
model_name=encoder_model_name,
|
||||
freeze=encoder_freeze,
|
||||
)
|
||||
|
||||
self.transformer = TriplaneTransformer(
|
||||
inner_dim=transformer_dim,
|
||||
num_layers=transformer_layers,
|
||||
num_heads=transformer_heads,
|
||||
image_feat_dim=encoder_feat_dim,
|
||||
triplane_low_res=triplane_low_res,
|
||||
triplane_high_res=triplane_high_res,
|
||||
triplane_dim=triplane_dim,
|
||||
)
|
||||
|
||||
self.synthesizer = TriplaneSynthesizer(
|
||||
triplane_dim=triplane_dim,
|
||||
samples_per_ray=rendering_samples_per_ray,
|
||||
)
|
||||
|
||||
def init_flexicubes_geometry(self, device, fovy=50.0):
|
||||
camera = PerspectiveCamera(fovy=fovy, device=device)
|
||||
renderer = NeuralRender(device, camera_model=camera)
|
||||
self.geometry = FlexiCubesGeometry(
|
||||
grid_res=self.grid_res,
|
||||
scale=self.grid_scale,
|
||||
renderer=renderer,
|
||||
render_type='neural_render',
|
||||
device=device,
|
||||
)
|
||||
|
||||
def forward_planes(self, images, cameras):
|
||||
# images: [B, V, C_img, H_img, W_img]
|
||||
# cameras: [B, V, 16]
|
||||
B = images.shape[0]
|
||||
|
||||
# encode images
|
||||
image_feats = self.encoder(images, cameras)
|
||||
image_feats = rearrange(image_feats, '(b v) l d -> b (v l) d', b=B)
|
||||
|
||||
# decode triplanes
|
||||
planes = self.transformer(image_feats)
|
||||
|
||||
return planes
|
||||
|
||||
def get_sdf_deformation_prediction(self, planes):
|
||||
'''
|
||||
Predict SDF and deformation for tetrahedron vertices
|
||||
:param planes: triplane feature map for the geometry
|
||||
'''
|
||||
init_position = self.geometry.verts.unsqueeze(0).expand(planes.shape[0], -1, -1)
|
||||
|
||||
# Step 1: predict the SDF and deformation
|
||||
sdf, deformation, weight = torch.utils.checkpoint.checkpoint(
|
||||
self.synthesizer.get_geometry_prediction,
|
||||
planes,
|
||||
init_position,
|
||||
self.geometry.indices,
|
||||
use_reentrant=False,
|
||||
)
|
||||
|
||||
# Step 2: Normalize the deformation to avoid the flipped triangles.
|
||||
deformation = 1.0 / (self.grid_res * self.deformation_multiplier) * torch.tanh(deformation)
|
||||
sdf_reg_loss = torch.zeros(sdf.shape[0], device=sdf.device, dtype=torch.float32)
|
||||
|
||||
####
|
||||
# Step 3: Fix some sdf if we observe empty shape (full positive or full negative)
|
||||
sdf_bxnxnxn = sdf.reshape((sdf.shape[0], self.grid_res + 1, self.grid_res + 1, self.grid_res + 1))
|
||||
sdf_less_boundary = sdf_bxnxnxn[:, 1:-1, 1:-1, 1:-1].reshape(sdf.shape[0], -1)
|
||||
pos_shape = torch.sum((sdf_less_boundary > 0).int(), dim=-1)
|
||||
neg_shape = torch.sum((sdf_less_boundary < 0).int(), dim=-1)
|
||||
zero_surface = torch.bitwise_or(pos_shape == 0, neg_shape == 0)
|
||||
if torch.sum(zero_surface).item() > 0:
|
||||
update_sdf = torch.zeros_like(sdf[0:1])
|
||||
max_sdf = sdf.max()
|
||||
min_sdf = sdf.min()
|
||||
update_sdf[:, self.geometry.center_indices] += (1.0 - min_sdf) # greater than zero
|
||||
update_sdf[:, self.geometry.boundary_indices] += (-1 - max_sdf) # smaller than zero
|
||||
new_sdf = torch.zeros_like(sdf)
|
||||
for i_batch in range(zero_surface.shape[0]):
|
||||
if zero_surface[i_batch]:
|
||||
new_sdf[i_batch:i_batch + 1] += update_sdf
|
||||
update_mask = (new_sdf == 0).float()
|
||||
# Regulraization here is used to push the sdf to be a different sign (make it not fully positive or fully negative)
|
||||
sdf_reg_loss = torch.abs(sdf).mean(dim=-1).mean(dim=-1)
|
||||
sdf_reg_loss = sdf_reg_loss * zero_surface.float()
|
||||
sdf = sdf * update_mask + new_sdf * (1 - update_mask)
|
||||
|
||||
# Step 4: Here we remove the gradient for the bad sdf (full positive or full negative)
|
||||
final_sdf = []
|
||||
final_def = []
|
||||
for i_batch in range(zero_surface.shape[0]):
|
||||
if zero_surface[i_batch]:
|
||||
final_sdf.append(sdf[i_batch: i_batch + 1].detach())
|
||||
final_def.append(deformation[i_batch: i_batch + 1].detach())
|
||||
else:
|
||||
final_sdf.append(sdf[i_batch: i_batch + 1])
|
||||
final_def.append(deformation[i_batch: i_batch + 1])
|
||||
sdf = torch.cat(final_sdf, dim=0)
|
||||
deformation = torch.cat(final_def, dim=0)
|
||||
return sdf, deformation, sdf_reg_loss, weight
|
||||
|
||||
def get_geometry_prediction(self, planes=None):
|
||||
'''
|
||||
Function to generate mesh with give triplanes
|
||||
:param planes: triplane features
|
||||
'''
|
||||
# Step 1: first get the sdf and deformation value for each vertices in the tetrahedon grid.
|
||||
sdf, deformation, sdf_reg_loss, weight = self.get_sdf_deformation_prediction(planes)
|
||||
v_deformed = self.geometry.verts.unsqueeze(dim=0).expand(sdf.shape[0], -1, -1) + deformation
|
||||
tets = self.geometry.indices
|
||||
n_batch = planes.shape[0]
|
||||
v_list = []
|
||||
f_list = []
|
||||
flexicubes_surface_reg_list = []
|
||||
|
||||
# Step 2: Using marching tet to obtain the mesh
|
||||
for i_batch in range(n_batch):
|
||||
verts, faces, flexicubes_surface_reg = self.geometry.get_mesh(
|
||||
v_deformed[i_batch],
|
||||
sdf[i_batch].squeeze(dim=-1),
|
||||
with_uv=False,
|
||||
indices=tets,
|
||||
weight_n=weight[i_batch].squeeze(dim=-1),
|
||||
is_training=self.training,
|
||||
)
|
||||
flexicubes_surface_reg_list.append(flexicubes_surface_reg)
|
||||
v_list.append(verts)
|
||||
f_list.append(faces)
|
||||
|
||||
flexicubes_surface_reg = torch.cat(flexicubes_surface_reg_list).mean()
|
||||
flexicubes_weight_reg = (weight ** 2).mean()
|
||||
|
||||
return v_list, f_list, sdf, deformation, v_deformed, (sdf_reg_loss, flexicubes_surface_reg, flexicubes_weight_reg)
|
||||
|
||||
def get_texture_prediction(self, planes, tex_pos, hard_mask=None):
|
||||
'''
|
||||
Predict Texture given triplanes
|
||||
:param planes: the triplane feature map
|
||||
:param tex_pos: Position we want to query the texture field
|
||||
:param hard_mask: 2D silhoueete of the rendered image
|
||||
'''
|
||||
tex_pos = torch.cat(tex_pos, dim=0)
|
||||
if not hard_mask is None:
|
||||
tex_pos = tex_pos * hard_mask.float()
|
||||
batch_size = tex_pos.shape[0]
|
||||
tex_pos = tex_pos.reshape(batch_size, -1, 3)
|
||||
###################
|
||||
# We use mask to get the texture location (to save the memory)
|
||||
if hard_mask is not None:
|
||||
n_point_list = torch.sum(hard_mask.long().reshape(hard_mask.shape[0], -1), dim=-1)
|
||||
sample_tex_pose_list = []
|
||||
max_point = n_point_list.max()
|
||||
expanded_hard_mask = hard_mask.reshape(batch_size, -1, 1).expand(-1, -1, 3) > 0.5
|
||||
for i in range(tex_pos.shape[0]):
|
||||
tex_pos_one_shape = tex_pos[i][expanded_hard_mask[i]].reshape(1, -1, 3)
|
||||
if tex_pos_one_shape.shape[1] < max_point:
|
||||
tex_pos_one_shape = torch.cat(
|
||||
[tex_pos_one_shape, torch.zeros(
|
||||
1, max_point - tex_pos_one_shape.shape[1], 3,
|
||||
device=tex_pos_one_shape.device, dtype=torch.float32)], dim=1)
|
||||
sample_tex_pose_list.append(tex_pos_one_shape)
|
||||
tex_pos = torch.cat(sample_tex_pose_list, dim=0)
|
||||
|
||||
tex_feat = torch.utils.checkpoint.checkpoint(
|
||||
self.synthesizer.get_texture_prediction,
|
||||
planes,
|
||||
tex_pos,
|
||||
use_reentrant=False,
|
||||
)
|
||||
|
||||
if hard_mask is not None:
|
||||
final_tex_feat = torch.zeros(
|
||||
planes.shape[0], hard_mask.shape[1] * hard_mask.shape[2], tex_feat.shape[-1], device=tex_feat.device)
|
||||
expanded_hard_mask = hard_mask.reshape(hard_mask.shape[0], -1, 1).expand(-1, -1, final_tex_feat.shape[-1]) > 0.5
|
||||
for i in range(planes.shape[0]):
|
||||
final_tex_feat[i][expanded_hard_mask[i]] = tex_feat[i][:n_point_list[i]].reshape(-1)
|
||||
tex_feat = final_tex_feat
|
||||
|
||||
return tex_feat.reshape(planes.shape[0], hard_mask.shape[1], hard_mask.shape[2], tex_feat.shape[-1])
|
||||
|
||||
def render_mesh(self, mesh_v, mesh_f, cam_mv, render_size=256):
|
||||
'''
|
||||
Function to render a generated mesh with nvdiffrast
|
||||
:param mesh_v: List of vertices for the mesh
|
||||
:param mesh_f: List of faces for the mesh
|
||||
:param cam_mv: 4x4 rotation matrix
|
||||
:return:
|
||||
'''
|
||||
return_value_list = []
|
||||
for i_mesh in range(len(mesh_v)):
|
||||
return_value = self.geometry.render_mesh(
|
||||
mesh_v[i_mesh],
|
||||
mesh_f[i_mesh].int(),
|
||||
cam_mv[i_mesh],
|
||||
resolution=render_size,
|
||||
hierarchical_mask=False
|
||||
)
|
||||
return_value_list.append(return_value)
|
||||
|
||||
return_keys = return_value_list[0].keys()
|
||||
return_value = dict()
|
||||
for k in return_keys:
|
||||
value = [v[k] for v in return_value_list]
|
||||
return_value[k] = value
|
||||
|
||||
mask = torch.cat(return_value['mask'], dim=0)
|
||||
hard_mask = torch.cat(return_value['hard_mask'], dim=0)
|
||||
tex_pos = return_value['tex_pos']
|
||||
depth = torch.cat(return_value['depth'], dim=0)
|
||||
normal = torch.cat(return_value['normal'], dim=0)
|
||||
return mask, hard_mask, tex_pos, depth, normal
|
||||
|
||||
def forward_geometry(self, planes, render_cameras, render_size=256):
|
||||
'''
|
||||
Main function of our Generator. It first generate 3D mesh, then render it into 2D image
|
||||
with given `render_cameras`.
|
||||
:param planes: triplane features
|
||||
:param render_cameras: cameras to render generated 3D shape
|
||||
'''
|
||||
B, NV = render_cameras.shape[:2]
|
||||
|
||||
# Generate 3D mesh first
|
||||
mesh_v, mesh_f, sdf, deformation, v_deformed, sdf_reg_loss = self.get_geometry_prediction(planes)
|
||||
|
||||
# Render the mesh into 2D image (get 3d position of each image plane)
|
||||
cam_mv = render_cameras
|
||||
run_n_view = cam_mv.shape[1]
|
||||
antilias_mask, hard_mask, tex_pos, depth, normal = self.render_mesh(mesh_v, mesh_f, cam_mv, render_size=render_size)
|
||||
|
||||
tex_hard_mask = hard_mask
|
||||
tex_pos = [torch.cat([pos[i_view:i_view + 1] for i_view in range(run_n_view)], dim=2) for pos in tex_pos]
|
||||
tex_hard_mask = torch.cat(
|
||||
[torch.cat(
|
||||
[tex_hard_mask[i * run_n_view + i_view: i * run_n_view + i_view + 1]
|
||||
for i_view in range(run_n_view)], dim=2)
|
||||
for i in range(planes.shape[0])], dim=0)
|
||||
|
||||
# Querying the texture field to predict the texture feature for each pixel on the image
|
||||
tex_feat = self.get_texture_prediction(planes, tex_pos, tex_hard_mask)
|
||||
background_feature = torch.ones_like(tex_feat) # white background
|
||||
|
||||
# Merge them together
|
||||
img_feat = tex_feat * tex_hard_mask + background_feature * (1 - tex_hard_mask)
|
||||
|
||||
# We should split it back to the original image shape
|
||||
img_feat = torch.cat(
|
||||
[torch.cat(
|
||||
[img_feat[i:i + 1, :, render_size * i_view: render_size * (i_view + 1)]
|
||||
for i_view in range(run_n_view)], dim=0) for i in range(len(tex_pos))], dim=0)
|
||||
|
||||
img = img_feat.clamp(0, 1).permute(0, 3, 1, 2).unflatten(0, (B, NV))
|
||||
antilias_mask = antilias_mask.permute(0, 3, 1, 2).unflatten(0, (B, NV))
|
||||
depth = -depth.permute(0, 3, 1, 2).unflatten(0, (B, NV)) # transform negative depth to positive
|
||||
normal = normal.permute(0, 3, 1, 2).unflatten(0, (B, NV))
|
||||
|
||||
out = {
|
||||
'img': img,
|
||||
'mask': antilias_mask,
|
||||
'depth': depth,
|
||||
'normal': normal,
|
||||
'sdf': sdf,
|
||||
'mesh_v': mesh_v,
|
||||
'mesh_f': mesh_f,
|
||||
'sdf_reg_loss': sdf_reg_loss,
|
||||
}
|
||||
return out
|
||||
|
||||
def forward(self, images, cameras, render_cameras, render_size: int):
|
||||
# images: [B, V, C_img, H_img, W_img]
|
||||
# cameras: [B, V, 16]
|
||||
# render_cameras: [B, M, D_cam_render]
|
||||
# render_size: int
|
||||
B, M = render_cameras.shape[:2]
|
||||
|
||||
planes = self.forward_planes(images, cameras)
|
||||
out = self.forward_geometry(planes, render_cameras, render_size=render_size)
|
||||
|
||||
return {
|
||||
'planes': planes,
|
||||
**out
|
||||
}
|
||||
|
||||
def extract_mesh(
|
||||
self,
|
||||
planes: torch.Tensor,
|
||||
use_texture_map: bool = False,
|
||||
texture_resolution: int = 1024,
|
||||
**kwargs,
|
||||
):
|
||||
'''
|
||||
Extract a 3D mesh from FlexiCubes. Only support batch_size 1.
|
||||
:param planes: triplane features
|
||||
:param use_texture_map: use texture map or vertex color
|
||||
:param texture_resolution: the resolution of texure map
|
||||
'''
|
||||
assert planes.shape[0] == 1
|
||||
device = planes.device
|
||||
|
||||
# predict geometry first
|
||||
mesh_v, mesh_f, sdf, deformation, v_deformed, sdf_reg_loss = self.get_geometry_prediction(planes)
|
||||
vertices, faces = mesh_v[0], mesh_f[0]
|
||||
|
||||
if not use_texture_map:
|
||||
# query vertex colors
|
||||
vertices_tensor = vertices.unsqueeze(0)
|
||||
vertices_colors = self.synthesizer.get_texture_prediction(
|
||||
planes, vertices_tensor).clamp(0, 1).squeeze(0).cpu().numpy()
|
||||
vertices_colors = (vertices_colors * 255).astype(np.uint8)
|
||||
|
||||
return vertices.cpu().numpy(), faces.cpu().numpy(), vertices_colors
|
||||
|
||||
# use x-atlas to get uv mapping for the mesh
|
||||
ctx = dr.RasterizeCudaContext(device=device)
|
||||
uvs, mesh_tex_idx, gb_pos, tex_hard_mask = xatlas_uvmap(
|
||||
self.geometry.renderer.ctx, vertices, faces, resolution=texture_resolution)
|
||||
tex_hard_mask = tex_hard_mask.float()
|
||||
|
||||
# query the texture field to get the RGB color for texture map
|
||||
tex_feat = self.get_texture_prediction(
|
||||
planes, [gb_pos], tex_hard_mask)
|
||||
background_feature = torch.zeros_like(tex_feat)
|
||||
img_feat = torch.lerp(background_feature, tex_feat, tex_hard_mask)
|
||||
texture_map = img_feat.permute(0, 3, 1, 2).squeeze(0)
|
||||
|
||||
return vertices, faces, uvs, mesh_tex_idx, texture_map
|
||||
@@ -0,0 +1,9 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
||||
#
|
||||
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
|
||||
# property and proprietary rights in and to this material, related
|
||||
# documentation and any modifications thereto. Any use, reproduction,
|
||||
# disclosure or distribution of this material and related documentation
|
||||
# without an express license agreement from NVIDIA CORPORATION or
|
||||
# its affiliates is strictly prohibited.
|
||||
@@ -0,0 +1,203 @@
|
||||
# ORIGINAL LICENSE
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
||||
#
|
||||
# Modified by Jiale Xu
|
||||
# The modifications are subject to the same license as the original.
|
||||
|
||||
|
||||
import itertools
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .utils.renderer import ImportanceRenderer
|
||||
from .utils.ray_sampler import RaySampler
|
||||
|
||||
|
||||
class OSGDecoder(nn.Module):
|
||||
"""
|
||||
Triplane decoder that gives RGB and sigma values from sampled features.
|
||||
Using ReLU here instead of Softplus in the original implementation.
|
||||
|
||||
Reference:
|
||||
EG3D: https://github.com/NVlabs/eg3d/blob/main/eg3d/training/triplane.py#L112
|
||||
"""
|
||||
def __init__(self, n_features: int,
|
||||
hidden_dim: int = 64, num_layers: int = 4, activation: nn.Module = nn.ReLU):
|
||||
super().__init__()
|
||||
self.net = nn.Sequential(
|
||||
nn.Linear(3 * n_features, hidden_dim),
|
||||
activation(),
|
||||
*itertools.chain(*[[
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
activation(),
|
||||
] for _ in range(num_layers - 2)]),
|
||||
nn.Linear(hidden_dim, 1 + 3),
|
||||
)
|
||||
# init all bias to zero
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def forward(self, sampled_features, ray_directions):
|
||||
# Aggregate features by mean
|
||||
# sampled_features = sampled_features.mean(1)
|
||||
# Aggregate features by concatenation
|
||||
_N, n_planes, _M, _C = sampled_features.shape
|
||||
sampled_features = sampled_features.permute(0, 2, 1, 3).reshape(_N, _M, n_planes*_C)
|
||||
x = sampled_features
|
||||
|
||||
N, M, C = x.shape
|
||||
x = x.contiguous().view(N*M, C)
|
||||
|
||||
x = self.net(x)
|
||||
x = x.view(N, M, -1)
|
||||
rgb = torch.sigmoid(x[..., 1:])*(1 + 2*0.001) - 0.001 # Uses sigmoid clamping from MipNeRF
|
||||
sigma = x[..., 0:1]
|
||||
|
||||
return {'rgb': rgb, 'sigma': sigma}
|
||||
|
||||
|
||||
class TriplaneSynthesizer(nn.Module):
|
||||
"""
|
||||
Synthesizer that renders a triplane volume with planes and a camera.
|
||||
|
||||
Reference:
|
||||
EG3D: https://github.com/NVlabs/eg3d/blob/main/eg3d/training/triplane.py#L19
|
||||
"""
|
||||
|
||||
DEFAULT_RENDERING_KWARGS = {
|
||||
'ray_start': 'auto',
|
||||
'ray_end': 'auto',
|
||||
'box_warp': 2.,
|
||||
'white_back': True,
|
||||
'disparity_space_sampling': False,
|
||||
'clamp_mode': 'softplus',
|
||||
'sampler_bbox_min': -1.,
|
||||
'sampler_bbox_max': 1.,
|
||||
}
|
||||
|
||||
def __init__(self, triplane_dim: int, samples_per_ray: int):
|
||||
super().__init__()
|
||||
|
||||
# attributes
|
||||
self.triplane_dim = triplane_dim
|
||||
self.rendering_kwargs = {
|
||||
**self.DEFAULT_RENDERING_KWARGS,
|
||||
'depth_resolution': samples_per_ray // 2,
|
||||
'depth_resolution_importance': samples_per_ray // 2,
|
||||
}
|
||||
|
||||
# renderings
|
||||
self.renderer = ImportanceRenderer()
|
||||
self.ray_sampler = RaySampler()
|
||||
|
||||
# modules
|
||||
self.decoder = OSGDecoder(n_features=triplane_dim)
|
||||
|
||||
def forward(self, planes, cameras, render_size=128, crop_params=None):
|
||||
# planes: (N, 3, D', H', W')
|
||||
# cameras: (N, M, D_cam)
|
||||
# render_size: int
|
||||
assert planes.shape[0] == cameras.shape[0], "Batch size mismatch for planes and cameras"
|
||||
N, M = cameras.shape[:2]
|
||||
|
||||
cam2world_matrix = cameras[..., :16].view(N, M, 4, 4)
|
||||
intrinsics = cameras[..., 16:25].view(N, M, 3, 3)
|
||||
|
||||
# Create a batch of rays for volume rendering
|
||||
ray_origins, ray_directions = self.ray_sampler(
|
||||
cam2world_matrix=cam2world_matrix.reshape(-1, 4, 4),
|
||||
intrinsics=intrinsics.reshape(-1, 3, 3),
|
||||
render_size=render_size,
|
||||
)
|
||||
assert N*M == ray_origins.shape[0], "Batch size mismatch for ray_origins"
|
||||
assert ray_origins.dim() == 3, "ray_origins should be 3-dimensional"
|
||||
|
||||
# Crop rays if crop_params is available
|
||||
if crop_params is not None:
|
||||
ray_origins = ray_origins.reshape(N*M, render_size, render_size, 3)
|
||||
ray_directions = ray_directions.reshape(N*M, render_size, render_size, 3)
|
||||
i, j, h, w = crop_params
|
||||
ray_origins = ray_origins[:, i:i+h, j:j+w, :].reshape(N*M, -1, 3)
|
||||
ray_directions = ray_directions[:, i:i+h, j:j+w, :].reshape(N*M, -1, 3)
|
||||
|
||||
# Perform volume rendering
|
||||
rgb_samples, depth_samples, weights_samples = self.renderer(
|
||||
planes.repeat_interleave(M, dim=0), self.decoder, ray_origins, ray_directions, self.rendering_kwargs,
|
||||
)
|
||||
|
||||
# Reshape into 'raw' neural-rendered image
|
||||
if crop_params is not None:
|
||||
Himg, Wimg = crop_params[2:]
|
||||
else:
|
||||
Himg = Wimg = render_size
|
||||
rgb_images = rgb_samples.permute(0, 2, 1).reshape(N, M, rgb_samples.shape[-1], Himg, Wimg).contiguous()
|
||||
depth_images = depth_samples.permute(0, 2, 1).reshape(N, M, 1, Himg, Wimg)
|
||||
weight_images = weights_samples.permute(0, 2, 1).reshape(N, M, 1, Himg, Wimg)
|
||||
|
||||
out = {
|
||||
'images_rgb': rgb_images,
|
||||
'images_depth': depth_images,
|
||||
'images_weight': weight_images,
|
||||
}
|
||||
return out
|
||||
|
||||
def forward_grid(self, planes, grid_size: int, aabb: torch.Tensor = None):
|
||||
# planes: (N, 3, D', H', W')
|
||||
# grid_size: int
|
||||
# aabb: (N, 2, 3)
|
||||
if aabb is None:
|
||||
aabb = torch.tensor([
|
||||
[self.rendering_kwargs['sampler_bbox_min']] * 3,
|
||||
[self.rendering_kwargs['sampler_bbox_max']] * 3,
|
||||
], device=planes.device, dtype=planes.dtype).unsqueeze(0).repeat(planes.shape[0], 1, 1)
|
||||
assert planes.shape[0] == aabb.shape[0], "Batch size mismatch for planes and aabb"
|
||||
N = planes.shape[0]
|
||||
|
||||
# create grid points for triplane query
|
||||
grid_points = []
|
||||
for i in range(N):
|
||||
grid_points.append(torch.stack(torch.meshgrid(
|
||||
torch.linspace(aabb[i, 0, 0], aabb[i, 1, 0], grid_size, device=planes.device),
|
||||
torch.linspace(aabb[i, 0, 1], aabb[i, 1, 1], grid_size, device=planes.device),
|
||||
torch.linspace(aabb[i, 0, 2], aabb[i, 1, 2], grid_size, device=planes.device),
|
||||
indexing='ij',
|
||||
), dim=-1).reshape(-1, 3))
|
||||
cube_grid = torch.stack(grid_points, dim=0).to(planes.device)
|
||||
|
||||
features = self.forward_points(planes, cube_grid)
|
||||
|
||||
# reshape into grid
|
||||
features = {
|
||||
k: v.reshape(N, grid_size, grid_size, grid_size, -1)
|
||||
for k, v in features.items()
|
||||
}
|
||||
return features
|
||||
|
||||
def forward_points(self, planes, points: torch.Tensor, chunk_size: int = 2**20):
|
||||
# planes: (N, 3, D', H', W')
|
||||
# points: (N, P, 3)
|
||||
N, P = points.shape[:2]
|
||||
|
||||
# query triplane in chunks
|
||||
outs = []
|
||||
for i in range(0, points.shape[1], chunk_size):
|
||||
chunk_points = points[:, i:i+chunk_size]
|
||||
|
||||
# query triplane
|
||||
chunk_out = self.renderer.run_model_activated(
|
||||
planes=planes,
|
||||
decoder=self.decoder,
|
||||
sample_coordinates=chunk_points,
|
||||
sample_directions=torch.zeros_like(chunk_points),
|
||||
options=self.rendering_kwargs,
|
||||
)
|
||||
outs.append(chunk_out)
|
||||
|
||||
# concatenate the outputs
|
||||
point_features = {
|
||||
k: torch.cat([out[k] for out in outs], dim=1)
|
||||
for k in outs[0].keys()
|
||||
}
|
||||
return point_features
|
||||
@@ -0,0 +1,141 @@
|
||||
# ORIGINAL LICENSE
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
||||
#
|
||||
# Modified by Jiale Xu
|
||||
# The modifications are subject to the same license as the original.
|
||||
|
||||
import itertools
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
|
||||
from .utils.renderer import generate_planes, project_onto_planes, sample_from_planes
|
||||
|
||||
|
||||
class OSGDecoder(nn.Module):
|
||||
"""
|
||||
Triplane decoder that gives RGB and sigma values from sampled features.
|
||||
Using ReLU here instead of Softplus in the original implementation.
|
||||
|
||||
Reference:
|
||||
EG3D: https://github.com/NVlabs/eg3d/blob/main/eg3d/training/triplane.py#L112
|
||||
"""
|
||||
def __init__(self, n_features: int,
|
||||
hidden_dim: int = 64, num_layers: int = 4, activation: nn.Module = nn.ReLU):
|
||||
super().__init__()
|
||||
|
||||
self.net_sdf = nn.Sequential(
|
||||
nn.Linear(3 * n_features, hidden_dim),
|
||||
activation(),
|
||||
*itertools.chain(*[[
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
activation(),
|
||||
] for _ in range(num_layers - 2)]),
|
||||
nn.Linear(hidden_dim, 1),
|
||||
)
|
||||
self.net_rgb = nn.Sequential(
|
||||
nn.Linear(3 * n_features, hidden_dim),
|
||||
activation(),
|
||||
*itertools.chain(*[[
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
activation(),
|
||||
] for _ in range(num_layers - 2)]),
|
||||
nn.Linear(hidden_dim, 3),
|
||||
)
|
||||
self.net_deformation = nn.Sequential(
|
||||
nn.Linear(3 * n_features, hidden_dim),
|
||||
activation(),
|
||||
*itertools.chain(*[[
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
activation(),
|
||||
] for _ in range(num_layers - 2)]),
|
||||
nn.Linear(hidden_dim, 3),
|
||||
)
|
||||
self.net_weight = nn.Sequential(
|
||||
nn.Linear(8 * 3 * n_features, hidden_dim),
|
||||
activation(),
|
||||
*itertools.chain(*[[
|
||||
nn.Linear(hidden_dim, hidden_dim),
|
||||
activation(),
|
||||
] for _ in range(num_layers - 2)]),
|
||||
nn.Linear(hidden_dim, 21),
|
||||
)
|
||||
|
||||
# init all bias to zero
|
||||
for m in self.modules():
|
||||
if isinstance(m, nn.Linear):
|
||||
nn.init.zeros_(m.bias)
|
||||
|
||||
def get_geometry_prediction(self, sampled_features, flexicubes_indices):
|
||||
_N, n_planes, _M, _C = sampled_features.shape
|
||||
sampled_features = sampled_features.permute(0, 2, 1, 3).reshape(_N, _M, n_planes*_C)
|
||||
|
||||
sdf = self.net_sdf(sampled_features)
|
||||
deformation = self.net_deformation(sampled_features)
|
||||
|
||||
grid_features = torch.index_select(input=sampled_features, index=flexicubes_indices.reshape(-1), dim=1)
|
||||
grid_features = grid_features.reshape(
|
||||
sampled_features.shape[0], flexicubes_indices.shape[0], flexicubes_indices.shape[1] * sampled_features.shape[-1])
|
||||
weight = self.net_weight(grid_features) * 0.1
|
||||
|
||||
return sdf, deformation, weight
|
||||
|
||||
def get_texture_prediction(self, sampled_features):
|
||||
_N, n_planes, _M, _C = sampled_features.shape
|
||||
sampled_features = sampled_features.permute(0, 2, 1, 3).reshape(_N, _M, n_planes*_C)
|
||||
|
||||
rgb = self.net_rgb(sampled_features)
|
||||
rgb = torch.sigmoid(rgb)*(1 + 2*0.001) - 0.001 # Uses sigmoid clamping from MipNeRF
|
||||
|
||||
return rgb
|
||||
|
||||
|
||||
class TriplaneSynthesizer(nn.Module):
|
||||
"""
|
||||
Synthesizer that renders a triplane volume with planes and a camera.
|
||||
|
||||
Reference:
|
||||
EG3D: https://github.com/NVlabs/eg3d/blob/main/eg3d/training/triplane.py#L19
|
||||
"""
|
||||
|
||||
DEFAULT_RENDERING_KWARGS = {
|
||||
'ray_start': 'auto',
|
||||
'ray_end': 'auto',
|
||||
'box_warp': 2.,
|
||||
'white_back': True,
|
||||
'disparity_space_sampling': False,
|
||||
'clamp_mode': 'softplus',
|
||||
'sampler_bbox_min': -1.,
|
||||
'sampler_bbox_max': 1.,
|
||||
}
|
||||
|
||||
def __init__(self, triplane_dim: int, samples_per_ray: int):
|
||||
super().__init__()
|
||||
|
||||
# attributes
|
||||
self.triplane_dim = triplane_dim
|
||||
self.rendering_kwargs = {
|
||||
**self.DEFAULT_RENDERING_KWARGS,
|
||||
'depth_resolution': samples_per_ray // 2,
|
||||
'depth_resolution_importance': samples_per_ray // 2,
|
||||
}
|
||||
|
||||
# modules
|
||||
self.plane_axes = generate_planes()
|
||||
self.decoder = OSGDecoder(n_features=triplane_dim)
|
||||
|
||||
def get_geometry_prediction(self, planes, sample_coordinates, flexicubes_indices):
|
||||
plane_axes = self.plane_axes.to(planes.device)
|
||||
sampled_features = sample_from_planes(
|
||||
plane_axes, planes, sample_coordinates, padding_mode='zeros', box_warp=self.rendering_kwargs['box_warp'])
|
||||
|
||||
sdf, deformation, weight = self.decoder.get_geometry_prediction(sampled_features, flexicubes_indices)
|
||||
return sdf, deformation, weight
|
||||
|
||||
def get_texture_prediction(self, planes, sample_coordinates):
|
||||
plane_axes = self.plane_axes.to(planes.device)
|
||||
sampled_features = sample_from_planes(
|
||||
plane_axes, planes, sample_coordinates, padding_mode='zeros', box_warp=self.rendering_kwargs['box_warp'])
|
||||
|
||||
rgb = self.decoder.get_texture_prediction(sampled_features)
|
||||
return rgb
|
||||
@@ -0,0 +1,9 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
||||
#
|
||||
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
|
||||
# property and proprietary rights in and to this material, related
|
||||
# documentation and any modifications thereto. Any use, reproduction,
|
||||
# disclosure or distribution of this material and related documentation
|
||||
# without an express license agreement from NVIDIA CORPORATION or
|
||||
# its affiliates is strictly prohibited.
|
||||
@@ -0,0 +1,118 @@
|
||||
# MIT License
|
||||
|
||||
# Copyright (c) 2022 Petr Kellnhofer
|
||||
|
||||
# Permission is hereby granted, free of charge, to any person obtaining a copy
|
||||
# of this software and associated documentation files (the "Software"), to deal
|
||||
# in the Software without restriction, including without limitation the rights
|
||||
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
||||
# copies of the Software, and to permit persons to whom the Software is
|
||||
# furnished to do so, subject to the following conditions:
|
||||
|
||||
# The above copyright notice and this permission notice shall be included in all
|
||||
# copies or substantial portions of the Software.
|
||||
|
||||
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
||||
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
||||
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
||||
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
||||
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
||||
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
||||
# SOFTWARE.
|
||||
|
||||
import torch
|
||||
|
||||
def transform_vectors(matrix: torch.Tensor, vectors4: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Left-multiplies MxM @ NxM. Returns NxM.
|
||||
"""
|
||||
res = torch.matmul(vectors4, matrix.T)
|
||||
return res
|
||||
|
||||
|
||||
def normalize_vecs(vectors: torch.Tensor) -> torch.Tensor:
|
||||
"""
|
||||
Normalize vector lengths.
|
||||
"""
|
||||
return vectors / (torch.norm(vectors, dim=-1, keepdim=True))
|
||||
|
||||
def torch_dot(x: torch.Tensor, y: torch.Tensor):
|
||||
"""
|
||||
Dot product of two tensors.
|
||||
"""
|
||||
return (x * y).sum(-1)
|
||||
|
||||
|
||||
def get_ray_limits_box(rays_o: torch.Tensor, rays_d: torch.Tensor, box_side_length):
|
||||
"""
|
||||
Author: Petr Kellnhofer
|
||||
Intersects rays with the [-1, 1] NDC volume.
|
||||
Returns min and max distance of entry.
|
||||
Returns -1 for no intersection.
|
||||
https://www.scratchapixel.com/lessons/3d-basic-rendering/minimal-ray-tracer-rendering-simple-shapes/ray-box-intersection
|
||||
"""
|
||||
o_shape = rays_o.shape
|
||||
rays_o = rays_o.detach().reshape(-1, 3)
|
||||
rays_d = rays_d.detach().reshape(-1, 3)
|
||||
|
||||
|
||||
bb_min = [-1*(box_side_length/2), -1*(box_side_length/2), -1*(box_side_length/2)]
|
||||
bb_max = [1*(box_side_length/2), 1*(box_side_length/2), 1*(box_side_length/2)]
|
||||
bounds = torch.tensor([bb_min, bb_max], dtype=rays_o.dtype, device=rays_o.device)
|
||||
is_valid = torch.ones(rays_o.shape[:-1], dtype=bool, device=rays_o.device)
|
||||
|
||||
# Precompute inverse for stability.
|
||||
invdir = 1 / rays_d
|
||||
sign = (invdir < 0).long()
|
||||
|
||||
# Intersect with YZ plane.
|
||||
tmin = (bounds.index_select(0, sign[..., 0])[..., 0] - rays_o[..., 0]) * invdir[..., 0]
|
||||
tmax = (bounds.index_select(0, 1 - sign[..., 0])[..., 0] - rays_o[..., 0]) * invdir[..., 0]
|
||||
|
||||
# Intersect with XZ plane.
|
||||
tymin = (bounds.index_select(0, sign[..., 1])[..., 1] - rays_o[..., 1]) * invdir[..., 1]
|
||||
tymax = (bounds.index_select(0, 1 - sign[..., 1])[..., 1] - rays_o[..., 1]) * invdir[..., 1]
|
||||
|
||||
# Resolve parallel rays.
|
||||
is_valid[torch.logical_or(tmin > tymax, tymin > tmax)] = False
|
||||
|
||||
# Use the shortest intersection.
|
||||
tmin = torch.max(tmin, tymin)
|
||||
tmax = torch.min(tmax, tymax)
|
||||
|
||||
# Intersect with XY plane.
|
||||
tzmin = (bounds.index_select(0, sign[..., 2])[..., 2] - rays_o[..., 2]) * invdir[..., 2]
|
||||
tzmax = (bounds.index_select(0, 1 - sign[..., 2])[..., 2] - rays_o[..., 2]) * invdir[..., 2]
|
||||
|
||||
# Resolve parallel rays.
|
||||
is_valid[torch.logical_or(tmin > tzmax, tzmin > tmax)] = False
|
||||
|
||||
# Use the shortest intersection.
|
||||
tmin = torch.max(tmin, tzmin)
|
||||
tmax = torch.min(tmax, tzmax)
|
||||
|
||||
# Mark invalid.
|
||||
tmin[torch.logical_not(is_valid)] = -1
|
||||
tmax[torch.logical_not(is_valid)] = -2
|
||||
|
||||
return tmin.reshape(*o_shape[:-1], 1), tmax.reshape(*o_shape[:-1], 1)
|
||||
|
||||
|
||||
def linspace(start: torch.Tensor, stop: torch.Tensor, num: int):
|
||||
"""
|
||||
Creates a tensor of shape [num, *start.shape] whose values are evenly spaced from start to end, inclusive.
|
||||
Replicates but the multi-dimensional bahaviour of numpy.linspace in PyTorch.
|
||||
"""
|
||||
# create a tensor of 'num' steps from 0 to 1
|
||||
steps = torch.arange(num, dtype=torch.float32, device=start.device) / (num - 1)
|
||||
|
||||
# reshape the 'steps' tensor to [-1, *([1]*start.ndim)] to allow for broadcastings
|
||||
# - using 'steps.reshape([-1, *([1]*start.ndim)])' would be nice here but torchscript
|
||||
# "cannot statically infer the expected size of a list in this contex", hence the code below
|
||||
for i in range(start.ndim):
|
||||
steps = steps.unsqueeze(-1)
|
||||
|
||||
# the output starts at 'start' and increments until 'stop' in each dimension
|
||||
out = start[None] + steps * (stop - start)[None]
|
||||
|
||||
return out
|
||||
@@ -0,0 +1,72 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
||||
#
|
||||
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
|
||||
# property and proprietary rights in and to this material, related
|
||||
# documentation and any modifications thereto. Any use, reproduction,
|
||||
# disclosure or distribution of this material and related documentation
|
||||
# without an express license agreement from NVIDIA CORPORATION or
|
||||
# its affiliates is strictly prohibited.
|
||||
#
|
||||
# Modified by Jiale Xu
|
||||
# The modifications are subject to the same license as the original.
|
||||
|
||||
|
||||
"""
|
||||
The ray marcher takes the raw output of the implicit representation and uses the volume rendering equation to produce composited colors and depths.
|
||||
Based off of the implementation in MipNeRF (this one doesn't do any cone tracing though!)
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
|
||||
class MipRayMarcher2(nn.Module):
|
||||
def __init__(self, activation_factory):
|
||||
super().__init__()
|
||||
self.activation_factory = activation_factory
|
||||
|
||||
def run_forward(self, colors, densities, depths, rendering_options, normals=None):
|
||||
dtype = colors.dtype
|
||||
deltas = depths[:, :, 1:] - depths[:, :, :-1]
|
||||
colors_mid = (colors[:, :, :-1] + colors[:, :, 1:]) / 2
|
||||
densities_mid = (densities[:, :, :-1] + densities[:, :, 1:]) / 2
|
||||
depths_mid = (depths[:, :, :-1] + depths[:, :, 1:]) / 2
|
||||
|
||||
# using factory mode for better usability
|
||||
densities_mid = self.activation_factory(rendering_options)(densities_mid).to(dtype)
|
||||
|
||||
density_delta = densities_mid * deltas
|
||||
|
||||
alpha = 1 - torch.exp(-density_delta).to(dtype)
|
||||
|
||||
alpha_shifted = torch.cat([torch.ones_like(alpha[:, :, :1]), 1-alpha + 1e-10], -2)
|
||||
weights = alpha * torch.cumprod(alpha_shifted, -2)[:, :, :-1]
|
||||
weights = weights.to(dtype)
|
||||
|
||||
composite_rgb = torch.sum(weights * colors_mid, -2)
|
||||
weight_total = weights.sum(2)
|
||||
# composite_depth = torch.sum(weights * depths_mid, -2) / weight_total
|
||||
composite_depth = torch.sum(weights * depths_mid, -2)
|
||||
|
||||
# clip the composite to min/max range of depths
|
||||
composite_depth = torch.nan_to_num(composite_depth, float('inf')).to(dtype)
|
||||
composite_depth = torch.clamp(composite_depth, torch.min(depths), torch.max(depths))
|
||||
|
||||
if rendering_options.get('white_back', False):
|
||||
composite_rgb = composite_rgb + 1 - weight_total
|
||||
|
||||
# rendered value scale is 0-1, comment out original mipnerf scaling
|
||||
# composite_rgb = composite_rgb * 2 - 1 # Scale to (-1, 1)
|
||||
|
||||
return composite_rgb, composite_depth, weights
|
||||
|
||||
|
||||
def forward(self, colors, densities, depths, rendering_options, normals=None):
|
||||
if normals is not None:
|
||||
composite_rgb, composite_depth, composite_normals, weights = self.run_forward(colors, densities, depths, rendering_options, normals)
|
||||
return composite_rgb, composite_depth, composite_normals, weights
|
||||
|
||||
composite_rgb, composite_depth, weights = self.run_forward(colors, densities, depths, rendering_options)
|
||||
return composite_rgb, composite_depth, weights
|
||||
@@ -0,0 +1,141 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
||||
#
|
||||
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
|
||||
# property and proprietary rights in and to this material, related
|
||||
# documentation and any modifications thereto. Any use, reproduction,
|
||||
# disclosure or distribution of this material and related documentation
|
||||
# without an express license agreement from NVIDIA CORPORATION or
|
||||
# its affiliates is strictly prohibited.
|
||||
#
|
||||
# Modified by Jiale Xu
|
||||
# The modifications are subject to the same license as the original.
|
||||
|
||||
|
||||
"""
|
||||
The ray sampler is a module that takes in camera matrices and resolution and batches of rays.
|
||||
Expects cam2world matrices that use the OpenCV camera coordinate system conventions.
|
||||
"""
|
||||
|
||||
import torch
|
||||
|
||||
class RaySampler(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.ray_origins_h, self.ray_directions, self.depths, self.image_coords, self.rendering_options = None, None, None, None, None
|
||||
|
||||
|
||||
def forward(self, cam2world_matrix, intrinsics, render_size):
|
||||
"""
|
||||
Create batches of rays and return origins and directions.
|
||||
|
||||
cam2world_matrix: (N, 4, 4)
|
||||
intrinsics: (N, 3, 3)
|
||||
render_size: int
|
||||
|
||||
ray_origins: (N, M, 3)
|
||||
ray_dirs: (N, M, 2)
|
||||
"""
|
||||
|
||||
dtype = cam2world_matrix.dtype
|
||||
device = cam2world_matrix.device
|
||||
N, M = cam2world_matrix.shape[0], render_size**2
|
||||
cam_locs_world = cam2world_matrix[:, :3, 3]
|
||||
fx = intrinsics[:, 0, 0]
|
||||
fy = intrinsics[:, 1, 1]
|
||||
cx = intrinsics[:, 0, 2]
|
||||
cy = intrinsics[:, 1, 2]
|
||||
sk = intrinsics[:, 0, 1]
|
||||
|
||||
uv = torch.stack(torch.meshgrid(
|
||||
torch.arange(render_size, dtype=dtype, device=device),
|
||||
torch.arange(render_size, dtype=dtype, device=device),
|
||||
indexing='ij',
|
||||
))
|
||||
uv = uv.flip(0).reshape(2, -1).transpose(1, 0)
|
||||
uv = uv.unsqueeze(0).repeat(cam2world_matrix.shape[0], 1, 1)
|
||||
|
||||
x_cam = uv[:, :, 0].view(N, -1) * (1./render_size) + (0.5/render_size)
|
||||
y_cam = uv[:, :, 1].view(N, -1) * (1./render_size) + (0.5/render_size)
|
||||
z_cam = torch.ones((N, M), dtype=dtype, device=device)
|
||||
|
||||
x_lift = (x_cam - cx.unsqueeze(-1) + cy.unsqueeze(-1)*sk.unsqueeze(-1)/fy.unsqueeze(-1) - sk.unsqueeze(-1)*y_cam/fy.unsqueeze(-1)) / fx.unsqueeze(-1) * z_cam
|
||||
y_lift = (y_cam - cy.unsqueeze(-1)) / fy.unsqueeze(-1) * z_cam
|
||||
|
||||
cam_rel_points = torch.stack((x_lift, y_lift, z_cam, torch.ones_like(z_cam)), dim=-1).to(dtype)
|
||||
|
||||
_opencv2blender = torch.tensor([
|
||||
[1, 0, 0, 0],
|
||||
[0, -1, 0, 0],
|
||||
[0, 0, -1, 0],
|
||||
[0, 0, 0, 1],
|
||||
], dtype=dtype, device=device).unsqueeze(0).repeat(N, 1, 1)
|
||||
|
||||
cam2world_matrix = torch.bmm(cam2world_matrix, _opencv2blender)
|
||||
|
||||
world_rel_points = torch.bmm(cam2world_matrix, cam_rel_points.permute(0, 2, 1)).permute(0, 2, 1)[:, :, :3]
|
||||
|
||||
ray_dirs = world_rel_points - cam_locs_world[:, None, :]
|
||||
ray_dirs = torch.nn.functional.normalize(ray_dirs, dim=2).to(dtype)
|
||||
|
||||
ray_origins = cam_locs_world.unsqueeze(1).repeat(1, ray_dirs.shape[1], 1)
|
||||
|
||||
return ray_origins, ray_dirs
|
||||
|
||||
|
||||
class OrthoRaySampler(torch.nn.Module):
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.ray_origins_h, self.ray_directions, self.depths, self.image_coords, self.rendering_options = None, None, None, None, None
|
||||
|
||||
|
||||
def forward(self, cam2world_matrix, ortho_scale, render_size):
|
||||
"""
|
||||
Create batches of rays and return origins and directions.
|
||||
|
||||
cam2world_matrix: (N, 4, 4)
|
||||
ortho_scale: float
|
||||
render_size: int
|
||||
|
||||
ray_origins: (N, M, 3)
|
||||
ray_dirs: (N, M, 3)
|
||||
"""
|
||||
|
||||
N, M = cam2world_matrix.shape[0], render_size**2
|
||||
|
||||
uv = torch.stack(torch.meshgrid(
|
||||
torch.arange(render_size, dtype=torch.float32, device=cam2world_matrix.device),
|
||||
torch.arange(render_size, dtype=torch.float32, device=cam2world_matrix.device),
|
||||
indexing='ij',
|
||||
))
|
||||
uv = uv.flip(0).reshape(2, -1).transpose(1, 0)
|
||||
uv = uv.unsqueeze(0).repeat(cam2world_matrix.shape[0], 1, 1)
|
||||
|
||||
x_cam = uv[:, :, 0].view(N, -1) * (1./render_size) + (0.5/render_size)
|
||||
y_cam = uv[:, :, 1].view(N, -1) * (1./render_size) + (0.5/render_size)
|
||||
z_cam = torch.zeros((N, M), device=cam2world_matrix.device)
|
||||
|
||||
x_lift = (x_cam - 0.5) * ortho_scale
|
||||
y_lift = (y_cam - 0.5) * ortho_scale
|
||||
|
||||
cam_rel_points = torch.stack((x_lift, y_lift, z_cam, torch.ones_like(z_cam)), dim=-1)
|
||||
|
||||
_opencv2blender = torch.tensor([
|
||||
[1, 0, 0, 0],
|
||||
[0, -1, 0, 0],
|
||||
[0, 0, -1, 0],
|
||||
[0, 0, 0, 1],
|
||||
], dtype=torch.float32, device=cam2world_matrix.device).unsqueeze(0).repeat(N, 1, 1)
|
||||
|
||||
cam2world_matrix = torch.bmm(cam2world_matrix, _opencv2blender)
|
||||
|
||||
ray_origins = torch.bmm(cam2world_matrix, cam_rel_points.permute(0, 2, 1)).permute(0, 2, 1)[:, :, :3]
|
||||
|
||||
ray_dirs_cam = torch.stack([
|
||||
torch.zeros((N, M), device=cam2world_matrix.device),
|
||||
torch.zeros((N, M), device=cam2world_matrix.device),
|
||||
torch.ones((N, M), device=cam2world_matrix.device),
|
||||
], dim=-1)
|
||||
ray_dirs = torch.bmm(cam2world_matrix[:, :3, :3], ray_dirs_cam.permute(0, 2, 1)).permute(0, 2, 1)
|
||||
|
||||
return ray_origins, ray_dirs
|
||||
@@ -0,0 +1,323 @@
|
||||
# SPDX-FileCopyrightText: Copyright (c) 2021-2022 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
# SPDX-License-Identifier: LicenseRef-NvidiaProprietary
|
||||
#
|
||||
# NVIDIA CORPORATION, its affiliates and licensors retain all intellectual
|
||||
# property and proprietary rights in and to this material, related
|
||||
# documentation and any modifications thereto. Any use, reproduction,
|
||||
# disclosure or distribution of this material and related documentation
|
||||
# without an express license agreement from NVIDIA CORPORATION or
|
||||
# its affiliates is strictly prohibited.
|
||||
#
|
||||
# Modified by Jiale Xu
|
||||
# The modifications are subject to the same license as the original.
|
||||
|
||||
|
||||
"""
|
||||
The renderer is a module that takes in rays, decides where to sample along each
|
||||
ray, and computes pixel colors using the volume rendering equation.
|
||||
"""
|
||||
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
|
||||
from .ray_marcher import MipRayMarcher2
|
||||
from . import math_utils
|
||||
|
||||
|
||||
def generate_planes():
|
||||
"""
|
||||
Defines planes by the three vectors that form the "axes" of the
|
||||
plane. Should work with arbitrary number of planes and planes of
|
||||
arbitrary orientation.
|
||||
|
||||
Bugfix reference: https://github.com/NVlabs/eg3d/issues/67
|
||||
"""
|
||||
return torch.tensor([[[1, 0, 0],
|
||||
[0, 1, 0],
|
||||
[0, 0, 1]],
|
||||
[[1, 0, 0],
|
||||
[0, 0, 1],
|
||||
[0, 1, 0]],
|
||||
[[0, 0, 1],
|
||||
[0, 1, 0],
|
||||
[1, 0, 0]]], dtype=torch.float32)
|
||||
|
||||
def project_onto_planes(planes, coordinates):
|
||||
"""
|
||||
Does a projection of a 3D point onto a batch of 2D planes,
|
||||
returning 2D plane coordinates.
|
||||
|
||||
Takes plane axes of shape n_planes, 3, 3
|
||||
# Takes coordinates of shape N, M, 3
|
||||
# returns projections of shape N*n_planes, M, 2
|
||||
"""
|
||||
N, M, C = coordinates.shape
|
||||
n_planes, _, _ = planes.shape
|
||||
coordinates = coordinates.unsqueeze(1).expand(-1, n_planes, -1, -1).reshape(N*n_planes, M, 3)
|
||||
inv_planes = torch.linalg.inv(planes).unsqueeze(0).expand(N, -1, -1, -1).reshape(N*n_planes, 3, 3)
|
||||
projections = torch.bmm(coordinates, inv_planes)
|
||||
return projections[..., :2]
|
||||
|
||||
def sample_from_planes(plane_axes, plane_features, coordinates, mode='bilinear', padding_mode='zeros', box_warp=None):
|
||||
assert padding_mode == 'zeros'
|
||||
N, n_planes, C, H, W = plane_features.shape
|
||||
_, M, _ = coordinates.shape
|
||||
plane_features = plane_features.view(N*n_planes, C, H, W)
|
||||
dtype = plane_features.dtype
|
||||
|
||||
coordinates = (2/box_warp) * coordinates # add specific box bounds
|
||||
|
||||
projected_coordinates = project_onto_planes(plane_axes, coordinates).unsqueeze(1)
|
||||
output_features = torch.nn.functional.grid_sample(
|
||||
plane_features,
|
||||
projected_coordinates.to(dtype),
|
||||
mode=mode,
|
||||
padding_mode=padding_mode,
|
||||
align_corners=False,
|
||||
).permute(0, 3, 2, 1).reshape(N, n_planes, M, C)
|
||||
return output_features
|
||||
|
||||
def sample_from_3dgrid(grid, coordinates):
|
||||
"""
|
||||
Expects coordinates in shape (batch_size, num_points_per_batch, 3)
|
||||
Expects grid in shape (1, channels, H, W, D)
|
||||
(Also works if grid has batch size)
|
||||
Returns sampled features of shape (batch_size, num_points_per_batch, feature_channels)
|
||||
"""
|
||||
batch_size, n_coords, n_dims = coordinates.shape
|
||||
sampled_features = torch.nn.functional.grid_sample(
|
||||
grid.expand(batch_size, -1, -1, -1, -1),
|
||||
coordinates.reshape(batch_size, 1, 1, -1, n_dims),
|
||||
mode='bilinear',
|
||||
padding_mode='zeros',
|
||||
align_corners=False,
|
||||
)
|
||||
N, C, H, W, D = sampled_features.shape
|
||||
sampled_features = sampled_features.permute(0, 4, 3, 2, 1).reshape(N, H*W*D, C)
|
||||
return sampled_features
|
||||
|
||||
class ImportanceRenderer(torch.nn.Module):
|
||||
"""
|
||||
Modified original version to filter out-of-box samples as TensoRF does.
|
||||
|
||||
Reference:
|
||||
TensoRF: https://github.com/apchenstu/TensoRF/blob/main/models/tensorBase.py#L277
|
||||
"""
|
||||
def __init__(self):
|
||||
super().__init__()
|
||||
self.activation_factory = self._build_activation_factory()
|
||||
self.ray_marcher = MipRayMarcher2(self.activation_factory)
|
||||
self.plane_axes = generate_planes()
|
||||
|
||||
def _build_activation_factory(self):
|
||||
def activation_factory(options: dict):
|
||||
if options['clamp_mode'] == 'softplus':
|
||||
return lambda x: F.softplus(x - 1) # activation bias of -1 makes things initialize better
|
||||
else:
|
||||
assert False, "Renderer only supports `clamp_mode`=`softplus`!"
|
||||
return activation_factory
|
||||
|
||||
def _forward_pass(self, depths: torch.Tensor, ray_directions: torch.Tensor, ray_origins: torch.Tensor,
|
||||
planes: torch.Tensor, decoder: nn.Module, rendering_options: dict):
|
||||
"""
|
||||
Additional filtering is applied to filter out-of-box samples.
|
||||
Modifications made by Zexin He.
|
||||
"""
|
||||
|
||||
# context related variables
|
||||
batch_size, num_rays, samples_per_ray, _ = depths.shape
|
||||
device = depths.device
|
||||
|
||||
# define sample points with depths
|
||||
sample_directions = ray_directions.unsqueeze(-2).expand(-1, -1, samples_per_ray, -1).reshape(batch_size, -1, 3)
|
||||
sample_coordinates = (ray_origins.unsqueeze(-2) + depths * ray_directions.unsqueeze(-2)).reshape(batch_size, -1, 3)
|
||||
|
||||
# filter out-of-box samples
|
||||
mask_inbox = \
|
||||
(rendering_options['sampler_bbox_min'] <= sample_coordinates) & \
|
||||
(sample_coordinates <= rendering_options['sampler_bbox_max'])
|
||||
mask_inbox = mask_inbox.all(-1)
|
||||
|
||||
# forward model according to all samples
|
||||
_out = self.run_model(planes, decoder, sample_coordinates, sample_directions, rendering_options)
|
||||
|
||||
# set out-of-box samples to zeros(rgb) & -inf(sigma)
|
||||
SAFE_GUARD = 3
|
||||
DATA_TYPE = _out['sigma'].dtype
|
||||
colors_pass = torch.zeros(batch_size, num_rays * samples_per_ray, 3, device=device, dtype=DATA_TYPE)
|
||||
densities_pass = torch.nan_to_num(torch.full((batch_size, num_rays * samples_per_ray, 1), -float('inf'), device=device, dtype=DATA_TYPE)) / SAFE_GUARD
|
||||
colors_pass[mask_inbox], densities_pass[mask_inbox] = _out['rgb'][mask_inbox], _out['sigma'][mask_inbox]
|
||||
|
||||
# reshape back
|
||||
colors_pass = colors_pass.reshape(batch_size, num_rays, samples_per_ray, colors_pass.shape[-1])
|
||||
densities_pass = densities_pass.reshape(batch_size, num_rays, samples_per_ray, densities_pass.shape[-1])
|
||||
|
||||
return colors_pass, densities_pass
|
||||
|
||||
def forward(self, planes, decoder, ray_origins, ray_directions, rendering_options):
|
||||
# self.plane_axes = self.plane_axes.to(ray_origins.device)
|
||||
|
||||
if rendering_options['ray_start'] == rendering_options['ray_end'] == 'auto':
|
||||
ray_start, ray_end = math_utils.get_ray_limits_box(ray_origins, ray_directions, box_side_length=rendering_options['box_warp'])
|
||||
is_ray_valid = ray_end > ray_start
|
||||
if torch.any(is_ray_valid).item():
|
||||
ray_start[~is_ray_valid] = ray_start[is_ray_valid].min()
|
||||
ray_end[~is_ray_valid] = ray_start[is_ray_valid].max()
|
||||
depths_coarse = self.sample_stratified(ray_origins, ray_start, ray_end, rendering_options['depth_resolution'], rendering_options['disparity_space_sampling'])
|
||||
else:
|
||||
# Create stratified depth samples
|
||||
depths_coarse = self.sample_stratified(ray_origins, rendering_options['ray_start'], rendering_options['ray_end'], rendering_options['depth_resolution'], rendering_options['disparity_space_sampling'])
|
||||
|
||||
# Coarse Pass
|
||||
colors_coarse, densities_coarse = self._forward_pass(
|
||||
depths=depths_coarse, ray_directions=ray_directions, ray_origins=ray_origins,
|
||||
planes=planes, decoder=decoder, rendering_options=rendering_options)
|
||||
|
||||
# Fine Pass
|
||||
N_importance = rendering_options['depth_resolution_importance']
|
||||
if N_importance > 0:
|
||||
_, _, weights = self.ray_marcher(colors_coarse, densities_coarse, depths_coarse, rendering_options)
|
||||
|
||||
depths_fine = self.sample_importance(depths_coarse, weights, N_importance)
|
||||
|
||||
colors_fine, densities_fine = self._forward_pass(
|
||||
depths=depths_fine, ray_directions=ray_directions, ray_origins=ray_origins,
|
||||
planes=planes, decoder=decoder, rendering_options=rendering_options)
|
||||
|
||||
all_depths, all_colors, all_densities = self.unify_samples(depths_coarse, colors_coarse, densities_coarse,
|
||||
depths_fine, colors_fine, densities_fine)
|
||||
|
||||
rgb_final, depth_final, weights = self.ray_marcher(all_colors, all_densities, all_depths, rendering_options)
|
||||
else:
|
||||
rgb_final, depth_final, weights = self.ray_marcher(colors_coarse, densities_coarse, depths_coarse, rendering_options)
|
||||
|
||||
return rgb_final, depth_final, weights.sum(2)
|
||||
|
||||
def run_model(self, planes, decoder, sample_coordinates, sample_directions, options):
|
||||
plane_axes = self.plane_axes.to(planes.device)
|
||||
sampled_features = sample_from_planes(plane_axes, planes, sample_coordinates, padding_mode='zeros', box_warp=options['box_warp'])
|
||||
|
||||
out = decoder(sampled_features, sample_directions)
|
||||
if options.get('density_noise', 0) > 0:
|
||||
out['sigma'] += torch.randn_like(out['sigma']) * options['density_noise']
|
||||
return out
|
||||
|
||||
def run_model_activated(self, planes, decoder, sample_coordinates, sample_directions, options):
|
||||
out = self.run_model(planes, decoder, sample_coordinates, sample_directions, options)
|
||||
out['sigma'] = self.activation_factory(options)(out['sigma'])
|
||||
return out
|
||||
|
||||
def sort_samples(self, all_depths, all_colors, all_densities):
|
||||
_, indices = torch.sort(all_depths, dim=-2)
|
||||
all_depths = torch.gather(all_depths, -2, indices)
|
||||
all_colors = torch.gather(all_colors, -2, indices.expand(-1, -1, -1, all_colors.shape[-1]))
|
||||
all_densities = torch.gather(all_densities, -2, indices.expand(-1, -1, -1, 1))
|
||||
return all_depths, all_colors, all_densities
|
||||
|
||||
def unify_samples(self, depths1, colors1, densities1, depths2, colors2, densities2, normals1=None, normals2=None):
|
||||
all_depths = torch.cat([depths1, depths2], dim = -2)
|
||||
all_colors = torch.cat([colors1, colors2], dim = -2)
|
||||
all_densities = torch.cat([densities1, densities2], dim = -2)
|
||||
|
||||
if normals1 is not None and normals2 is not None:
|
||||
all_normals = torch.cat([normals1, normals2], dim = -2)
|
||||
else:
|
||||
all_normals = None
|
||||
|
||||
_, indices = torch.sort(all_depths, dim=-2)
|
||||
all_depths = torch.gather(all_depths, -2, indices)
|
||||
all_colors = torch.gather(all_colors, -2, indices.expand(-1, -1, -1, all_colors.shape[-1]))
|
||||
all_densities = torch.gather(all_densities, -2, indices.expand(-1, -1, -1, 1))
|
||||
|
||||
if all_normals is not None:
|
||||
all_normals = torch.gather(all_normals, -2, indices.expand(-1, -1, -1, all_normals.shape[-1]))
|
||||
return all_depths, all_colors, all_normals, all_densities
|
||||
|
||||
return all_depths, all_colors, all_densities
|
||||
|
||||
def sample_stratified(self, ray_origins, ray_start, ray_end, depth_resolution, disparity_space_sampling=False):
|
||||
"""
|
||||
Return depths of approximately uniformly spaced samples along rays.
|
||||
"""
|
||||
N, M, _ = ray_origins.shape
|
||||
if disparity_space_sampling:
|
||||
depths_coarse = torch.linspace(0,
|
||||
1,
|
||||
depth_resolution,
|
||||
device=ray_origins.device).reshape(1, 1, depth_resolution, 1).repeat(N, M, 1, 1)
|
||||
depth_delta = 1/(depth_resolution - 1)
|
||||
depths_coarse += torch.rand_like(depths_coarse) * depth_delta
|
||||
depths_coarse = 1./(1./ray_start * (1. - depths_coarse) + 1./ray_end * depths_coarse)
|
||||
else:
|
||||
if type(ray_start) == torch.Tensor:
|
||||
depths_coarse = math_utils.linspace(ray_start, ray_end, depth_resolution).permute(1,2,0,3)
|
||||
depth_delta = (ray_end - ray_start) / (depth_resolution - 1)
|
||||
depths_coarse += torch.rand_like(depths_coarse) * depth_delta[..., None]
|
||||
else:
|
||||
depths_coarse = torch.linspace(ray_start, ray_end, depth_resolution, device=ray_origins.device).reshape(1, 1, depth_resolution, 1).repeat(N, M, 1, 1)
|
||||
depth_delta = (ray_end - ray_start)/(depth_resolution - 1)
|
||||
depths_coarse += torch.rand_like(depths_coarse) * depth_delta
|
||||
|
||||
return depths_coarse
|
||||
|
||||
def sample_importance(self, z_vals, weights, N_importance):
|
||||
"""
|
||||
Return depths of importance sampled points along rays. See NeRF importance sampling for more.
|
||||
"""
|
||||
with torch.no_grad():
|
||||
batch_size, num_rays, samples_per_ray, _ = z_vals.shape
|
||||
|
||||
z_vals = z_vals.reshape(batch_size * num_rays, samples_per_ray)
|
||||
weights = weights.reshape(batch_size * num_rays, -1) # -1 to account for loss of 1 sample in MipRayMarcher
|
||||
|
||||
# smooth weights
|
||||
weights = torch.nn.functional.max_pool1d(weights.unsqueeze(1), 2, 1, padding=1)
|
||||
weights = torch.nn.functional.avg_pool1d(weights, 2, 1).squeeze()
|
||||
weights = weights + 0.01
|
||||
|
||||
z_vals_mid = 0.5 * (z_vals[: ,:-1] + z_vals[: ,1:])
|
||||
importance_z_vals = self.sample_pdf(z_vals_mid, weights[:, 1:-1],
|
||||
N_importance).detach().reshape(batch_size, num_rays, N_importance, 1)
|
||||
return importance_z_vals
|
||||
|
||||
def sample_pdf(self, bins, weights, N_importance, det=False, eps=1e-5):
|
||||
"""
|
||||
Sample @N_importance samples from @bins with distribution defined by @weights.
|
||||
Inputs:
|
||||
bins: (N_rays, N_samples_+1) where N_samples_ is "the number of coarse samples per ray - 2"
|
||||
weights: (N_rays, N_samples_)
|
||||
N_importance: the number of samples to draw from the distribution
|
||||
det: deterministic or not
|
||||
eps: a small number to prevent division by zero
|
||||
Outputs:
|
||||
samples: the sampled samples
|
||||
"""
|
||||
N_rays, N_samples_ = weights.shape
|
||||
weights = weights + eps # prevent division by zero (don't do inplace op!)
|
||||
pdf = weights / torch.sum(weights, -1, keepdim=True) # (N_rays, N_samples_)
|
||||
cdf = torch.cumsum(pdf, -1) # (N_rays, N_samples), cumulative distribution function
|
||||
cdf = torch.cat([torch.zeros_like(cdf[: ,:1]), cdf], -1) # (N_rays, N_samples_+1)
|
||||
# padded to 0~1 inclusive
|
||||
|
||||
if det:
|
||||
u = torch.linspace(0, 1, N_importance, device=bins.device)
|
||||
u = u.expand(N_rays, N_importance)
|
||||
else:
|
||||
u = torch.rand(N_rays, N_importance, device=bins.device)
|
||||
u = u.contiguous()
|
||||
|
||||
inds = torch.searchsorted(cdf, u, right=True)
|
||||
below = torch.clamp_min(inds-1, 0)
|
||||
above = torch.clamp_max(inds, N_samples_)
|
||||
|
||||
inds_sampled = torch.stack([below, above], -1).view(N_rays, 2*N_importance)
|
||||
cdf_g = torch.gather(cdf, 1, inds_sampled).view(N_rays, N_importance, 2)
|
||||
bins_g = torch.gather(bins, 1, inds_sampled).view(N_rays, N_importance, 2)
|
||||
|
||||
denom = cdf_g[...,1]-cdf_g[...,0]
|
||||
denom[denom<eps] = 1 # denom equals 0 means a bin has weight 0, in which case it will not be sampled
|
||||
# anyway, therefore any value for it is fine (set to 1 here)
|
||||
|
||||
samples = bins_g[...,0] + (u-cdf_g[...,0])/denom * (bins_g[...,1]-bins_g[...,0])
|
||||
return samples
|
||||
@@ -0,0 +1,116 @@
|
||||
import torch
|
||||
import torch.nn.functional as F
|
||||
import numpy as np
|
||||
|
||||
|
||||
def pad_camera_extrinsics_4x4(extrinsics):
|
||||
if extrinsics.shape[-2] == 4:
|
||||
return extrinsics
|
||||
padding = torch.tensor([[0, 0, 0, 1]]).to(extrinsics)
|
||||
if extrinsics.ndim == 3:
|
||||
padding = padding.unsqueeze(0).repeat(extrinsics.shape[0], 1, 1)
|
||||
extrinsics = torch.cat([extrinsics, padding], dim=-2)
|
||||
return extrinsics
|
||||
|
||||
|
||||
def center_looking_at_camera_pose(camera_position: torch.Tensor, look_at: torch.Tensor = None, up_world: torch.Tensor = None):
|
||||
"""
|
||||
Create OpenGL camera extrinsics from camera locations and look-at position.
|
||||
|
||||
camera_position: (M, 3) or (3,)
|
||||
look_at: (3)
|
||||
up_world: (3)
|
||||
return: (M, 3, 4) or (3, 4)
|
||||
"""
|
||||
# by default, looking at the origin and world up is z-axis
|
||||
if look_at is None:
|
||||
look_at = torch.tensor([0, 0, 0], dtype=torch.float32)
|
||||
if up_world is None:
|
||||
up_world = torch.tensor([0, 0, 1], dtype=torch.float32)
|
||||
if camera_position.ndim == 2:
|
||||
look_at = look_at.unsqueeze(0).repeat(camera_position.shape[0], 1)
|
||||
up_world = up_world.unsqueeze(0).repeat(camera_position.shape[0], 1)
|
||||
|
||||
# OpenGL camera: z-backward, x-right, y-up
|
||||
z_axis = camera_position - look_at
|
||||
z_axis = F.normalize(z_axis, dim=-1).float()
|
||||
x_axis = torch.linalg.cross(up_world, z_axis, dim=-1)
|
||||
x_axis = F.normalize(x_axis, dim=-1).float()
|
||||
y_axis = torch.linalg.cross(z_axis, x_axis, dim=-1)
|
||||
y_axis = F.normalize(y_axis, dim=-1).float()
|
||||
|
||||
extrinsics = torch.stack([x_axis, y_axis, z_axis, camera_position], dim=-1)
|
||||
extrinsics = pad_camera_extrinsics_4x4(extrinsics)
|
||||
return extrinsics
|
||||
|
||||
|
||||
def spherical_camera_pose(azimuths: np.ndarray, elevations: np.ndarray, radius=2.5):
|
||||
azimuths = np.deg2rad(azimuths)
|
||||
elevations = np.deg2rad(elevations)
|
||||
|
||||
xs = radius * np.cos(elevations) * np.cos(azimuths)
|
||||
ys = radius * np.cos(elevations) * np.sin(azimuths)
|
||||
zs = radius * np.sin(elevations)
|
||||
|
||||
cam_locations = np.stack([xs, ys, zs], axis=-1)
|
||||
cam_locations = torch.from_numpy(cam_locations).float()
|
||||
|
||||
c2ws = center_looking_at_camera_pose(cam_locations)
|
||||
return c2ws
|
||||
|
||||
|
||||
def get_circular_camera_poses(M=120, radius=2.5, elevation=30.0):
|
||||
# M: number of circular views
|
||||
# radius: camera dist to center
|
||||
# elevation: elevation degrees of the camera
|
||||
# return: (M, 4, 4)
|
||||
assert M > 0 and radius > 0
|
||||
|
||||
elevation = np.deg2rad(elevation)
|
||||
|
||||
camera_positions = []
|
||||
for i in range(M):
|
||||
azimuth = 2 * np.pi * i / M
|
||||
x = radius * np.cos(elevation) * np.cos(azimuth)
|
||||
y = radius * np.cos(elevation) * np.sin(azimuth)
|
||||
z = radius * np.sin(elevation)
|
||||
camera_positions.append([x, y, z])
|
||||
camera_positions = np.array(camera_positions)
|
||||
camera_positions = torch.from_numpy(camera_positions).float()
|
||||
extrinsics = center_looking_at_camera_pose(camera_positions)
|
||||
return extrinsics
|
||||
|
||||
def FOV_to_intrinsics(fov, device='cpu'):
|
||||
"""
|
||||
Creates a 3x3 camera intrinsics matrix from the camera field of view, specified in degrees.
|
||||
Note the intrinsics are returned as normalized by image size, rather than in pixel units.
|
||||
Assumes principal point is at image center.
|
||||
"""
|
||||
focal_length = 0.5 / np.tan(np.deg2rad(fov) * 0.5)
|
||||
intrinsics = torch.tensor([[focal_length, 0, 0.5], [0, focal_length, 0.5], [0, 0, 1]], device=device)
|
||||
return intrinsics
|
||||
|
||||
def oribt_camera_poses_to_input_cameras(azimuths, elevations, batch_size=1, radius=4.0, fov=30.0):
|
||||
"""
|
||||
Args:
|
||||
azimuths (list)
|
||||
elevations (list)
|
||||
radius (float or list)
|
||||
"""
|
||||
# (0, 180/-180, 0) -> (0, 180, 360)
|
||||
azimuths = [360 + angle if angle < 0 else angle for angle in azimuths]
|
||||
azimuths = np.array(azimuths).astype(float)
|
||||
elevations = -np.array(elevations).astype(float)
|
||||
if isinstance(radius, list):
|
||||
radius = np.array(radius).astype(float)
|
||||
|
||||
c2ws = spherical_camera_pose(azimuths, elevations, radius)
|
||||
c2ws = c2ws.float().flatten(-2)
|
||||
|
||||
Ks = FOV_to_intrinsics(fov).unsqueeze(0).repeat(azimuths.shape[0], 1, 1).float().flatten(-2)
|
||||
|
||||
extrinsics = c2ws[:, :12]
|
||||
intrinsics = torch.stack([Ks[:, 0], Ks[:, 4], Ks[:, 2], Ks[:, 5]], dim=-1)
|
||||
cameras = torch.cat([extrinsics, intrinsics], dim=-1)
|
||||
|
||||
return cameras.unsqueeze(0).repeat(batch_size, 1, 1)
|
||||
@@ -0,0 +1,97 @@
|
||||
import os
|
||||
import imageio
|
||||
import rembg
|
||||
import torch
|
||||
import numpy as np
|
||||
import PIL.Image
|
||||
from PIL import Image
|
||||
from typing import Any
|
||||
|
||||
|
||||
def remove_background(image: PIL.Image.Image,
|
||||
rembg_session: Any = None,
|
||||
force: bool = False,
|
||||
**rembg_kwargs,
|
||||
) -> PIL.Image.Image:
|
||||
do_remove = True
|
||||
if image.mode == "RGBA" and image.getextrema()[3][0] < 255:
|
||||
do_remove = False
|
||||
do_remove = do_remove or force
|
||||
if do_remove:
|
||||
image = rembg.remove(image, session=rembg_session, **rembg_kwargs)
|
||||
return image
|
||||
|
||||
|
||||
def resize_foreground(
|
||||
image: PIL.Image.Image,
|
||||
ratio: float,
|
||||
) -> PIL.Image.Image:
|
||||
image = np.array(image)
|
||||
assert image.shape[-1] == 4
|
||||
alpha = np.where(image[..., 3] > 0)
|
||||
y1, y2, x1, x2 = (
|
||||
alpha[0].min(),
|
||||
alpha[0].max(),
|
||||
alpha[1].min(),
|
||||
alpha[1].max(),
|
||||
)
|
||||
# crop the foreground
|
||||
fg = image[y1:y2, x1:x2]
|
||||
# pad to square
|
||||
size = max(fg.shape[0], fg.shape[1])
|
||||
ph0, pw0 = (size - fg.shape[0]) // 2, (size - fg.shape[1]) // 2
|
||||
ph1, pw1 = size - fg.shape[0] - ph0, size - fg.shape[1] - pw0
|
||||
new_image = np.pad(
|
||||
fg,
|
||||
((ph0, ph1), (pw0, pw1), (0, 0)),
|
||||
mode="constant",
|
||||
constant_values=((0, 0), (0, 0), (0, 0)),
|
||||
)
|
||||
|
||||
# compute padding according to the ratio
|
||||
new_size = int(new_image.shape[0] / ratio)
|
||||
# pad to size, double side
|
||||
ph0, pw0 = (new_size - size) // 2, (new_size - size) // 2
|
||||
ph1, pw1 = new_size - size - ph0, new_size - size - pw0
|
||||
new_image = np.pad(
|
||||
new_image,
|
||||
((ph0, ph1), (pw0, pw1), (0, 0)),
|
||||
mode="constant",
|
||||
constant_values=((0, 0), (0, 0), (0, 0)),
|
||||
)
|
||||
new_image = PIL.Image.fromarray(new_image)
|
||||
return new_image
|
||||
|
||||
|
||||
def images_to_video(
|
||||
images: torch.Tensor,
|
||||
output_path: str,
|
||||
fps: int = 30,
|
||||
) -> None:
|
||||
# images: (N, C, H, W)
|
||||
video_dir = os.path.dirname(output_path)
|
||||
video_name = os.path.basename(output_path)
|
||||
os.makedirs(video_dir, exist_ok=True)
|
||||
|
||||
frames = []
|
||||
for i in range(len(images)):
|
||||
frame = (images[i].permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8)
|
||||
assert frame.shape[0] == images.shape[2] and frame.shape[1] == images.shape[3], \
|
||||
f"Frame shape mismatch: {frame.shape} vs {images.shape}"
|
||||
assert frame.min() >= 0 and frame.max() <= 255, \
|
||||
f"Frame value out of range: {frame.min()} ~ {frame.max()}"
|
||||
frames.append(frame)
|
||||
imageio.mimwrite(output_path, np.stack(frames), fps=fps, quality=10)
|
||||
|
||||
|
||||
def save_video(
|
||||
frames: torch.Tensor,
|
||||
output_path: str,
|
||||
fps: int = 30,
|
||||
) -> None:
|
||||
# images: (N, C, H, W)
|
||||
frames = [(frame.permute(1, 2, 0).cpu().numpy() * 255).astype(np.uint8) for frame in frames]
|
||||
writer = imageio.get_writer(output_path, fps=fps)
|
||||
for frame in frames:
|
||||
writer.append_data(frame)
|
||||
writer.close()
|
||||
@@ -0,0 +1,181 @@
|
||||
# Copyright (c) 2022, NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
||||
#
|
||||
# NVIDIA CORPORATION & AFFILIATES and its licensors retain all intellectual property
|
||||
# and proprietary rights in and to this software, related documentation
|
||||
# and any modifications thereto. Any use, reproduction, disclosure or
|
||||
# distribution of this software and related documentation without an express
|
||||
# license agreement from NVIDIA CORPORATION & AFFILIATES is strictly prohibited.
|
||||
|
||||
import torch
|
||||
import xatlas
|
||||
import trimesh
|
||||
import cv2
|
||||
import numpy as np
|
||||
import nvdiffrast.torch as dr
|
||||
from PIL import Image
|
||||
|
||||
|
||||
def save_obj(pointnp_px3, facenp_fx3, colornp_px3, fpath):
|
||||
|
||||
pointnp_px3 = pointnp_px3 @ np.array([[1, 0, 0], [0, 1, 0], [0, 0, -1]])
|
||||
facenp_fx3 = facenp_fx3[:, [2, 1, 0]]
|
||||
|
||||
mesh = trimesh.Trimesh(
|
||||
vertices=pointnp_px3,
|
||||
faces=facenp_fx3,
|
||||
vertex_colors=colornp_px3,
|
||||
)
|
||||
mesh.export(fpath, 'obj')
|
||||
|
||||
|
||||
def save_glb(pointnp_px3, facenp_fx3, colornp_px3, fpath):
|
||||
|
||||
pointnp_px3 = pointnp_px3 @ np.array([[-1, 0, 0], [0, 1, 0], [0, 0, -1]])
|
||||
|
||||
mesh = trimesh.Trimesh(
|
||||
vertices=pointnp_px3,
|
||||
faces=facenp_fx3,
|
||||
vertex_colors=colornp_px3,
|
||||
)
|
||||
mesh.export(fpath, 'glb')
|
||||
|
||||
|
||||
def save_obj_with_mtl(pointnp_px3, tcoords_px2, facenp_fx3, facetex_fx3, texmap_hxwx3, fname):
|
||||
import os
|
||||
fol, na = os.path.split(fname)
|
||||
na, _ = os.path.splitext(na)
|
||||
|
||||
matname = '%s/%s.mtl' % (fol, na)
|
||||
fid = open(matname, 'w')
|
||||
fid.write('newmtl material_0\n')
|
||||
fid.write('Kd 1 1 1\n')
|
||||
fid.write('Ka 0 0 0\n')
|
||||
fid.write('Ks 0.4 0.4 0.4\n')
|
||||
fid.write('Ns 10\n')
|
||||
fid.write('illum 2\n')
|
||||
fid.write('map_Kd %s.png\n' % na)
|
||||
fid.close()
|
||||
####
|
||||
|
||||
fid = open(fname, 'w')
|
||||
fid.write('mtllib %s.mtl\n' % na)
|
||||
|
||||
for pidx, p in enumerate(pointnp_px3):
|
||||
pp = p
|
||||
fid.write('v %f %f %f\n' % (pp[0], pp[1], pp[2]))
|
||||
|
||||
for pidx, p in enumerate(tcoords_px2):
|
||||
pp = p
|
||||
fid.write('vt %f %f\n' % (pp[0], pp[1]))
|
||||
|
||||
fid.write('usemtl material_0\n')
|
||||
for i, f in enumerate(facenp_fx3):
|
||||
f1 = f + 1
|
||||
f2 = facetex_fx3[i] + 1
|
||||
fid.write('f %d/%d %d/%d %d/%d\n' % (f1[0], f2[0], f1[1], f2[1], f1[2], f2[2]))
|
||||
fid.close()
|
||||
|
||||
# save texture map
|
||||
lo, hi = 0, 1
|
||||
img = np.asarray(texmap_hxwx3, dtype=np.float32)
|
||||
img = (img - lo) * (255 / (hi - lo))
|
||||
img = img.clip(0, 255)
|
||||
mask = np.sum(img.astype(np.float32), axis=-1, keepdims=True)
|
||||
mask = (mask <= 3.0).astype(np.float32)
|
||||
kernel = np.ones((3, 3), 'uint8')
|
||||
dilate_img = cv2.dilate(img, kernel, iterations=1)
|
||||
img = img * (1 - mask) + dilate_img * mask
|
||||
img = img.clip(0, 255).astype(np.uint8)
|
||||
Image.fromarray(np.ascontiguousarray(img[::-1, :, :]), 'RGB').save(f'{fol}/{na}.png')
|
||||
|
||||
|
||||
def loadobj(meshfile):
|
||||
v = []
|
||||
f = []
|
||||
meshfp = open(meshfile, 'r')
|
||||
for line in meshfp.readlines():
|
||||
data = line.strip().split(' ')
|
||||
data = [da for da in data if len(da) > 0]
|
||||
if len(data) != 4:
|
||||
continue
|
||||
if data[0] == 'v':
|
||||
v.append([float(d) for d in data[1:]])
|
||||
if data[0] == 'f':
|
||||
data = [da.split('/')[0] for da in data]
|
||||
f.append([int(d) for d in data[1:]])
|
||||
meshfp.close()
|
||||
|
||||
# torch need int64
|
||||
facenp_fx3 = np.array(f, dtype=np.int64) - 1
|
||||
pointnp_px3 = np.array(v, dtype=np.float32)
|
||||
return pointnp_px3, facenp_fx3
|
||||
|
||||
|
||||
def loadobjtex(meshfile):
|
||||
v = []
|
||||
vt = []
|
||||
f = []
|
||||
ft = []
|
||||
meshfp = open(meshfile, 'r')
|
||||
for line in meshfp.readlines():
|
||||
data = line.strip().split(' ')
|
||||
data = [da for da in data if len(da) > 0]
|
||||
if not ((len(data) == 3) or (len(data) == 4) or (len(data) == 5)):
|
||||
continue
|
||||
if data[0] == 'v':
|
||||
assert len(data) == 4
|
||||
|
||||
v.append([float(d) for d in data[1:]])
|
||||
if data[0] == 'vt':
|
||||
if len(data) == 3 or len(data) == 4:
|
||||
vt.append([float(d) for d in data[1:3]])
|
||||
if data[0] == 'f':
|
||||
data = [da.split('/') for da in data]
|
||||
if len(data) == 4:
|
||||
f.append([int(d[0]) for d in data[1:]])
|
||||
ft.append([int(d[1]) for d in data[1:]])
|
||||
elif len(data) == 5:
|
||||
idx1 = [1, 2, 3]
|
||||
data1 = [data[i] for i in idx1]
|
||||
f.append([int(d[0]) for d in data1])
|
||||
ft.append([int(d[1]) for d in data1])
|
||||
idx2 = [1, 3, 4]
|
||||
data2 = [data[i] for i in idx2]
|
||||
f.append([int(d[0]) for d in data2])
|
||||
ft.append([int(d[1]) for d in data2])
|
||||
meshfp.close()
|
||||
|
||||
# torch need int64
|
||||
facenp_fx3 = np.array(f, dtype=np.int64) - 1
|
||||
ftnp_fx3 = np.array(ft, dtype=np.int64) - 1
|
||||
pointnp_px3 = np.array(v, dtype=np.float32)
|
||||
uvs = np.array(vt, dtype=np.float32)
|
||||
return pointnp_px3, facenp_fx3, uvs, ftnp_fx3
|
||||
|
||||
|
||||
# ==============================================================================================
|
||||
def interpolate(attr, rast, attr_idx, rast_db=None):
|
||||
return dr.interpolate(attr.contiguous(), rast, attr_idx, rast_db=rast_db, diff_attrs=None if rast_db is None else 'all')
|
||||
|
||||
|
||||
def xatlas_uvmap(ctx, mesh_v, mesh_pos_idx, resolution):
|
||||
vmapping, indices, uvs = xatlas.parametrize(mesh_v.detach().cpu().numpy(), mesh_pos_idx.detach().cpu().numpy())
|
||||
|
||||
# Convert to tensors
|
||||
indices_int64 = indices.astype(np.uint64, casting='same_kind').view(np.int64)
|
||||
|
||||
uvs = torch.tensor(uvs, dtype=torch.float32, device=mesh_v.device)
|
||||
mesh_tex_idx = torch.tensor(indices_int64, dtype=torch.int64, device=mesh_v.device)
|
||||
# mesh_v_tex. ture
|
||||
uv_clip = uvs[None, ...] * 2.0 - 1.0
|
||||
|
||||
# pad to four component coordinate
|
||||
uv_clip4 = torch.cat((uv_clip, torch.zeros_like(uv_clip[..., 0:1]), torch.ones_like(uv_clip[..., 0:1])), dim=-1)
|
||||
|
||||
# rasterize
|
||||
rast, _ = dr.rasterize(ctx, uv_clip4, mesh_tex_idx.int(), (resolution, resolution))
|
||||
|
||||
# Interpolate world space position
|
||||
gb_pos, _ = interpolate(mesh_v[None, ...], rast, mesh_pos_idx.int())
|
||||
mask = rast[..., 3:4] > 0
|
||||
return uvs, mesh_tex_idx, gb_pos, mask
|
||||
@@ -0,0 +1,26 @@
|
||||
import importlib
|
||||
|
||||
|
||||
def count_params(model, verbose=False):
|
||||
total_params = sum(p.numel() for p in model.parameters())
|
||||
if verbose:
|
||||
print(f"{model.__class__.__name__} has {total_params*1.e-6:.2f} M params.")
|
||||
return total_params
|
||||
|
||||
|
||||
def instantiate_from_config(config):
|
||||
if not "target" in config:
|
||||
if config == '__is_first_stage__':
|
||||
return None
|
||||
elif config == "__is_unconditional__":
|
||||
return None
|
||||
raise KeyError("Expected key `target` to instantiate.")
|
||||
return get_obj_from_str(config["target"])(**config.get("params", dict()))
|
||||
|
||||
|
||||
def get_obj_from_str(string, reload=False):
|
||||
module, cls = string.rsplit(".", 1)
|
||||
if reload:
|
||||
module_imp = importlib.import_module(module)
|
||||
importlib.reload(module_imp)
|
||||
return getattr(importlib.import_module(module, package=None), cls)
|
||||
@@ -55,7 +55,7 @@ class Mesh:
|
||||
self.ft = ft
|
||||
# will first see if there is vertex color to use
|
||||
self.vc = vc
|
||||
# only support a single albedo image
|
||||
# only support a single albedo image, (H, W, 3)
|
||||
self.albedo = albedo
|
||||
# pbr extension, metallic(Blue) = metallicRoughness[..., 2], roughness(Green) = metallicRoughness[..., 1]
|
||||
# ref: https://registry.khronos.org/glTF/specs/2.0/glTF-2.0.html
|
||||
|
||||
@@ -11,11 +11,27 @@ import json
|
||||
|
||||
import torch
|
||||
from torch.utils.data import DataLoader
|
||||
from torchvision.transforms import v2
|
||||
import torchvision.transforms.functional as TF
|
||||
import numpy as np
|
||||
from safetensors.torch import load_file
|
||||
from einops import rearrange
|
||||
|
||||
from diffusers import (
|
||||
DiffusionPipeline,
|
||||
StableDiffusionPipeline
|
||||
)
|
||||
|
||||
from diffusers import (
|
||||
EulerAncestralDiscreteScheduler,
|
||||
EulerDiscreteScheduler,
|
||||
DDIMScheduler,
|
||||
DDIMParallelScheduler,
|
||||
LCMScheduler,
|
||||
KDPM2AncestralDiscreteScheduler,
|
||||
KDPM2DiscreteScheduler,
|
||||
)
|
||||
|
||||
from plyfile import PlyData
|
||||
from PIL import Image
|
||||
|
||||
@@ -32,7 +48,8 @@ from .mesh_processer.mesh_utils import (
|
||||
)
|
||||
from .algorithms.main_3DGS import GaussianSplatting, GaussianSplattingCameraController, GSParams
|
||||
from .algorithms.main_3DGS_renderer import GaussianSplattingRenderer
|
||||
from .algorithms.diff_mesh import DiffMesh
|
||||
from .algorithms.diff_mesh import DiffMesh, DiffMeshCameraController
|
||||
from .algorithms.diff_mesh import DiffRastRenderer
|
||||
from .algorithms.dmtet import DMTetMesh
|
||||
from .algorithms.triplane_gaussian_transformers import TGS
|
||||
from .algorithms.large_multiview_gaussian_model import LGM
|
||||
@@ -52,13 +69,29 @@ from .mvdiffusion.data.single_image_dataset import SingleImageDataset
|
||||
from .mvdiffusion.utils.misc import load_config as load_config_wonder3d
|
||||
from .tsr.system import TSR
|
||||
from .crm.model.crm.model import CRM
|
||||
from .zero123plus.pipeline import Zero123PlusPipeline
|
||||
from .instant_mesh.utils.camera_util import oribt_camera_poses_to_input_cameras
|
||||
|
||||
from .shared_utils.image_utils import prepare_torch_img, torch_img_to_pil_rgba
|
||||
from .shared_utils.common_utils import cstr, parse_save_filename, get_list_filenames
|
||||
from .shared_utils.image_utils import prepare_torch_img, torch_img_to_pil_rgba, troch_image_dilate
|
||||
from .shared_utils.common_utils import cstr, parse_save_filename, get_list_filenames, resume_or_download_model_from_hf
|
||||
|
||||
DIFFUSERS_PIPE_DICT = OrderedDict([
|
||||
("MVDreamPipeline", MVDreamPipeline),
|
||||
("Wonder3DMVDiffusionPipeline", MVDiffusionImagePipeline),
|
||||
("Zero123PlusPipeline", Zero123PlusPipeline),
|
||||
("DiffusionPipeline", DiffusionPipeline),
|
||||
("StableDiffusionPipeline", StableDiffusionPipeline),
|
||||
])
|
||||
|
||||
DIFFUSERS_SCHEDULER_DICT = OrderedDict([
|
||||
("EulerAncestralDiscreteScheduler", EulerAncestralDiscreteScheduler),
|
||||
("Wonder3DMVDiffusionPipeline", MVDiffusionImagePipeline),
|
||||
("EulerDiscreteScheduler,", EulerDiscreteScheduler),
|
||||
("DDIMScheduler,", DDIMScheduler),
|
||||
("DDIMParallelScheduler,", DDIMParallelScheduler),
|
||||
("LCMScheduler,", LCMScheduler),
|
||||
("KDPM2AncestralDiscreteScheduler,", KDPM2AncestralDiscreteScheduler),
|
||||
("KDPM2DiscreteScheduler,", KDPM2DiscreteScheduler),
|
||||
])
|
||||
|
||||
ROOT_PATH = os.path.join(comfy_paths.get_folder_paths("custom_nodes")[0], "ComfyUI-3D-Pack")
|
||||
@@ -94,6 +127,8 @@ AZIMUTH_MAX = 180.0
|
||||
|
||||
WEIGHT_DTYPE = torch.float16
|
||||
|
||||
DEVICE = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
class Preview_3DGS:
|
||||
|
||||
@classmethod
|
||||
@@ -437,7 +472,7 @@ class Stack_Orbit_Camera_Poses:
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"ORBIT_CAMPOSES", # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
|
||||
"ORBIT_CAMPOSES", # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
"FLOAT",
|
||||
"FLOAT",
|
||||
"FLOAT",
|
||||
@@ -446,7 +481,7 @@ class Stack_Orbit_Camera_Poses:
|
||||
"FLOAT",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"orbit_camposes",
|
||||
"orbit_camposes", # List of 6 lists
|
||||
"orbit_radius_list",
|
||||
"elevation_list",
|
||||
"azimuth_list",
|
||||
@@ -612,7 +647,7 @@ class Generate_Orbit_Camera_Poses:
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"ORBIT_CAMPOSES", # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
|
||||
"ORBIT_CAMPOSES", # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"orbit_camposes",
|
||||
@@ -673,6 +708,95 @@ class Generate_Orbit_Camera_Poses:
|
||||
cstr(f"[{self.__class__.__name__}] Last end_reference_image_index: {end_reference_image_index} plus 1 must equal to current start_reference_image_index: {start_reference_image_index}").error.print()
|
||||
|
||||
return (orbit_camposes, )
|
||||
|
||||
class Get_Camposes_From_List_Indexed:
|
||||
|
||||
RETURN_TYPES = ("ORBIT_CAMPOSES",)
|
||||
FUNCTION = "get_indexed_camposes"
|
||||
CATEGORY = "Comfy3D/Preprocessor"
|
||||
DESCRIPTION = """
|
||||
Selects and returns the camera poses at the specified indices as an list.
|
||||
"""
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"original_orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
"indexes": ("STRING", {"default": "0, 1, 2", "multiline": True}),
|
||||
},
|
||||
}
|
||||
|
||||
def get_indexed_camposes(self, original_orbit_camera_poses, indexes):
|
||||
|
||||
# Parse the indexes string into a list of integers
|
||||
index_list = [int(index.strip()) for index in indexes.split(',')]
|
||||
|
||||
# Select the camposes at the specified indices
|
||||
orbit_camera_poses = []
|
||||
for pose_list in original_orbit_camera_poses:
|
||||
new_pose_list = [pose_list[i] for i in index_list]
|
||||
orbit_camera_poses.append(new_pose_list)
|
||||
|
||||
return (orbit_camera_poses,)
|
||||
|
||||
class Mesh_Orbit_Renderer:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"mesh": ("MESH",),
|
||||
"render_image_width": ("INT", {"default": 1024, "min": 128, "max": 8192}),
|
||||
"render_image_height": ("INT", {"default": 1024, "min": 128, "max": 8192}),
|
||||
"render_orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
"render_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
|
||||
"render_background_color_r": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"render_background_color_g": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"render_background_color_b": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"force_cuda_rasterize": ("BOOLEAN", {"default": False},),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"MASK",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"rendered_mesh_images", # [Number of Poses, H, W, 3]
|
||||
"rendered_mesh_masks", # [Number of Poses, H, W, 1]
|
||||
)
|
||||
|
||||
FUNCTION = "render_mesh"
|
||||
CATEGORY = "Comfy3D/Preprocessor"
|
||||
|
||||
def render_mesh(
|
||||
self,
|
||||
mesh,
|
||||
render_image_width,
|
||||
render_image_height,
|
||||
render_orbit_camera_poses,
|
||||
render_orbit_camera_fovy,
|
||||
render_background_color_r,
|
||||
render_background_color_g,
|
||||
render_background_color_b,
|
||||
force_cuda_rasterize,
|
||||
):
|
||||
|
||||
renderer = DiffRastRenderer(mesh, force_cuda_rasterize)
|
||||
|
||||
cam_controller = DiffMeshCameraController(
|
||||
renderer,
|
||||
render_image_width,
|
||||
render_image_height,
|
||||
render_orbit_camera_fovy,
|
||||
static_bg=[render_background_color_r, render_background_color_g, render_background_color_b]
|
||||
)
|
||||
|
||||
all_rendered_images, all_rendered_masks = cam_controller.render_all_pose(render_orbit_camera_poses)
|
||||
all_rendered_masks = all_rendered_masks.squeeze(-1) # [N, H, W, 1] -> [N, H, W]
|
||||
|
||||
return (all_rendered_images, all_rendered_masks) # [N, H, W, 3], [N, H, W]
|
||||
|
||||
|
||||
class Gaussian_Splatting_Orbit_Renderer:
|
||||
@classmethod
|
||||
@@ -682,7 +806,7 @@ class Gaussian_Splatting_Orbit_Renderer:
|
||||
"gs_ply": ("GS_PLY",),
|
||||
"render_image_width": ("INT", {"default": 1024, "min": 128, "max": 8192}),
|
||||
"render_image_height": ("INT", {"default": 1024, "min": 128, "max": 8192}),
|
||||
"render_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
|
||||
"render_orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
"render_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
|
||||
"render_background_color_r": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
"render_background_color_g": ("FLOAT", {"default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001}),
|
||||
@@ -740,7 +864,7 @@ class Gaussian_Splatting:
|
||||
"required": {
|
||||
"reference_images": ("IMAGE",),
|
||||
"reference_masks": ("MASK",),
|
||||
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
|
||||
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
"reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
|
||||
"training_iterations": ("INT", {"default": 30_000, "min": 1, "max": 0xffffffffffffffff}),
|
||||
"batch_size": ("INT", {"default": 1, "min": 1, "max": 0xffffffffffffffff}),
|
||||
@@ -893,7 +1017,7 @@ class Fitting_Mesh_With_Multiview_Images:
|
||||
"required": {
|
||||
"reference_images": ("IMAGE",),
|
||||
"reference_masks": ("MASK",),
|
||||
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
|
||||
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
"reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
|
||||
"mesh": ("MESH",),
|
||||
"mesh_albedo_width": ("INT", {"default": 1024, "min": 128, "max": 8192}),
|
||||
@@ -1066,15 +1190,8 @@ class Load_Triplane_Gaussian_Transformers:
|
||||
|
||||
config_path = os.path.join(ROOT_PATH, self.tgs_config_path)
|
||||
cfg: ExperimentConfig = load_config_tgs(config_path)
|
||||
|
||||
# Download pre-trained model if it not exist locally
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
if not os.path.isfile(ckpt_path):
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=model_name, repo_type="model")
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
|
||||
ckpt_path = resume_or_download_model_from_hf(self.checkpoints_dir_abs, self.default_repo_id, model_name, self.__class__.__name__)
|
||||
|
||||
cfg.system.weights=ckpt_path
|
||||
tgs_model = TGS(cfg=cfg.system).to(device)
|
||||
@@ -1154,6 +1271,7 @@ class Load_Diffusers_Pipeline:
|
||||
"required": {
|
||||
"diffusers_pipeline_name": (list(DIFFUSERS_PIPE_DICT.keys()),),
|
||||
"model_name": ("STRING", {"default": "ashawkey/imagedream-ipmv-diffusers", "multiline": False}),
|
||||
"custom_pipeline": ("STRING", {"default": "", "multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
@@ -1166,7 +1284,7 @@ class Load_Diffusers_Pipeline:
|
||||
FUNCTION = "load_diffusers_pipe"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def load_diffusers_pipe(self, diffusers_pipeline_name, model_name):
|
||||
def load_diffusers_pipe(self, diffusers_pipeline_name, model_name, custom_pipeline):
|
||||
|
||||
# resume pretrained checkpoint
|
||||
ckpt_path = os.path.join(ROOT_PATH, "checkpoints", model_name)
|
||||
@@ -1176,22 +1294,87 @@ class Load_Diffusers_Pipeline:
|
||||
from huggingface_hub import snapshot_download
|
||||
snapshot_download(repo_id=model_name, local_dir=ckpt_path, repo_type="model")
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
diffusers_pipeline_class = DIFFUSERS_PIPE_DICT[diffusers_pipeline_name]
|
||||
|
||||
# load image dream
|
||||
|
||||
# load diffusers pipeline
|
||||
if not custom_pipeline:
|
||||
custom_pipeline = None
|
||||
|
||||
pipe = diffusers_pipeline_class.from_pretrained(
|
||||
ckpt_path,
|
||||
torch_dtype=WEIGHT_DTYPE,
|
||||
custom_pipeline=custom_pipeline,
|
||||
trust_remote_code=True,
|
||||
).to(device)
|
||||
).to(DEVICE)
|
||||
|
||||
pipe.enable_xformers_memory_efficient_attention()
|
||||
|
||||
return (pipe, )
|
||||
|
||||
class Set_Diffusers_Pipeline_Scheduler:
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"pipe": ("DIFFUSERS_PIPE",),
|
||||
"diffusers_scheduler_name": (list(DIFFUSERS_SCHEDULER_DICT.keys()),),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"DIFFUSERS_PIPE",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"pipe",
|
||||
)
|
||||
FUNCTION = "set_pipe_scheduler"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def set_pipe_scheduler(self, pipe, diffusers_scheduler_name):
|
||||
|
||||
diffusers_scheduler_class = DIFFUSERS_SCHEDULER_DICT[diffusers_scheduler_name]
|
||||
|
||||
pipe.scheduler = diffusers_scheduler_class.from_config(
|
||||
pipe.scheduler.config, timestep_spacing='trailing'
|
||||
)
|
||||
return (pipe, )
|
||||
|
||||
class Set_Diffusers_Pipeline_State_Dict:
|
||||
|
||||
checkpoints_dir = "checkpoints"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
cls.checkpoints_base_dir_abs = os.path.join(ROOT_PATH, cls.checkpoints_dir)
|
||||
return {
|
||||
"required": {
|
||||
"pipe": ("DIFFUSERS_PIPE",),
|
||||
"repo_id": ("STRING", {"default": "TencentARC/InstantMesh", "multiline": False}),
|
||||
"model_name": ("STRING", {"default": "diffusion_pytorch_model.bin", "multiline": False}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"DIFFUSERS_PIPE",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"pipe",
|
||||
)
|
||||
FUNCTION = "set_pipe_state_dict"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def set_pipe_state_dict(self, pipe, repo_id, model_name):
|
||||
|
||||
checkpoints_dir_abs = os.path.join(self.checkpoints_base_dir_abs, repo_id)
|
||||
ckpt_path = resume_or_download_model_from_hf(checkpoints_dir_abs, repo_id, model_name, self.__class__.__name__)
|
||||
|
||||
state_dict = torch.load(ckpt_path, map_location='cpu')
|
||||
pipe.unet.load_state_dict(state_dict, strict=True)
|
||||
|
||||
pipe = pipe.to(DEVICE)
|
||||
|
||||
return (pipe, )
|
||||
|
||||
class Wonder3D_MVDiffusion_Model:
|
||||
|
||||
wonder3d_config_path = "configs/wonder3d_config.yaml"
|
||||
@@ -1373,14 +1556,7 @@ class Load_Large_Multiview_Gaussian_Model:
|
||||
|
||||
lgm_model = LGM(config_defaults[lgb_config])
|
||||
|
||||
# resume pretrained checkpoint
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
if not os.path.isfile(ckpt_path):
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=model_name, repo_type="model")
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
ckpt_path = resume_or_download_model_from_hf(self.checkpoints_dir_abs, self.default_repo_id, model_name, self.__class__.__name__)
|
||||
|
||||
if ckpt_path.endswith('safetensors'):
|
||||
ckpt = load_file(ckpt_path, device='cpu')
|
||||
@@ -1389,9 +1565,7 @@ class Load_Large_Multiview_Gaussian_Model:
|
||||
|
||||
lgm_model.load_state_dict(ckpt, strict=False)
|
||||
|
||||
# device
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
lgm_model = lgm_model.half().to(device)
|
||||
lgm_model = lgm_model.half().to(DEVICE)
|
||||
lgm_model.eval()
|
||||
|
||||
cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
|
||||
@@ -1605,16 +1779,7 @@ class Load_TripoSR_Model:
|
||||
|
||||
def load_TSR(self, model_name, chunk_size):
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
# resume pretrained checkpoint
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
if not os.path.isfile(ckpt_path):
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=model_name, repo_type="model")
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
ckpt_path = resume_or_download_model_from_hf(self.checkpoints_dir_abs, self.default_repo_id, model_name, self.__class__.__name__)
|
||||
|
||||
tsr_model = TSR.from_pretrained(
|
||||
weight_path = ckpt_path,
|
||||
@@ -1622,7 +1787,7 @@ class Load_TripoSR_Model:
|
||||
)
|
||||
|
||||
tsr_model.renderer.set_chunk_size(chunk_size)
|
||||
tsr_model.to(device)
|
||||
tsr_model.to(DEVICE)
|
||||
|
||||
cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
|
||||
|
||||
@@ -1656,8 +1821,6 @@ class TripoSR:
|
||||
def run_TSR(self, tsr_model, reference_image, reference_mask, geometry_extract_resolution, marching_cude_threshold):
|
||||
cstr(f"[{self.__class__.__name__}] Running TripoSR...").msg.print()
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
mesh = None
|
||||
|
||||
image = reference_image[0]
|
||||
@@ -1667,7 +1830,8 @@ class TripoSR:
|
||||
image = Image.fromarray(np.clip(255. * image, 0, 255).astype(np.uint8))
|
||||
image = self.fill_background(image)
|
||||
image = image.convert('RGB')
|
||||
scene_codes = tsr_model([image], device)
|
||||
|
||||
scene_codes = tsr_model([image], DEVICE)
|
||||
meshes = tsr_model.extract_mesh(scene_codes, resolution=geometry_extract_resolution, threshold=marching_cude_threshold)
|
||||
mesh = Mesh.load_trimesh(given_mesh=meshes[0])
|
||||
|
||||
@@ -1720,26 +1884,17 @@ class Load_CRM_MVDiffusion_Model:
|
||||
if not os.path.isabs(crm_config_path):
|
||||
crm_config_path = os.path.join(ROOT_PATH, crm_config_path)
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
# resume pretrained checkpoint
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
if not os.path.isfile(ckpt_path):
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=model_name, repo_type="model")
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
ckpt_path = resume_or_download_model_from_hf(self.checkpoints_dir_abs, self.default_repo_id, model_name, self.__class__.__name__)
|
||||
|
||||
crm_config = OmegaConf.load(crm_config_path)
|
||||
|
||||
crm_mvdiffusion_model = instantiate_from_config(crm_config.model)
|
||||
crm_mvdiffusion_model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
|
||||
crm_mvdiffusion_model = crm_mvdiffusion_model.to(device).to(WEIGHT_DTYPE)
|
||||
crm_mvdiffusion_model.device = device
|
||||
crm_mvdiffusion_model = crm_mvdiffusion_model.to(DEVICE).to(WEIGHT_DTYPE)
|
||||
crm_mvdiffusion_model.device = DEVICE
|
||||
|
||||
crm_mvdiffusion_sampler = get_obj_from_str(crm_config.sampler.target)(
|
||||
crm_mvdiffusion_model, device=device, dtype=WEIGHT_DTYPE, **crm_config.sampler.params
|
||||
crm_mvdiffusion_model, device=DEVICE, dtype=WEIGHT_DTYPE, **crm_config.sampler.params
|
||||
)
|
||||
|
||||
cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
|
||||
@@ -1766,14 +1921,17 @@ class CRM_Images_MVDiffusion_Model:
|
||||
"seed": ("INT", {"default": 1234, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"mv_guidance_scale": ("FLOAT", {"default": 5.5, "min": 0.0, "step": 0.01}),
|
||||
"num_inference_steps": ("INT", {"default": 50, "min": 1}),
|
||||
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"ORBIT_CAMPOSES", # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"multiview_images",
|
||||
"orbit_camposes",
|
||||
)
|
||||
FUNCTION = "run_model"
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
@@ -1801,8 +1959,15 @@ class CRM_Images_MVDiffusion_Model:
|
||||
mv_guidance_scale,
|
||||
num_inference_steps
|
||||
)
|
||||
|
||||
return (multiview_images, )
|
||||
|
||||
azimuths = [-90, 0, 180, 90, 0, 0]
|
||||
elevations = [0, 90, 0, 0, -90, 0]
|
||||
radius = [4.0] * 6
|
||||
center = [0.0] * 6
|
||||
|
||||
orbit_camposes = [azimuths, elevations, radius, center, center, center]
|
||||
|
||||
return (multiview_images, orbit_camposes)
|
||||
|
||||
class CRM_CCMs_MVDiffusion_Model:
|
||||
|
||||
@@ -1837,6 +2002,7 @@ class CRM_CCMs_MVDiffusion_Model:
|
||||
FUNCTION = "run_model"
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
|
||||
@torch.no_grad()
|
||||
def run_model(
|
||||
self,
|
||||
crm_mvdiffusion_sampler,
|
||||
@@ -1869,7 +2035,7 @@ class Load_Convolutional_Reconstruction_Model:
|
||||
checkpoints_dir = "checkpoints/crm"
|
||||
default_ckpt_name = "CRM.pth"
|
||||
default_repo_id = "Zhengyi/CRM"
|
||||
crm_config_path = "configs/crm_configs/specs_objaverse_total.json"
|
||||
config_path = "configs/crm_configs/specs_objaverse_total.json"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
@@ -1895,19 +2061,10 @@ class Load_Convolutional_Reconstruction_Model:
|
||||
|
||||
def load_CRM(self, model_name):
|
||||
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
ckpt_path = resume_or_download_model_from_hf(self.checkpoints_dir_abs, self.default_repo_id, model_name, self.__class__.__name__)
|
||||
|
||||
# resume pretrained checkpoint
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
if not os.path.isfile(ckpt_path):
|
||||
cstr(f"[{self.__class__.__name__}] can't find checkpoint {ckpt_path}, will download it from repo {self.default_repo_id} instead").warning.print()
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(repo_id=self.default_repo_id, local_dir=self.checkpoints_dir_abs, filename=model_name, repo_type="model")
|
||||
ckpt_path = os.path.join(self.checkpoints_dir_abs, model_name)
|
||||
|
||||
crm_conf = json.load(open(os.path.join(ROOT_PATH, self.crm_config_path)))
|
||||
crm_model = CRM(crm_conf).to(device)
|
||||
crm_conf = json.load(open(os.path.join(ROOT_PATH, self.config_path)))
|
||||
crm_model = CRM(crm_conf).to(DEVICE)
|
||||
crm_model.load_state_dict(torch.load(ckpt_path, map_location="cpu"), strict=False)
|
||||
|
||||
cstr(f"[{self.__class__.__name__}] loaded model ckpt from {ckpt_path}").msg.print()
|
||||
@@ -1938,15 +2095,183 @@ class Convolutional_Reconstruction_Model:
|
||||
|
||||
@torch.no_grad()
|
||||
def run_CRM(self, crm_model, multiview_images, multiview_CCMs):
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
|
||||
np_imgs = np.concatenate(multiview_images.cpu().numpy(), 1) # (256, 256*6==1536, 3)
|
||||
np_xyzs = np.concatenate(multiview_CCMs.cpu().numpy(), 1) # (256, 1536, 3)
|
||||
|
||||
mesh = CRMSampler.generate3d(crm_model, np_imgs, np_xyzs, device)
|
||||
mesh = CRMSampler.generate3d(crm_model, np_imgs, np_xyzs, DEVICE)
|
||||
|
||||
return (mesh,)
|
||||
|
||||
class Zero123Plus_Diffusion_Model:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
return {
|
||||
"required": {
|
||||
"zero123plus_pipe": ("DIFFUSERS_PIPE",),
|
||||
"reference_image": ("IMAGE",),
|
||||
"reference_mask": ("MASK",),
|
||||
"seed": ("INT", {"default": 1234, "min": 0, "max": 0xffffffffffffffff}),
|
||||
"guidance_scale": ("FLOAT", {"default": 4.0, "min": 0.0, "step": 0.01}),
|
||||
"num_inference_steps": ("INT", {"default": 28, "min": 1}),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"IMAGE",
|
||||
"ORBIT_CAMPOSES", # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"multiviews",
|
||||
"orbit_camposes",
|
||||
)
|
||||
FUNCTION = "run_model"
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
|
||||
@torch.no_grad()
|
||||
def run_model(
|
||||
self,
|
||||
zero123plus_pipe,
|
||||
reference_image,
|
||||
reference_mask,
|
||||
seed,
|
||||
guidance_scale,
|
||||
num_inference_steps,
|
||||
):
|
||||
|
||||
single_image = torch_img_to_pil_rgba(reference_image, reference_mask)
|
||||
|
||||
seed = int(seed)
|
||||
generator = torch.Generator(device=zero123plus_pipe.unet.device).manual_seed(seed)
|
||||
|
||||
# sampling
|
||||
output_image = zero123plus_pipe(
|
||||
single_image,
|
||||
generator=generator,
|
||||
guidance_scale=guidance_scale,
|
||||
num_inference_steps=num_inference_steps,
|
||||
).images[0]
|
||||
|
||||
multiview_images = np.asarray(output_image, dtype=np.float32) / 255.0
|
||||
multiview_images = torch.from_numpy(multiview_images).permute(2, 0, 1).contiguous().float() # (3, 960, 640)
|
||||
multiview_images = rearrange(multiview_images, 'c (n h) (m w) -> (n m) h w c', n=3, m=2) # (6, 320, 320, 3)
|
||||
|
||||
azimuths = [30, 90, 150, -150, -90, -30]
|
||||
elevations = [-20, 10, -20, 10, -20, 10]
|
||||
radius = [4.0] * 6
|
||||
center = [0.0] * 6
|
||||
|
||||
orbit_camposes = [azimuths, elevations, radius, center, center, center]
|
||||
|
||||
return (multiview_images, orbit_camposes)
|
||||
|
||||
class Load_InstantMesh_Reconstruction_Model:
|
||||
checkpoints_dir = "checkpoints/crm"
|
||||
default_ckpt_names = ["instant_mesh_large.ckpt", "instant_mesh_base.ckpt", "instant_nerf_large.ckpt", "instant_nerf_base.ckpt"]
|
||||
default_repo_id = "TencentARC/InstantMesh"
|
||||
config_root_dir = "configs/InstantMesh_configs"
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(cls):
|
||||
cls.checkpoints_dir_abs = os.path.join(ROOT_PATH, cls.checkpoints_dir)
|
||||
all_models_names = get_list_filenames(cls.checkpoints_dir_abs, SUPPORTED_CHECKPOINTS_EXTENSIONS)
|
||||
for ckpt_name in cls.default_ckpt_names:
|
||||
if ckpt_name not in all_models_names:
|
||||
all_models_names += [ckpt_name]
|
||||
|
||||
return {
|
||||
"required": {
|
||||
"model_name": (all_models_names, ),
|
||||
},
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"LRM_MODEL",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"lrm_model",
|
||||
)
|
||||
FUNCTION = "load_LRM"
|
||||
CATEGORY = "Comfy3D/Import|Export"
|
||||
|
||||
def load_LRM(self, model_name):
|
||||
|
||||
from .instant_mesh.utils.train_util import instantiate_from_config
|
||||
|
||||
is_flexicubes = True if model_name.startswith('instant_mesh') else False
|
||||
|
||||
config_name = model_name.split(".")[0] + ".yaml"
|
||||
config_path = os.path.join(ROOT_PATH, self.config_root_dir, config_name)
|
||||
config = OmegaConf.load(config_path)
|
||||
|
||||
lrm_model = instantiate_from_config(config.model_config)
|
||||
ckpt_path = resume_or_download_model_from_hf(self.checkpoints_dir_abs, self.default_repo_id, model_name, self.__class__.__name__)
|
||||
|
||||
state_dict = torch.load(ckpt_path, map_location='cpu')['state_dict']
|
||||
state_dict = {k[14:]: v for k, v in state_dict.items() if k.startswith('lrm_generator.')}
|
||||
lrm_model.load_state_dict(state_dict, strict=True)
|
||||
|
||||
lrm_model = lrm_model.to(DEVICE)
|
||||
if is_flexicubes:
|
||||
lrm_model.init_flexicubes_geometry(DEVICE, fovy=30.0)
|
||||
lrm_model = lrm_model.eval()
|
||||
|
||||
return (lrm_model, )
|
||||
|
||||
class InstantMesh_Reconstruction_Model:
|
||||
|
||||
@classmethod
|
||||
def INPUT_TYPES(s):
|
||||
return {
|
||||
"required": {
|
||||
"lrm_model": ("LRM_MODEL", ),
|
||||
"multiview_images": ("IMAGE",),
|
||||
"orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
"orbit_camera_fovy": ("FLOAT", {"default": 30.0, "min": 0.0, "max": 180.0, "step": 0.1}),
|
||||
"texture_resolution": ("INT", {"default": 1024, "min": 128, "max": 8192}),
|
||||
}
|
||||
}
|
||||
|
||||
RETURN_TYPES = (
|
||||
"MESH",
|
||||
)
|
||||
RETURN_NAMES = (
|
||||
"mesh",
|
||||
)
|
||||
|
||||
FUNCTION = "run_LRM"
|
||||
CATEGORY = "Comfy3D/Algorithm"
|
||||
|
||||
@torch.no_grad()
|
||||
def run_LRM(self, lrm_model, multiview_images, orbit_camera_poses, orbit_camera_fovy, texture_resolution):
|
||||
|
||||
multiview_images
|
||||
|
||||
images = multiview_images.permute(0, 3, 1, 2).unsqueeze(0).to(DEVICE) # [N, H, W, 3] -> [1, N, 3, H, W]
|
||||
images = v2.functional.resize(images, 320, interpolation=3, antialias=True).clamp(0, 1)
|
||||
|
||||
# convert camera format from orbit to lrm inputs
|
||||
azimuths, elevations, radius = orbit_camera_poses[0], orbit_camera_poses[1], orbit_camera_poses[2]
|
||||
input_cameras = oribt_camera_poses_to_input_cameras(azimuths, elevations, radius=radius, fov=orbit_camera_fovy).to(DEVICE)
|
||||
|
||||
# get triplane
|
||||
planes = lrm_model.forward_planes(images, input_cameras)
|
||||
|
||||
# get mesh
|
||||
mesh_out = lrm_model.extract_mesh(
|
||||
planes,
|
||||
use_texture_map=True,
|
||||
texture_resolution=texture_resolution,
|
||||
)
|
||||
|
||||
vertices, faces, uvs, mesh_tex_idx, tex_map = mesh_out
|
||||
tex_map = troch_image_dilate(tex_map.permute(1, 2, 0)) # [3, H, W] -> [H, W, 3]
|
||||
|
||||
mesh = Mesh(v=vertices, f=faces, vt=uvs, ft=mesh_tex_idx, albedo=tex_map, device=DEVICE)
|
||||
mesh.auto_normal()
|
||||
return (mesh,)
|
||||
|
||||
class Instant_NGP:
|
||||
|
||||
@classmethod
|
||||
@@ -1955,7 +2280,7 @@ class Instant_NGP:
|
||||
"required": {
|
||||
"reference_image": ("IMAGE",),
|
||||
"reference_mask": ("MASK",),
|
||||
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
|
||||
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
"reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
|
||||
"training_iterations": ("INT", {"default": 512, "min": 1, "max": 0xffffffffffffffff}),
|
||||
"training_resolution": ("INT", {"default": 128, "min": 128, "max": 8192}),
|
||||
@@ -1991,22 +2316,21 @@ class Instant_NGP:
|
||||
force_cuda_rast
|
||||
):
|
||||
with torch.inference_mode(False):
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
ngp = InstantNGP(training_resolution).to(device)
|
||||
ngp = InstantNGP(training_resolution).to(DEVICE)
|
||||
ngp.prepare_training(reference_image, reference_mask, reference_orbit_camera_poses, reference_orbit_camera_fovy)
|
||||
ngp.fit_nerf(training_iterations)
|
||||
|
||||
vertices, triangles = marching_cubes_density_to_mesh(ngp.get_density, marching_cude_grids_resolution, marching_cude_grids_batch_size, marching_cude_threshold)
|
||||
|
||||
v = torch.from_numpy(vertices).contiguous().float().to(device)
|
||||
f = torch.from_numpy(triangles).contiguous().int().to(device)
|
||||
v = torch.from_numpy(vertices).contiguous().float().to(DEVICE)
|
||||
f = torch.from_numpy(triangles).contiguous().int().to(DEVICE)
|
||||
|
||||
mesh = Mesh(v=v, f=f, device=device)
|
||||
mesh = Mesh(v=v, f=f, device=DEVICE)
|
||||
mesh.auto_normal()
|
||||
mesh.auto_uv()
|
||||
|
||||
mesh.albedo = color_func_to_albedo(mesh, ngp.get_color, texture_resolution, device=device, force_cuda_rast=force_cuda_rast)
|
||||
mesh.albedo = color_func_to_albedo(mesh, ngp.get_color, texture_resolution, device=DEVICE, force_cuda_rast=force_cuda_rast)
|
||||
|
||||
return (mesh, )
|
||||
|
||||
@@ -2018,7 +2342,7 @@ class FlexiCubes_MVS:
|
||||
"required": {
|
||||
"reference_depth_maps": ("IMAGE",),
|
||||
"reference_masks": ("MASK",),
|
||||
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # (orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z)
|
||||
"reference_orbit_camera_poses": ("ORBIT_CAMPOSES",), # [orbit radius, elevation, azimuth, orbit center X, orbit center Y, orbit center Z]
|
||||
"reference_orbit_camera_fovy": ("FLOAT", {"default": 49.1, "min": 0.0, "max": 180.0, "step": 0.1}),
|
||||
"training_iterations": ("INT", {"default": 512, "min": 1, "max": 0xffffffffffffffff}),
|
||||
"batch_size": ("INT", {"default": 4, "min": 1, "max": 0xffffffffffffffff}),
|
||||
@@ -2071,7 +2395,6 @@ class FlexiCubes_MVS:
|
||||
):
|
||||
|
||||
with torch.inference_mode(False):
|
||||
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
|
||||
|
||||
fc_trainer = FlexiCubesTrainer(
|
||||
training_iterations,
|
||||
@@ -2087,7 +2410,7 @@ class FlexiCubes_MVS:
|
||||
remove_floaters_weight,
|
||||
cube_stabilizer_weight,
|
||||
force_cuda_rast,
|
||||
device=device
|
||||
device=DEVICE
|
||||
)
|
||||
|
||||
fc_trainer.prepare_training(reference_depth_maps, reference_masks, reference_orbit_camera_poses, reference_orbit_camera_fovy, reference_normal_maps)
|
||||
|
||||
@@ -21,7 +21,9 @@ safetensors
|
||||
open_clip_torch
|
||||
|
||||
# for training differentiable tensors
|
||||
torchmetrics
|
||||
pytorch_msssim
|
||||
pytorch-lightning
|
||||
|
||||
# for process images & videos
|
||||
imageio
|
||||
|
||||
@@ -1,9 +1,13 @@
|
||||
from abc import ABC, abstractmethod
|
||||
|
||||
import numpy as np
|
||||
import math
|
||||
from scipy.spatial.transform import Rotation as R
|
||||
|
||||
import torch
|
||||
|
||||
from kiui.cam import orbit_camera
|
||||
|
||||
def dot(x, y):
|
||||
if isinstance(x, np.ndarray):
|
||||
return np.sum(x * y, -1, keepdims=True)
|
||||
@@ -191,4 +195,55 @@ class MiniCam:
|
||||
.cuda()
|
||||
) if projection_matrix is None else projection_matrix
|
||||
self.full_proj_transform = self.world_view_transform @ self.projection_matrix
|
||||
self.camera_center = -torch.tensor(c2w[:3, 3]).cuda()
|
||||
self.camera_center = -torch.tensor(c2w[:3, 3]).cuda()
|
||||
|
||||
class BaseCameraController(ABC):
|
||||
def __init__(self, renderer, cam_size_W, cam_size_H, reference_orbit_camera_fovy, invert_bg_prob=1.0, static_bg=None, device='cuda'):
|
||||
self.device = torch.device(device)
|
||||
|
||||
self.renderer = renderer
|
||||
self.cam = OrbitCamera(cam_size_W, cam_size_H, fovy=reference_orbit_camera_fovy)
|
||||
|
||||
self.invert_bg_prob = invert_bg_prob
|
||||
self.black_bg = torch.tensor([0, 0, 0], dtype=torch.float32, device=self.device)
|
||||
self.white_bg = torch.tensor([1, 1, 1], dtype=torch.float32, device=self.device)
|
||||
self.static_bg = None if static_bg is None else torch.tensor(static_bg, dtype=torch.float32, device=self.device)
|
||||
|
||||
self.pose_init()
|
||||
|
||||
super().__init__()
|
||||
|
||||
def pose_init(self):
|
||||
# Calls after default initialize at the end of __init__()
|
||||
pass
|
||||
|
||||
@abstractmethod
|
||||
def get_render_result(self, render_pose, bg_color):
|
||||
pass
|
||||
|
||||
def render_at_pose(self, cam_pose):
|
||||
radius, elevation, azimuth, center_X, center_Y, center_Z = cam_pose
|
||||
|
||||
orbit_target = np.array([center_X, center_Y, center_Z], dtype=np.float32)
|
||||
render_pose = orbit_camera(elevation, azimuth, radius, target=orbit_target)
|
||||
|
||||
if self.static_bg is None:
|
||||
bg_color = self.white_bg if np.random.rand() > self.invert_bg_prob else self.black_bg
|
||||
else:
|
||||
bg_color = self.static_bg
|
||||
|
||||
return self.get_render_result(render_pose, bg_color)
|
||||
|
||||
def render_all_pose(self, all_cam_poses):
|
||||
all_rendered_images, all_rendered_masks = [], []
|
||||
for cam_pose in all_cam_poses:
|
||||
out = self.render_at_pose(cam_pose)
|
||||
|
||||
image = out["image"] # [3, H, W] in [0, 1]
|
||||
mask = out["alpha"] # [1, H, W] in [0, 1]
|
||||
|
||||
all_rendered_images.append(image)
|
||||
all_rendered_masks.append(mask)
|
||||
|
||||
# [Number of Poses, 3, H, W], [Number of Poses, 1, H, W] both in [0, 1]
|
||||
return torch.stack(all_rendered_images, dim=0), torch.stack(all_rendered_masks, dim=0)
|
||||
@@ -122,4 +122,16 @@ def get_list_filenames(directory, extension_filter=None):
|
||||
if exists(directory):
|
||||
return [f for f in listdir(directory) if isfile(join(directory, f)) and (extension_filter is None or f.lower().endswith(extension_filter))]
|
||||
else:
|
||||
return []
|
||||
return []
|
||||
|
||||
# Download pre-trained model if it not exist locally
|
||||
def resume_or_download_model_from_hf(checkpoints_dir_abs, default_repo_id, model_name, class_name=""):
|
||||
|
||||
ckpt_path = os.path.join(checkpoints_dir_abs, model_name)
|
||||
if not os.path.isfile(ckpt_path):
|
||||
cstr(f"[{class_name}] can't find checkpoint {ckpt_path}, will download it from repo {default_repo_id} instead").warning.print()
|
||||
|
||||
from huggingface_hub import hf_hub_download
|
||||
hf_hub_download(repo_id=default_repo_id, local_dir=checkpoints_dir_abs, filename=model_name, repo_type="model")
|
||||
|
||||
return ckpt_path
|
||||
@@ -24,4 +24,21 @@ def torch_img_to_pil_rgba(img, mask):
|
||||
single_image = torch.cat((img, mask), dim=2).detach().cpu().numpy()
|
||||
single_image = Image.fromarray((single_image * 255).astype(np.uint8), mode="RGBA")
|
||||
|
||||
return single_image
|
||||
return single_image
|
||||
|
||||
def troch_image_dilate(img):
|
||||
"""
|
||||
Remove thin seams on generated texture
|
||||
img (torch): [H, W, C]
|
||||
"""
|
||||
import cv2
|
||||
img = np.asarray(img.cpu().numpy(), dtype=np.float32)
|
||||
img = img * 255
|
||||
img = img.clip(0, 255)
|
||||
mask = np.sum(img.astype(np.float32), axis=-1, keepdims=True)
|
||||
mask = (mask <= 3.0).astype(np.float32)
|
||||
kernel = np.ones((3, 3), 'uint8')
|
||||
dilate_img = cv2.dilate(img, kernel, iterations=1)
|
||||
img = img * (1 - mask) + dilate_img * mask
|
||||
img = (img.clip(0, 255) / 255).astype(np.float32)
|
||||
return torch.from_numpy(img)
|
||||
@@ -0,0 +1,271 @@
|
||||
import os
|
||||
import numpy as np
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.nn.functional as F
|
||||
import pytorch_lightning as pl
|
||||
from tqdm import tqdm
|
||||
from torchvision.transforms import v2
|
||||
from torchvision.utils import make_grid, save_image
|
||||
from einops import rearrange
|
||||
|
||||
from diffusers import DiffusionPipeline, EulerAncestralDiscreteScheduler, DDPMScheduler, UNet2DConditionModel
|
||||
from .pipeline import RefOnlyNoisedUNet
|
||||
|
||||
|
||||
def scale_latents(latents):
|
||||
latents = (latents - 0.22) * 0.75
|
||||
return latents
|
||||
|
||||
|
||||
def unscale_latents(latents):
|
||||
latents = latents / 0.75 + 0.22
|
||||
return latents
|
||||
|
||||
|
||||
def scale_image(image):
|
||||
image = image * 0.5 / 0.8
|
||||
return image
|
||||
|
||||
|
||||
def unscale_image(image):
|
||||
image = image / 0.5 * 0.8
|
||||
return image
|
||||
|
||||
|
||||
def extract_into_tensor(a, t, x_shape):
|
||||
b, *_ = t.shape
|
||||
out = a.gather(-1, t)
|
||||
return out.reshape(b, *((1,) * (len(x_shape) - 1)))
|
||||
|
||||
|
||||
class MVDiffusion(pl.LightningModule):
|
||||
def __init__(
|
||||
self,
|
||||
stable_diffusion_config,
|
||||
drop_cond_prob=0.1,
|
||||
):
|
||||
super(MVDiffusion, self).__init__()
|
||||
|
||||
self.drop_cond_prob = drop_cond_prob
|
||||
|
||||
self.register_schedule()
|
||||
|
||||
# init modules
|
||||
pipeline = DiffusionPipeline.from_pretrained(**stable_diffusion_config)
|
||||
pipeline.scheduler = EulerAncestralDiscreteScheduler.from_config(
|
||||
pipeline.scheduler.config, timestep_spacing='trailing'
|
||||
)
|
||||
self.pipeline = pipeline
|
||||
|
||||
train_sched = DDPMScheduler.from_config(self.pipeline.scheduler.config)
|
||||
if isinstance(self.pipeline.unet, UNet2DConditionModel):
|
||||
self.pipeline.unet = RefOnlyNoisedUNet(self.pipeline.unet, train_sched, self.pipeline.scheduler)
|
||||
|
||||
self.train_scheduler = train_sched # use ddpm scheduler during training
|
||||
|
||||
self.unet = pipeline.unet
|
||||
|
||||
# validation output buffer
|
||||
self.validation_step_outputs = []
|
||||
|
||||
def register_schedule(self):
|
||||
self.num_timesteps = 1000
|
||||
|
||||
# replace scaled_linear schedule with linear schedule as Zero123++
|
||||
beta_start = 0.00085
|
||||
beta_end = 0.0120
|
||||
betas = torch.linspace(beta_start, beta_end, 1000, dtype=torch.float32)
|
||||
|
||||
alphas = 1. - betas
|
||||
alphas_cumprod = torch.cumprod(alphas, dim=0)
|
||||
alphas_cumprod_prev = torch.cat([torch.ones(1, dtype=torch.float64), alphas_cumprod[:-1]], 0)
|
||||
|
||||
self.register_buffer('betas', betas.float())
|
||||
self.register_buffer('alphas_cumprod', alphas_cumprod.float())
|
||||
self.register_buffer('alphas_cumprod_prev', alphas_cumprod_prev.float())
|
||||
|
||||
# calculations for diffusion q(x_t | x_{t-1}) and others
|
||||
self.register_buffer('sqrt_alphas_cumprod', torch.sqrt(alphas_cumprod).float())
|
||||
self.register_buffer('sqrt_one_minus_alphas_cumprod', torch.sqrt(1 - alphas_cumprod).float())
|
||||
|
||||
self.register_buffer('sqrt_recip_alphas_cumprod', torch.sqrt(1. / alphas_cumprod).float())
|
||||
self.register_buffer('sqrt_recipm1_alphas_cumprod', torch.sqrt(1. / alphas_cumprod - 1).float())
|
||||
|
||||
def on_fit_start(self):
|
||||
device = torch.device(f'cuda:{self.global_rank}')
|
||||
self.pipeline.to(device)
|
||||
if self.global_rank == 0:
|
||||
os.makedirs(os.path.join(self.logdir, 'images'), exist_ok=True)
|
||||
os.makedirs(os.path.join(self.logdir, 'images_val'), exist_ok=True)
|
||||
|
||||
def prepare_batch_data(self, batch):
|
||||
# prepare stable diffusion input
|
||||
cond_imgs = batch['cond_imgs'] # (B, C, H, W)
|
||||
cond_imgs = cond_imgs.to(self.device)
|
||||
|
||||
# random resize the condition image
|
||||
cond_size = np.random.randint(128, 513)
|
||||
cond_imgs = v2.functional.resize(cond_imgs, cond_size, interpolation=3, antialias=True).clamp(0, 1)
|
||||
|
||||
target_imgs = batch['target_imgs'] # (B, 6, C, H, W)
|
||||
target_imgs = v2.functional.resize(target_imgs, 320, interpolation=3, antialias=True).clamp(0, 1)
|
||||
target_imgs = rearrange(target_imgs, 'b (x y) c h w -> b c (x h) (y w)', x=3, y=2) # (B, C, 3H, 2W)
|
||||
target_imgs = target_imgs.to(self.device)
|
||||
|
||||
return cond_imgs, target_imgs
|
||||
|
||||
@torch.no_grad()
|
||||
def forward_vision_encoder(self, images):
|
||||
dtype = next(self.pipeline.vision_encoder.parameters()).dtype
|
||||
image_pil = [v2.functional.to_pil_image(images[i]) for i in range(images.shape[0])]
|
||||
image_pt = self.pipeline.feature_extractor_clip(images=image_pil, return_tensors="pt").pixel_values
|
||||
image_pt = image_pt.to(device=self.device, dtype=dtype)
|
||||
global_embeds = self.pipeline.vision_encoder(image_pt, output_hidden_states=False).image_embeds
|
||||
global_embeds = global_embeds.unsqueeze(-2)
|
||||
|
||||
encoder_hidden_states = self.pipeline._encode_prompt("", self.device, 1, False)[0]
|
||||
ramp = global_embeds.new_tensor(self.pipeline.config.ramping_coefficients).unsqueeze(-1)
|
||||
encoder_hidden_states = encoder_hidden_states + global_embeds * ramp
|
||||
|
||||
return encoder_hidden_states
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_condition_image(self, images):
|
||||
dtype = next(self.pipeline.vae.parameters()).dtype
|
||||
image_pil = [v2.functional.to_pil_image(images[i]) for i in range(images.shape[0])]
|
||||
image_pt = self.pipeline.feature_extractor_vae(images=image_pil, return_tensors="pt").pixel_values
|
||||
image_pt = image_pt.to(device=self.device, dtype=dtype)
|
||||
latents = self.pipeline.vae.encode(image_pt).latent_dist.sample()
|
||||
return latents
|
||||
|
||||
@torch.no_grad()
|
||||
def encode_target_images(self, images):
|
||||
dtype = next(self.pipeline.vae.parameters()).dtype
|
||||
# equals to scaling images to [-1, 1] first and then call scale_image
|
||||
images = (images - 0.5) / 0.8 # [-0.625, 0.625]
|
||||
posterior = self.pipeline.vae.encode(images.to(dtype)).latent_dist
|
||||
latents = posterior.sample() * self.pipeline.vae.config.scaling_factor
|
||||
latents = scale_latents(latents)
|
||||
return latents
|
||||
|
||||
def forward_unet(self, latents, t, prompt_embeds, cond_latents):
|
||||
dtype = next(self.pipeline.unet.parameters()).dtype
|
||||
latents = latents.to(dtype)
|
||||
prompt_embeds = prompt_embeds.to(dtype)
|
||||
cond_latents = cond_latents.to(dtype)
|
||||
cross_attention_kwargs = dict(cond_lat=cond_latents)
|
||||
pred_noise = self.pipeline.unet(
|
||||
latents,
|
||||
t,
|
||||
encoder_hidden_states=prompt_embeds,
|
||||
cross_attention_kwargs=cross_attention_kwargs,
|
||||
return_dict=False,
|
||||
)[0]
|
||||
return pred_noise
|
||||
|
||||
def predict_start_from_z_and_v(self, x_t, t, v):
|
||||
return (
|
||||
extract_into_tensor(self.sqrt_alphas_cumprod, t, x_t.shape) * x_t -
|
||||
extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x_t.shape) * v
|
||||
)
|
||||
|
||||
def get_v(self, x, noise, t):
|
||||
return (
|
||||
extract_into_tensor(self.sqrt_alphas_cumprod, t, x.shape) * noise -
|
||||
extract_into_tensor(self.sqrt_one_minus_alphas_cumprod, t, x.shape) * x
|
||||
)
|
||||
|
||||
def training_step(self, batch, batch_idx):
|
||||
# get input
|
||||
cond_imgs, target_imgs = self.prepare_batch_data(batch)
|
||||
|
||||
# sample random timestep
|
||||
B = cond_imgs.shape[0]
|
||||
|
||||
t = torch.randint(0, self.num_timesteps, size=(B,)).long().to(self.device)
|
||||
|
||||
# classifier-free guidance
|
||||
if np.random.rand() < self.drop_cond_prob:
|
||||
prompt_embeds = self.pipeline._encode_prompt([""]*B, self.device, 1, False)
|
||||
cond_latents = self.encode_condition_image(torch.zeros_like(cond_imgs))
|
||||
else:
|
||||
prompt_embeds = self.forward_vision_encoder(cond_imgs)
|
||||
cond_latents = self.encode_condition_image(cond_imgs)
|
||||
|
||||
latents = self.encode_target_images(target_imgs)
|
||||
noise = torch.randn_like(latents)
|
||||
latents_noisy = self.train_scheduler.add_noise(latents, noise, t)
|
||||
|
||||
v_pred = self.forward_unet(latents_noisy, t, prompt_embeds, cond_latents)
|
||||
v_target = self.get_v(latents, noise, t)
|
||||
|
||||
loss, loss_dict = self.compute_loss(v_pred, v_target)
|
||||
|
||||
# logging
|
||||
self.log_dict(loss_dict, prog_bar=True, logger=True, on_step=True, on_epoch=True)
|
||||
self.log("global_step", self.global_step, prog_bar=True, logger=True, on_step=True, on_epoch=False)
|
||||
lr = self.optimizers().param_groups[0]['lr']
|
||||
self.log('lr_abs', lr, prog_bar=True, logger=True, on_step=True, on_epoch=False)
|
||||
|
||||
if self.global_step % 500 == 0 and self.global_rank == 0:
|
||||
with torch.no_grad():
|
||||
latents_pred = self.predict_start_from_z_and_v(latents_noisy, t, v_pred)
|
||||
|
||||
latents = unscale_latents(latents_pred)
|
||||
images = unscale_image(self.pipeline.vae.decode(latents / self.pipeline.vae.config.scaling_factor, return_dict=False)[0]) # [-1, 1]
|
||||
images = (images * 0.5 + 0.5).clamp(0, 1)
|
||||
images = torch.cat([target_imgs, images], dim=-2)
|
||||
|
||||
grid = make_grid(images, nrow=images.shape[0], normalize=True, value_range=(0, 1))
|
||||
save_image(grid, os.path.join(self.logdir, 'images', f'train_{self.global_step:07d}.png'))
|
||||
|
||||
return loss
|
||||
|
||||
def compute_loss(self, noise_pred, noise_gt):
|
||||
loss = F.mse_loss(noise_pred, noise_gt)
|
||||
|
||||
prefix = 'train'
|
||||
loss_dict = {}
|
||||
loss_dict.update({f'{prefix}/loss': loss})
|
||||
|
||||
return loss, loss_dict
|
||||
|
||||
@torch.no_grad()
|
||||
def validation_step(self, batch, batch_idx):
|
||||
# get input
|
||||
cond_imgs, target_imgs = self.prepare_batch_data(batch)
|
||||
|
||||
images_pil = [v2.functional.to_pil_image(cond_imgs[i]) for i in range(cond_imgs.shape[0])]
|
||||
|
||||
outputs = []
|
||||
for cond_img in images_pil:
|
||||
latent = self.pipeline(cond_img, num_inference_steps=75, output_type='latent').images
|
||||
image = unscale_image(self.pipeline.vae.decode(latent / self.pipeline.vae.config.scaling_factor, return_dict=False)[0]) # [-1, 1]
|
||||
image = (image * 0.5 + 0.5).clamp(0, 1)
|
||||
outputs.append(image)
|
||||
outputs = torch.cat(outputs, dim=0).to(self.device)
|
||||
images = torch.cat([target_imgs, outputs], dim=-2)
|
||||
|
||||
self.validation_step_outputs.append(images)
|
||||
|
||||
@torch.no_grad()
|
||||
def on_validation_epoch_end(self):
|
||||
images = torch.cat(self.validation_step_outputs, dim=0)
|
||||
|
||||
all_images = self.all_gather(images)
|
||||
all_images = rearrange(all_images, 'r b c h w -> (r b) c h w')
|
||||
|
||||
if self.global_rank == 0:
|
||||
grid = make_grid(all_images, nrow=8, normalize=True, value_range=(0, 1))
|
||||
save_image(grid, os.path.join(self.logdir, 'images_val', f'val_{self.global_step:07d}.png'))
|
||||
|
||||
self.validation_step_outputs.clear() # free memory
|
||||
|
||||
def configure_optimizers(self):
|
||||
lr = self.learning_rate
|
||||
|
||||
optimizer = torch.optim.AdamW(self.unet.parameters(), lr=lr)
|
||||
scheduler = torch.optim.lr_scheduler.CosineAnnealingWarmRestarts(optimizer, 3000, eta_min=lr/4)
|
||||
|
||||
return {'optimizer': optimizer, 'lr_scheduler': scheduler}
|
||||
@@ -0,0 +1,406 @@
|
||||
from typing import Any, Dict, Optional
|
||||
from diffusers.models import AutoencoderKL, UNet2DConditionModel
|
||||
from diffusers.schedulers import KarrasDiffusionSchedulers
|
||||
|
||||
import numpy
|
||||
import torch
|
||||
import torch.nn as nn
|
||||
import torch.utils.checkpoint
|
||||
import torch.distributed
|
||||
import transformers
|
||||
from collections import OrderedDict
|
||||
from PIL import Image
|
||||
from torchvision import transforms
|
||||
from transformers import CLIPImageProcessor, CLIPTextModel, CLIPTokenizer
|
||||
|
||||
import diffusers
|
||||
from diffusers import (
|
||||
AutoencoderKL,
|
||||
DDPMScheduler,
|
||||
DiffusionPipeline,
|
||||
EulerAncestralDiscreteScheduler,
|
||||
UNet2DConditionModel,
|
||||
ImagePipelineOutput
|
||||
)
|
||||
from diffusers.image_processor import VaeImageProcessor
|
||||
from diffusers.models.attention_processor import Attention, AttnProcessor, XFormersAttnProcessor, AttnProcessor2_0
|
||||
from diffusers.utils.import_utils import is_xformers_available
|
||||
|
||||
|
||||
def to_rgb_image(maybe_rgba: Image.Image):
|
||||
if maybe_rgba.mode == 'RGB':
|
||||
return maybe_rgba
|
||||
elif maybe_rgba.mode == 'RGBA':
|
||||
rgba = maybe_rgba
|
||||
img = numpy.random.randint(255, 256, size=[rgba.size[1], rgba.size[0], 3], dtype=numpy.uint8)
|
||||
img = Image.fromarray(img, 'RGB')
|
||||
img.paste(rgba, mask=rgba.getchannel('A'))
|
||||
return img
|
||||
else:
|
||||
raise ValueError("Unsupported image type.", maybe_rgba.mode)
|
||||
|
||||
|
||||
class ReferenceOnlyAttnProc(torch.nn.Module):
|
||||
def __init__(
|
||||
self,
|
||||
chained_proc,
|
||||
enabled=False,
|
||||
name=None
|
||||
) -> None:
|
||||
super().__init__()
|
||||
self.enabled = enabled
|
||||
self.chained_proc = chained_proc
|
||||
self.name = name
|
||||
|
||||
def __call__(
|
||||
self, attn: Attention, hidden_states, encoder_hidden_states=None, attention_mask=None,
|
||||
mode="w", ref_dict: dict = None, is_cfg_guidance = False
|
||||
) -> Any:
|
||||
if encoder_hidden_states is None:
|
||||
encoder_hidden_states = hidden_states
|
||||
if self.enabled and is_cfg_guidance:
|
||||
res0 = self.chained_proc(attn, hidden_states[:1], encoder_hidden_states[:1], attention_mask)
|
||||
hidden_states = hidden_states[1:]
|
||||
encoder_hidden_states = encoder_hidden_states[1:]
|
||||
if self.enabled:
|
||||
if mode == 'w':
|
||||
ref_dict[self.name] = encoder_hidden_states
|
||||
elif mode == 'r':
|
||||
encoder_hidden_states = torch.cat([encoder_hidden_states, ref_dict.pop(self.name)], dim=1)
|
||||
elif mode == 'm':
|
||||
encoder_hidden_states = torch.cat([encoder_hidden_states, ref_dict[self.name]], dim=1)
|
||||
else:
|
||||
assert False, mode
|
||||
res = self.chained_proc(attn, hidden_states, encoder_hidden_states, attention_mask)
|
||||
if self.enabled and is_cfg_guidance:
|
||||
res = torch.cat([res0, res])
|
||||
return res
|
||||
|
||||
|
||||
class RefOnlyNoisedUNet(torch.nn.Module):
|
||||
def __init__(self, unet: UNet2DConditionModel, train_sched: DDPMScheduler, val_sched: EulerAncestralDiscreteScheduler) -> None:
|
||||
super().__init__()
|
||||
self.unet = unet
|
||||
self.train_sched = train_sched
|
||||
self.val_sched = val_sched
|
||||
|
||||
unet_lora_attn_procs = dict()
|
||||
for name, _ in unet.attn_processors.items():
|
||||
if torch.__version__ >= '2.0':
|
||||
default_attn_proc = AttnProcessor2_0()
|
||||
elif is_xformers_available():
|
||||
default_attn_proc = XFormersAttnProcessor()
|
||||
else:
|
||||
default_attn_proc = AttnProcessor()
|
||||
unet_lora_attn_procs[name] = ReferenceOnlyAttnProc(
|
||||
default_attn_proc, enabled=name.endswith("attn1.processor"), name=name
|
||||
)
|
||||
unet.set_attn_processor(unet_lora_attn_procs)
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
try:
|
||||
return super().__getattr__(name)
|
||||
except AttributeError:
|
||||
return getattr(self.unet, name)
|
||||
|
||||
def forward_cond(self, noisy_cond_lat, timestep, encoder_hidden_states, class_labels, ref_dict, is_cfg_guidance, **kwargs):
|
||||
if is_cfg_guidance:
|
||||
encoder_hidden_states = encoder_hidden_states[1:]
|
||||
class_labels = class_labels[1:]
|
||||
self.unet(
|
||||
noisy_cond_lat, timestep,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
class_labels=class_labels,
|
||||
cross_attention_kwargs=dict(mode="w", ref_dict=ref_dict),
|
||||
**kwargs
|
||||
)
|
||||
|
||||
def forward(
|
||||
self, sample, timestep, encoder_hidden_states, class_labels=None,
|
||||
*args, cross_attention_kwargs,
|
||||
down_block_res_samples=None, mid_block_res_sample=None,
|
||||
**kwargs
|
||||
):
|
||||
cond_lat = cross_attention_kwargs['cond_lat']
|
||||
is_cfg_guidance = cross_attention_kwargs.get('is_cfg_guidance', False)
|
||||
noise = torch.randn_like(cond_lat)
|
||||
if self.training:
|
||||
noisy_cond_lat = self.train_sched.add_noise(cond_lat, noise, timestep)
|
||||
noisy_cond_lat = self.train_sched.scale_model_input(noisy_cond_lat, timestep)
|
||||
else:
|
||||
noisy_cond_lat = self.val_sched.add_noise(cond_lat, noise, timestep.reshape(-1))
|
||||
noisy_cond_lat = self.val_sched.scale_model_input(noisy_cond_lat, timestep.reshape(-1))
|
||||
ref_dict = {}
|
||||
self.forward_cond(
|
||||
noisy_cond_lat, timestep,
|
||||
encoder_hidden_states, class_labels,
|
||||
ref_dict, is_cfg_guidance, **kwargs
|
||||
)
|
||||
weight_dtype = self.unet.dtype
|
||||
return self.unet(
|
||||
sample, timestep,
|
||||
encoder_hidden_states, *args,
|
||||
class_labels=class_labels,
|
||||
cross_attention_kwargs=dict(mode="r", ref_dict=ref_dict, is_cfg_guidance=is_cfg_guidance),
|
||||
down_block_additional_residuals=[
|
||||
sample.to(dtype=weight_dtype) for sample in down_block_res_samples
|
||||
] if down_block_res_samples is not None else None,
|
||||
mid_block_additional_residual=(
|
||||
mid_block_res_sample.to(dtype=weight_dtype)
|
||||
if mid_block_res_sample is not None else None
|
||||
),
|
||||
**kwargs
|
||||
)
|
||||
|
||||
|
||||
def scale_latents(latents):
|
||||
latents = (latents - 0.22) * 0.75
|
||||
return latents
|
||||
|
||||
|
||||
def unscale_latents(latents):
|
||||
latents = latents / 0.75 + 0.22
|
||||
return latents
|
||||
|
||||
|
||||
def scale_image(image):
|
||||
image = image * 0.5 / 0.8
|
||||
return image
|
||||
|
||||
|
||||
def unscale_image(image):
|
||||
image = image / 0.5 * 0.8
|
||||
return image
|
||||
|
||||
|
||||
class DepthControlUNet(torch.nn.Module):
|
||||
def __init__(self, unet: RefOnlyNoisedUNet, controlnet: Optional[diffusers.ControlNetModel] = None, conditioning_scale=1.0) -> None:
|
||||
super().__init__()
|
||||
self.unet = unet
|
||||
if controlnet is None:
|
||||
self.controlnet = diffusers.ControlNetModel.from_unet(unet.unet)
|
||||
else:
|
||||
self.controlnet = controlnet
|
||||
DefaultAttnProc = AttnProcessor2_0
|
||||
if is_xformers_available():
|
||||
DefaultAttnProc = XFormersAttnProcessor
|
||||
self.controlnet.set_attn_processor(DefaultAttnProc())
|
||||
self.conditioning_scale = conditioning_scale
|
||||
|
||||
def __getattr__(self, name: str):
|
||||
try:
|
||||
return super().__getattr__(name)
|
||||
except AttributeError:
|
||||
return getattr(self.unet, name)
|
||||
|
||||
def forward(self, sample, timestep, encoder_hidden_states, class_labels=None, *args, cross_attention_kwargs: dict, **kwargs):
|
||||
cross_attention_kwargs = dict(cross_attention_kwargs)
|
||||
control_depth = cross_attention_kwargs.pop('control_depth')
|
||||
down_block_res_samples, mid_block_res_sample = self.controlnet(
|
||||
sample,
|
||||
timestep,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
controlnet_cond=control_depth,
|
||||
conditioning_scale=self.conditioning_scale,
|
||||
return_dict=False,
|
||||
)
|
||||
return self.unet(
|
||||
sample,
|
||||
timestep,
|
||||
encoder_hidden_states=encoder_hidden_states,
|
||||
down_block_res_samples=down_block_res_samples,
|
||||
mid_block_res_sample=mid_block_res_sample,
|
||||
cross_attention_kwargs=cross_attention_kwargs
|
||||
)
|
||||
|
||||
|
||||
class ModuleListDict(torch.nn.Module):
|
||||
def __init__(self, procs: dict) -> None:
|
||||
super().__init__()
|
||||
self.keys = sorted(procs.keys())
|
||||
self.values = torch.nn.ModuleList(procs[k] for k in self.keys)
|
||||
|
||||
def __getitem__(self, key):
|
||||
return self.values[self.keys.index(key)]
|
||||
|
||||
|
||||
class SuperNet(torch.nn.Module):
|
||||
def __init__(self, state_dict: Dict[str, torch.Tensor]):
|
||||
super().__init__()
|
||||
state_dict = OrderedDict((k, state_dict[k]) for k in sorted(state_dict.keys()))
|
||||
self.layers = torch.nn.ModuleList(state_dict.values())
|
||||
self.mapping = dict(enumerate(state_dict.keys()))
|
||||
self.rev_mapping = {v: k for k, v in enumerate(state_dict.keys())}
|
||||
|
||||
# .processor for unet, .self_attn for text encoder
|
||||
self.split_keys = [".processor", ".self_attn"]
|
||||
|
||||
# we add a hook to state_dict() and load_state_dict() so that the
|
||||
# naming fits with `unet.attn_processors`
|
||||
def map_to(module, state_dict, *args, **kwargs):
|
||||
new_state_dict = {}
|
||||
for key, value in state_dict.items():
|
||||
num = int(key.split(".")[1]) # 0 is always "layers"
|
||||
new_key = key.replace(f"layers.{num}", module.mapping[num])
|
||||
new_state_dict[new_key] = value
|
||||
|
||||
return new_state_dict
|
||||
|
||||
def remap_key(key, state_dict):
|
||||
for k in self.split_keys:
|
||||
if k in key:
|
||||
return key.split(k)[0] + k
|
||||
return key.split('.')[0]
|
||||
|
||||
def map_from(module, state_dict, *args, **kwargs):
|
||||
all_keys = list(state_dict.keys())
|
||||
for key in all_keys:
|
||||
replace_key = remap_key(key, state_dict)
|
||||
new_key = key.replace(replace_key, f"layers.{module.rev_mapping[replace_key]}")
|
||||
state_dict[new_key] = state_dict[key]
|
||||
del state_dict[key]
|
||||
|
||||
self._register_state_dict_hook(map_to)
|
||||
self._register_load_state_dict_pre_hook(map_from, with_module=True)
|
||||
|
||||
|
||||
class Zero123PlusPipeline(diffusers.StableDiffusionPipeline):
|
||||
tokenizer: transformers.CLIPTokenizer
|
||||
text_encoder: transformers.CLIPTextModel
|
||||
vision_encoder: transformers.CLIPVisionModelWithProjection
|
||||
|
||||
feature_extractor_clip: transformers.CLIPImageProcessor
|
||||
unet: UNet2DConditionModel
|
||||
scheduler: diffusers.schedulers.KarrasDiffusionSchedulers
|
||||
|
||||
vae: AutoencoderKL
|
||||
ramping: nn.Linear
|
||||
|
||||
feature_extractor_vae: transformers.CLIPImageProcessor
|
||||
|
||||
depth_transforms_multi = transforms.Compose([
|
||||
transforms.ToTensor(),
|
||||
transforms.Normalize([0.5], [0.5])
|
||||
])
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
vae: AutoencoderKL,
|
||||
text_encoder: CLIPTextModel,
|
||||
tokenizer: CLIPTokenizer,
|
||||
unet: UNet2DConditionModel,
|
||||
scheduler: KarrasDiffusionSchedulers,
|
||||
vision_encoder: transformers.CLIPVisionModelWithProjection,
|
||||
feature_extractor_clip: CLIPImageProcessor,
|
||||
feature_extractor_vae: CLIPImageProcessor,
|
||||
ramping_coefficients: Optional[list] = None,
|
||||
safety_checker=None,
|
||||
):
|
||||
DiffusionPipeline.__init__(self)
|
||||
|
||||
self.register_modules(
|
||||
vae=vae, text_encoder=text_encoder, tokenizer=tokenizer,
|
||||
unet=unet, scheduler=scheduler, safety_checker=None,
|
||||
vision_encoder=vision_encoder,
|
||||
feature_extractor_clip=feature_extractor_clip,
|
||||
feature_extractor_vae=feature_extractor_vae
|
||||
)
|
||||
self.register_to_config(ramping_coefficients=ramping_coefficients)
|
||||
self.vae_scale_factor = 2 ** (len(self.vae.config.block_out_channels) - 1)
|
||||
self.image_processor = VaeImageProcessor(vae_scale_factor=self.vae_scale_factor)
|
||||
|
||||
def prepare(self):
|
||||
train_sched = DDPMScheduler.from_config(self.scheduler.config)
|
||||
if isinstance(self.unet, UNet2DConditionModel):
|
||||
self.unet = RefOnlyNoisedUNet(self.unet, train_sched, self.scheduler).eval()
|
||||
|
||||
def add_controlnet(self, controlnet: Optional[diffusers.ControlNetModel] = None, conditioning_scale=1.0):
|
||||
self.prepare()
|
||||
self.unet = DepthControlUNet(self.unet, controlnet, conditioning_scale)
|
||||
return SuperNet(OrderedDict([('controlnet', self.unet.controlnet)]))
|
||||
|
||||
def encode_condition_image(self, image: torch.Tensor):
|
||||
image = self.vae.encode(image).latent_dist.sample()
|
||||
return image
|
||||
|
||||
@torch.no_grad()
|
||||
def __call__(
|
||||
self,
|
||||
image: Image.Image = None,
|
||||
prompt = "",
|
||||
*args,
|
||||
num_images_per_prompt: Optional[int] = 1,
|
||||
guidance_scale=4.0,
|
||||
depth_image: Image.Image = None,
|
||||
output_type: Optional[str] = "pil",
|
||||
width=640,
|
||||
height=960,
|
||||
num_inference_steps=28,
|
||||
return_dict=True,
|
||||
**kwargs
|
||||
):
|
||||
self.prepare()
|
||||
if image is None:
|
||||
raise ValueError("Inputting embeddings not supported for this pipeline. Please pass an image.")
|
||||
assert not isinstance(image, torch.Tensor)
|
||||
image = to_rgb_image(image)
|
||||
image_1 = self.feature_extractor_vae(images=image, return_tensors="pt").pixel_values
|
||||
image_2 = self.feature_extractor_clip(images=image, return_tensors="pt").pixel_values
|
||||
if depth_image is not None and hasattr(self.unet, "controlnet"):
|
||||
depth_image = to_rgb_image(depth_image)
|
||||
depth_image = self.depth_transforms_multi(depth_image).to(
|
||||
device=self.unet.controlnet.device, dtype=self.unet.controlnet.dtype
|
||||
)
|
||||
image = image_1.to(device=self.vae.device, dtype=self.vae.dtype)
|
||||
image_2 = image_2.to(device=self.vae.device, dtype=self.vae.dtype)
|
||||
cond_lat = self.encode_condition_image(image)
|
||||
if guidance_scale > 1:
|
||||
negative_lat = self.encode_condition_image(torch.zeros_like(image))
|
||||
cond_lat = torch.cat([negative_lat, cond_lat])
|
||||
encoded = self.vision_encoder(image_2, output_hidden_states=False)
|
||||
global_embeds = encoded.image_embeds
|
||||
global_embeds = global_embeds.unsqueeze(-2)
|
||||
|
||||
if hasattr(self, "encode_prompt"):
|
||||
encoder_hidden_states = self.encode_prompt(
|
||||
prompt,
|
||||
self.device,
|
||||
num_images_per_prompt,
|
||||
False
|
||||
)[0]
|
||||
else:
|
||||
encoder_hidden_states = self._encode_prompt(
|
||||
prompt,
|
||||
self.device,
|
||||
num_images_per_prompt,
|
||||
False
|
||||
)
|
||||
ramp = global_embeds.new_tensor(self.config.ramping_coefficients).unsqueeze(-1)
|
||||
encoder_hidden_states = encoder_hidden_states + global_embeds * ramp
|
||||
cak = dict(cond_lat=cond_lat)
|
||||
if hasattr(self.unet, "controlnet"):
|
||||
cak['control_depth'] = depth_image
|
||||
latents: torch.Tensor = super().__call__(
|
||||
None,
|
||||
*args,
|
||||
cross_attention_kwargs=cak,
|
||||
guidance_scale=guidance_scale,
|
||||
num_images_per_prompt=num_images_per_prompt,
|
||||
prompt_embeds=encoder_hidden_states,
|
||||
num_inference_steps=num_inference_steps,
|
||||
output_type='latent',
|
||||
width=width,
|
||||
height=height,
|
||||
**kwargs
|
||||
).images
|
||||
latents = unscale_latents(latents)
|
||||
if not output_type == "latent":
|
||||
image = unscale_image(self.vae.decode(latents / self.vae.config.scaling_factor, return_dict=False)[0])
|
||||
else:
|
||||
image = latents
|
||||
|
||||
image = self.image_processor.postprocess(image, output_type=output_type)
|
||||
if not return_dict:
|
||||
return (image,)
|
||||
|
||||
return ImagePipelineOutput(images=image)
|
||||
Reference in New Issue
Block a user