nodes/visualbruno-ComfyUI-Trellis2
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

Added the nodes "MeshWithVoxel Multi-View Generator" and "Mesh Texturing Multi-View"

Browse Source 
Updated/Added workflows
This commit is contained in:
Bruno Fargnoli
2026-02-12 14:20:37 +01:00
parent a19110a28a
commit f75fc3ac8c
9 changed files with 2837 additions and 555 deletions
Download Patch File Download Diff File Expand all files Collapse all files
README.md
+3
View File
@@ -14,6 +14,9 @@
| Date | Description |
| --- | --- |
| **2026-02-12** | Added the node "Mesh With Voxel Multi-View Generator" |
|| Added the node "Mesh Texturing Multi-View" |
|| Added new example workflows |
| **2026-02-10** | Improved progress bar when filling holes with meshlib |
| **2026-02-09** | Fixed "Mesh Texturing" node<br>"mesh_cluster_threshold_cone_half_angle_rad" was not used |
| **2026-02-08** | Fixed "Fill Holes" node progress bar<br>Updated Cumesh package<br>Added "Remesh with Quad" node<br>Added "Batch Simplify Mesh and Export" node|
example_workflows/MultiViews.json
+486 -552
View File
File diff suppressed because it is too large Load Diff
example_workflows/MultiViews_MeshOnly.json
+690
View File
@@ -0,0 +1,690 @@
{
"id": "440427ce-0c6a-4462-af51-639ed2f16dec",
"revision": 0,
"last_node_id": 121,
"last_link_id": 230,
"nodes": [
{
"id": 50,
"type": "Trellis2PreProcessImage",
"pos": [
-514.210025695023,
250.92630997998222
],
"size": [
297.84375,
107.328125
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 110
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
223
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "15854282a73cf231b81d52ada22e652f414a078b",
"Node name for S&R": "Trellis2PreProcessImage",
"widget_ue_connectable": {}
},
"widgets_values": [
25,
false
]
},
{
"id": 6,
"type": "Trellis2LoadImageWithTransparency",
"pos": [
-1354.11041692169,
-103.28960088005408
],
"size": [
610.640625,
788.515625
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": null
},
{
"name": "mask",
"type": "MASK",
"links": null
},
{
"name": "image_with_alpha",
"type": "IMAGE",
"links": [
110
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "e7f9b30df7a09bcedf1c955e176754c73f983254",
"Node name for S&R": "Trellis2LoadImageWithTransparency",
"widget_ue_connectable": {}
},
"widgets_values": [
"Image_1024_00010_.png",
"image"
]
},
{
"id": 39,
"type": "Trellis2LoadModel",
"pos": [
-502.8470854177922,
-74.32512914672895
],
"size": [
336.0625,
207.328125
],
"flags": {
"collapsed": false
},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "pipeline",
"type": "TRELLIS2PIPELINE",
"links": [
222
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "c5d66aa46bf1bbf903fd4fe9ff086ac774bb7e03",
"Node name for S&R": "Trellis2LoadModel",
"widget_ue_connectable": {}
},
"widgets_values": [
"TRELLIS.2-4B",
"flash_attn",
"cuda",
true,
false
]
},
{
"id": 103,
"type": "Trellis2RemeshWithQuad",
"pos": [
467.2386692680435,
243.87402592407818
],
"size": [
455.3125,
262
],
"flags": {
"collapsed": false
},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESHWITHVOXEL",
"link": 225
}
],
"outputs": [
{
"name": "mesh",
"type": "MESHWITHVOXEL",
"links": [
226
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "258cd607667d64b01b2abdaded4e59016ffd5cb6",
"Node name for S&R": "Trellis2RemeshWithQuad",
"widget_ue_connectable": {}
},
"widgets_values": [
1,
0,
false,
0.03,
"1024",
true,
true
]
},
{
"id": 114,
"type": "Trellis2LoadImageWithTransparency",
"pos": [
-1359.0614927706545,
735.4082082769835
],
"size": [
609.125,
759.125
],
"flags": {
"collapsed": false
},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": []
},
{
"name": "mask",
"type": "MASK",
"links": []
},
{
"name": "image_with_alpha",
"type": "IMAGE",
"links": [
221
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "e7f9b30df7a09bcedf1c955e176754c73f983254",
"widget_ue_connectable": {},
"Node name for S&R": "Trellis2LoadImageWithTransparency"
},
"widgets_values": [
"Image_1024_00021_.png",
"image"
]
},
{
"id": 115,
"type": "Trellis2PreProcessImage",
"pos": [
-520.4300253912838,
427.5271012938298
],
"size": [
311.921875,
107.328125
],
"flags": {
"collapsed": false
},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 221
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
224
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "15854282a73cf231b81d52ada22e652f414a078b",
"widget_ue_connectable": {},
"Node name for S&R": "Trellis2PreProcessImage"
},
"widgets_values": [
50,
false
]
},
{
"id": 116,
"type": "Trellis2MeshWithVoxelMultiViewGenerator",
"pos": [
-76.41514399946482,
241.60506031220336
],
"size": [
468.21875,
1000.53125
],
"flags": {},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "pipeline",
"type": "TRELLIS2PIPELINE",
"link": 222
},
{
"name": "front_image",
"type": "IMAGE",
"link": 223
},
{
"name": "back_image",
"shape": 7,
"type": "IMAGE",
"link": 224
},
{
"name": "left_image",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"name": "right_image",
"shape": 7,
"type": "IMAGE",
"link": null
}
],
"outputs": [
{
"name": "mesh",
"type": "MESHWITHVOXEL",
"links": [
225
]
},
{
"name": "bvh",
"type": "BVH",
"links": null
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "a19110a28a5c5c434386af0421005dd8edae82db",
"widget_ue_connectable": {},
"Node name for S&R": "Trellis2MeshWithVoxelMultiViewGenerator"
},
"widgets_values": [
12345,
"randomize",
"1024_cascade",
25,
6.5,
0.2,
4,
25,
6.5,
0.2,
4,
12,
3,
0.2,
3,
999999,
32,
false,
0.1,
1,
0.1,
1,
0,
0.9,
true,
"z",
2
]
},
{
"id": 117,
"type": "Trellis2SimplifyMesh",
"pos": [
472.141197337784,
564.9329628175803
],
"size": [
317,
115.046875
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESHWITHVOXEL",
"link": 226
}
],
"outputs": [
{
"name": "mesh",
"type": "MESHWITHVOXEL",
"links": [
227
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "a19110a28a5c5c434386af0421005dd8edae82db",
"widget_ue_connectable": {},
"Node name for S&R": "Trellis2SimplifyMesh"
},
"widgets_values": [
500000,
"Cumesh"
]
},
{
"id": 118,
"type": "Trellis2FillHolesWithMeshlib",
"pos": [
473.1986723221411,
743.6571759057757
],
"size": [
348.71875,
71.328125
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESHWITHVOXEL",
"link": 227
}
],
"outputs": [
{
"name": "mesh",
"type": "MESHWITHVOXEL",
"links": [
228
]
},
{
"name": "holes_filled",
"type": "INT",
"links": null
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "a19110a28a5c5c434386af0421005dd8edae82db",
"widget_ue_connectable": {},
"Node name for S&R": "Trellis2FillHolesWithMeshlib"
},
"widgets_values": []
},
{
"id": 119,
"type": "Trellis2ExportMesh",
"pos": [
471.08304460897625,
1037.6532301953084
],
"size": [
348.71875,
148.109375
],
"flags": {},
"order": 10,
"mode": 0,
"inputs": [
{
"name": "trimesh",
"type": "TRIMESH",
"link": 229
}
],
"outputs": [
{
"name": "glb_path",
"type": "STRING",
"links": [
230
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "a19110a28a5c5c434386af0421005dd8edae82db",
"widget_ue_connectable": {},
"Node name for S&R": "Trellis2ExportMesh"
},
"widgets_values": [
"Trellis2MV",
"glb",
true
]
},
{
"id": 120,
"type": "Trellis2MeshWithVoxelToTrimesh",
"pos": [
472.14061641396916,
874.7920616685096
],
"size": [
332.859375,
92.5625
],
"flags": {},
"order": 9,
"mode": 0,
"inputs": [
{
"name": "mesh",
"type": "MESHWITHVOXEL",
"link": 228
}
],
"outputs": [
{
"name": "trimesh",
"type": "TRIMESH",
"links": [
229
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "a19110a28a5c5c434386af0421005dd8edae82db",
"widget_ue_connectable": {},
"Node name for S&R": "Trellis2MeshWithVoxelToTrimesh"
},
"widgets_values": [
"90 degrees"
]
},
{
"id": 121,
"type": "Preview3D",
"pos": [
960.7538631636339,
252.95888156302374
],
"size": [
892.296875,
936.96875
],
"flags": {},
"order": 11,
"mode": 0,
"inputs": [
{
"name": "camera_info",
"shape": 7,
"type": "LOAD3D_CAMERA",
"link": null
},
{
"name": "bg_image",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"name": "model_file",
"type": "STRING,FILE_3D_GLB,FILE_3D_GLTF,FILE_3D_FBX,FILE_3D_OBJ,FILE_3D_STL,FILE_3D_USDZ,FILE_3D",
"widget": {
"name": "model_file"
},
"link": 230
}
],
"outputs": [],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.13.0",
"widget_ue_connectable": {},
"Node name for S&R": "Preview3D"
},
"widgets_values": [
"",
""
]
}
],
"links": [
[
110,
6,
2,
50,
0,
"IMAGE"
],
[
221,
114,
2,
115,
0,
"IMAGE"
],
[
222,
39,
0,
116,
0,
"TRELLIS2PIPELINE"
],
[
223,
50,
0,
116,
1,
"IMAGE"
],
[
224,
115,
0,
116,
2,
"IMAGE"
],
[
225,
116,
0,
103,
0,
"MESHWITHVOXEL"
],
[
226,
103,
0,
117,
0,
"MESHWITHVOXEL"
],
[
227,
117,
0,
118,
0,
"MESHWITHVOXEL"
],
[
228,
118,
0,
120,
0,
"MESHWITHVOXEL"
],
[
229,
120,
0,
119,
0,
"TRIMESH"
],
[
230,
119,
0,
121,
2,
"STRING"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "Vue",
"ue_links": [],
"ds": {
"scale": 0.4736244074476824,
"offset": [
1666.544822537185,
301.4636979965617
]
},
"links_added_by_ue": [],
"frontendVersion": "1.38.13",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
example_workflows/MultiViews_TextureMesh.json
+561
View File
@@ -0,0 +1,561 @@
{
"id": "9b1321c3-d206-474a-bf17-498b285fb45b",
"revision": 0,
"last_node_id": 37,
"last_link_id": 83,
"nodes": [
{
"id": 12,
"type": "Preview3D",
"pos": [
1360.4666161520581,
1420.8767073529182
],
"size": [
1144.84375,
1103.03125
],
"flags": {},
"order": 8,
"mode": 0,
"inputs": [
{
"name": "camera_info",
"shape": 7,
"type": "LOAD3D_CAMERA",
"link": null
},
{
"name": "bg_image",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"name": "model_file",
"type": "STRING,FILE_3D_GLB,FILE_3D_GLTF,FILE_3D_FBX,FILE_3D_OBJ,FILE_3D_STL,FILE_3D_USDZ,FILE_3D",
"widget": {
"name": "model_file"
},
"link": 12
}
],
"outputs": [],
"properties": {
"cnr_id": "comfy-core",
"ver": "0.5.1",
"Node name for S&R": "Preview3D",
"widget_ue_connectable": {},
"Last Time Model File": "DwarfMV_fromMV_00001_.glb",
"Scene Config": {
"showGrid": false,
"backgroundColor": "#282828",
"backgroundImage": "",
"backgroundRenderMode": "tiled"
},
"Camera Config": {
"cameraType": "perspective",
"fov": 35,
"state": {
"position": {
"x": -2.356233924150018,
"y": 5.529086349255764,
"z": 8.664564277092959
},
"target": {
"x": 0,
"y": 2.5,
"z": 0
},
"zoom": 1,
"cameraType": "perspective"
}
},
"Light Config": {
"intensity": 3
},
"Model Config": {
"upDirection": "original",
"materialMode": "original",
"showSkeleton": false
}
},
"widgets_values": [
"DwarfMV_fromMV_00001_.glb",
""
]
},
{
"id": 2,
"type": "Trellis2LoadModel",
"pos": [
264.5598759967819,
993.6174889527508
],
"size": [
362.0625,
207.328125
],
"flags": {},
"order": 0,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "pipeline",
"type": "TRELLIS2PIPELINE",
"links": [
78
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "c5d66aa46bf1bbf903fd4fe9ff086ac774bb7e03",
"Node name for S&R": "Trellis2LoadModel",
"widget_ue_connectable": {}
},
"widgets_values": [
"TRELLIS.2-4B",
"flash_attn",
"cuda",
true,
false
]
},
{
"id": 18,
"type": "Trellis2PreProcessImage",
"pos": [
263.0397099736384,
1453.040730438439
],
"size": [
360.984375,
107.328125
],
"flags": {
"collapsed": false
},
"order": 4,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 32
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
79
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "15854282a73cf231b81d52ada22e652f414a078b",
"Node name for S&R": "Trellis2PreProcessImage",
"widget_ue_connectable": {}
},
"widgets_values": [
25,
false
]
},
{
"id": 1,
"type": "Trellis2LoadImageWithTransparency",
"pos": [
-293.9350983507024,
1366.6009200665326
],
"size": [
470.78125,
620.78125
],
"flags": {},
"order": 1,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": []
},
{
"name": "mask",
"type": "MASK",
"links": []
},
{
"name": "image_with_alpha",
"type": "IMAGE",
"links": [
32
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "e7f9b30df7a09bcedf1c955e176754c73f983254",
"Node name for S&R": "Trellis2LoadImageWithTransparency",
"widget_ue_connectable": {}
},
"widgets_values": [
"Image_1024_00020_.png",
"image"
]
},
{
"id": 21,
"type": "Trellis2LoadImageWithTransparency",
"pos": [
-285.4316533126076,
2056.8396735019032
],
"size": [
465.203125,
615.203125
],
"flags": {
"collapsed": false
},
"order": 2,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": []
},
{
"name": "mask",
"type": "MASK",
"links": []
},
{
"name": "image_with_alpha",
"type": "IMAGE",
"links": [
47
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "e7f9b30df7a09bcedf1c955e176754c73f983254",
"Node name for S&R": "Trellis2LoadImageWithTransparency",
"widget_ue_connectable": {}
},
"widgets_values": [
"Image_1024_00021_.png",
"image"
]
},
{
"id": 11,
"type": "Trellis2ExportMesh",
"pos": [
963.3009624015439,
1226.3665857311485
],
"size": [
324,
146
],
"flags": {},
"order": 7,
"mode": 0,
"inputs": [
{
"name": "trimesh",
"type": "TRIMESH",
"link": 83
}
],
"outputs": [
{
"name": "glb_path",
"type": "STRING",
"links": [
12
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "20651320d20ed56a52168c1bd3f29d6886c4be06",
"Node name for S&R": "Trellis2ExportMesh",
"widget_ue_connectable": {}
},
"widgets_values": [
"TexturedMeshMV",
"glb",
true
]
},
{
"id": 26,
"type": "Trellis2PreProcessImage",
"pos": [
273.345930367051,
1625.0892297322023
],
"size": [
386.1875,
107.328125
],
"flags": {
"collapsed": false
},
"order": 5,
"mode": 0,
"inputs": [
{
"name": "image",
"type": "IMAGE",
"link": 47
}
],
"outputs": [
{
"name": "image",
"type": "IMAGE",
"links": [
81
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "15854282a73cf231b81d52ada22e652f414a078b",
"Node name for S&R": "Trellis2PreProcessImage",
"widget_ue_connectable": {}
},
"widgets_values": [
25,
false
]
},
{
"id": 8,
"type": "Trellis2LoadMesh",
"pos": [
-288.7042468259506,
1215.4568912264674
],
"size": [
514.34375,
83.328125
],
"flags": {},
"order": 3,
"mode": 0,
"inputs": [],
"outputs": [
{
"name": "trimesh",
"type": "TRIMESH",
"links": [
80
]
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "20651320d20ed56a52168c1bd3f29d6886c4be06",
"Node name for S&R": "Trellis2LoadMesh",
"widget_ue_connectable": {}
},
"widgets_values": [
"C:\\Git\\ComfyUI\\output\\DwarfMV_Textured_00001_.glb"
]
},
{
"id": 37,
"type": "Trellis2MeshTexturingMultiView",
"pos": [
742.5264619145277,
1433.3435382751823
],
"size": [
539.515625,
675.328125
],
"flags": {},
"order": 6,
"mode": 0,
"inputs": [
{
"name": "pipeline",
"type": "TRELLIS2PIPELINE",
"link": 78
},
{
"name": "front_image",
"type": "IMAGE",
"link": 79
},
{
"name": "trimesh",
"type": "TRIMESH",
"link": 80
},
{
"name": "back_image",
"shape": 7,
"type": "IMAGE",
"link": 81
},
{
"name": "left_image",
"shape": 7,
"type": "IMAGE",
"link": null
},
{
"name": "right_image",
"shape": 7,
"type": "IMAGE",
"link": null
}
],
"outputs": [
{
"name": "trimesh",
"type": "TRIMESH",
"links": [
83
]
},
{
"name": "base_color_texture",
"type": "IMAGE",
"links": null
},
{
"name": "metallic_roughness_texture",
"type": "IMAGE",
"links": null
}
],
"properties": {
"aux_id": "visualbruno/ComfyUI-Trellis2",
"ver": "a19110a28a5c5c434386af0421005dd8edae82db",
"Node name for S&R": "Trellis2MeshTexturingMultiView",
"widget_ue_connectable": {}
},
"widgets_values": [
12345,
"fixed",
25,
3,
0.2,
3,
1024,
4096,
"OPAQUE",
false,
0,
0.9,
false,
false,
60,
"z",
2
]
}
],
"links": [
[
12,
11,
0,
12,
2,
"STRING"
],
[
32,
1,
2,
18,
0,
"IMAGE"
],
[
47,
21,
2,
26,
0,
"IMAGE"
],
[
78,
2,
0,
37,
0,
"TRELLIS2PIPELINE"
],
[
79,
18,
0,
37,
1,
"IMAGE"
],
[
80,
8,
0,
37,
2,
"TRIMESH"
],
[
81,
26,
0,
37,
3,
"IMAGE"
],
[
83,
37,
0,
11,
0,
"TRIMESH"
]
],
"groups": [],
"config": {},
"extra": {
"workflowRendererVersion": "Vue",
"ue_links": [],
"ds": {
"scale": 0.4305676431342665,
"offset": [
903.5556585357364,
-713.9325537000318
]
},
"links_added_by_ue": [],
"frontendVersion": "1.38.13",
"VHS_latentpreview": false,
"VHS_latentpreviewrate": 0,
"VHS_MetadataImage": true,
"VHS_KeepIntermediate": true
},
"version": 0.4
}
nodes.py
+237
View File
@@ -1253,6 +1253,140 @@ class Trellis2MeshWithVoxelAdvancedGenerator:
return (mesh,bvh,)
class Trellis2MeshWithVoxelMultiViewGenerator:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"pipeline": ("TRELLIS2PIPELINE",),
"front_image": ("IMAGE",),
"seed": ("INT", {"default": 12345, "min": 0, "max": 0x7fffffff}),
"pipeline_type": (["512","1024","1024_cascade","1536_cascade"],{"default":"1024_cascade"}),
"sparse_structure_steps": ("INT",{"default":12, "min":1, "max":100},),
"sparse_structure_guidance_strength": ("FLOAT",{"default":6.50}),
"sparse_structure_guidance_rescale": ("FLOAT",{"default":0.20}),
"sparse_structure_rescale_t": ("FLOAT",{"default":4.00}),
"shape_steps": ("INT",{"default":12, "min":1, "max":100},),
"shape_guidance_strength": ("FLOAT",{"default":6.50}),
"shape_guidance_rescale": ("FLOAT",{"default":0.20}),
"shape_rescale_t": ("FLOAT",{"default":4.00}),
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
"texture_guidance_strength": ("FLOAT",{"default":3.00}),
"texture_guidance_rescale": ("FLOAT",{"default":0.20}),
"texture_rescale_t": ("FLOAT",{"default":3.00}),
"max_num_tokens": ("INT",{"default":999999,"min":0,"max":999999}),
"sparse_structure_resolution": ("INT", {"default":32,"min":8,"max":128,"step":8}),
"generate_texture_slat": ("BOOLEAN", {"default":True}),
"sparse_structure_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
"sparse_structure_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
"shape_guidance_interval_start": ("FLOAT",{"default":0.10,"min":0.00,"max":1.00,"step":0.01}),
"shape_guidance_interval_end": ("FLOAT",{"default":1.00,"min":0.00,"max":1.00,"step":0.01}),
"texture_guidance_interval_start": ("FLOAT",{"default":0.00,"min":0.00,"max":1.00,"step":0.01}),
"texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}),
"use_tiled_decoder": ("BOOLEAN", {"default":True}),
"front_axis": (["z", "x"], {"default": "z"}),
"blend_temperature": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 10.0, "step": 0.1}),
},
"optional": {
"back_image": ("IMAGE",),
"left_image": ("IMAGE",),
"right_image": ("IMAGE",),
},
}
RETURN_TYPES = ("MESHWITHVOXEL","BVH", )
RETURN_NAMES = ("mesh", "bvh", )
FUNCTION = "process"
CATEGORY = "Trellis2Wrapper"
OUTPUT_NODE = True
def process(self, pipeline, front_image, seed, pipeline_type, sparse_structure_steps,
sparse_structure_guidance_strength,
sparse_structure_guidance_rescale,
sparse_structure_rescale_t,
shape_steps,
shape_guidance_strength,
shape_guidance_rescale,
shape_rescale_t,
texture_steps,
texture_guidance_strength,
texture_guidance_rescale,
texture_rescale_t,
max_num_tokens,
sparse_structure_resolution,
generate_texture_slat,
sparse_structure_guidance_interval_start,
sparse_structure_guidance_interval_end,
shape_guidance_interval_start,
shape_guidance_interval_end,
texture_guidance_interval_start,
texture_guidance_interval_end,
use_tiled_decoder,
front_axis,
blend_temperature,
back_image=None,
left_image=None,
right_image=None):
reset_cuda()
# Convert front image tensor to PIL
front_pil = tensor2pil(front_image)
# Convert optional view image tensors to PIL
back_pil = tensor2pil(back_image) if back_image is not None else None
left_pil = tensor2pil(left_image) if left_image is not None else None
right_pil = tensor2pil(right_image) if right_image is not None else None
sparse_structure_guidance_interval = [sparse_structure_guidance_interval_start,sparse_structure_guidance_interval_end]
shape_guidance_interval = [shape_guidance_interval_start,shape_guidance_interval_end]
texture_guidance_interval = [texture_guidance_interval_start,texture_guidance_interval_end]
sparse_structure_sampler_params = {"steps":sparse_structure_steps,"guidance_strength":sparse_structure_guidance_strength,"guidance_rescale":sparse_structure_guidance_rescale,"guidance_interval":sparse_structure_guidance_interval,"rescale_t":sparse_structure_rescale_t}
shape_slat_sampler_params = {"steps":shape_steps,"guidance_strength":shape_guidance_strength,"guidance_rescale":shape_guidance_rescale,"guidance_interval":shape_guidance_interval,"rescale_t":shape_rescale_t}
tex_slat_sampler_params = {"steps":texture_steps,"guidance_strength":texture_guidance_strength,"guidance_rescale":texture_guidance_rescale,"guidance_interval":texture_guidance_interval,"rescale_t":texture_rescale_t}
if generate_texture_slat:
num_steps = 5
else:
num_steps = 4
pbar = ProgressBar(num_steps)
mesh = pipeline.run_multiview(
front=front_pil,
back=back_pil,
left=left_pil,
right=right_pil,
seed=seed,
pipeline_type=pipeline_type,
sparse_structure_sampler_params=sparse_structure_sampler_params,
shape_slat_sampler_params=shape_slat_sampler_params,
tex_slat_sampler_params=tex_slat_sampler_params,
max_num_tokens=max_num_tokens,
sparse_structure_resolution=sparse_structure_resolution,
generate_texture_slat=generate_texture_slat,
use_tiled=use_tiled_decoder,
pbar=pbar,
front_axis=front_axis,
blend_temperature=blend_temperature,
)[0]
vertices = mesh.vertices.cuda()
faces = mesh.faces.cuda()
if generate_texture_slat:
# Build BVH for the current mesh to guide remeshing
print("Building BVH for current mesh...")
bvh = CuMesh.cuBVH(vertices.detach().clone(), faces.detach().clone())
bvh.vertices = vertices.detach().clone()
bvh.faces = faces.detach().clone()
else:
print("Not building BVH : only used for texturing")
bvh = None
return (mesh,bvh,)
class Trellis2PostProcessAndUnWrapAndRasterizer:
@classmethod
def INPUT_TYPES(s):
@@ -1922,6 +2056,105 @@ class Trellis2MeshTexturing:
return (textured_mesh, baseColorTexture, metallicRoughnessTexture, )
class Trellis2MeshTexturingMultiView:
@classmethod
def INPUT_TYPES(s):
return {
"required": {
"pipeline": ("TRELLIS2PIPELINE",),
"front_image": ("IMAGE",),
"trimesh": ("TRIMESH",),
"seed": ("INT", {"default": 0, "min": 0, "max": 0x7fffffff}),
"texture_steps": ("INT",{"default":12, "min":1, "max":100},),
"texture_guidance_strength": ("FLOAT",{"default":3.0}),
"texture_guidance_rescale": ("FLOAT",{"default":0.2}),
"texture_rescale_t": ("FLOAT",{"default":3.0}),
"resolution": ([512,1024],{"default":1024}),
"texture_size": ("INT",{"default":4096,"min":512,"max":16384}),
"texture_alpha_mode": (["OPAQUE","MASK","BLEND"],{"default":"OPAQUE"}),
"double_side_material": ("BOOLEAN",{"default":False}),
"texture_guidance_interval_start": ("FLOAT",{"default":0.00,"min":0.00,"max":1.00,"step":0.01}),
"texture_guidance_interval_end": ("FLOAT",{"default":0.90,"min":0.00,"max":1.00,"step":0.01}),
"bake_on_vertices": ("BOOLEAN",{"default":False}),
"use_custom_normals": ("BOOLEAN",{"default":False}),
"mesh_cluster_threshold_cone_half_angle_rad": ("FLOAT",{"default":60.0,"min":0.0,"max":359.9}),
"front_axis": (["z", "x"], {"default": "z"}),
"blend_temperature": ("FLOAT", {"default": 2.0, "min": 0.1, "max": 10.0, "step": 0.1}),
},
"optional": {
"back_image": ("IMAGE",),
"left_image": ("IMAGE",),
"right_image": ("IMAGE",),
}
}
RETURN_TYPES = ("TRIMESH","IMAGE","IMAGE",)
RETURN_NAMES = ("trimesh","base_color_texture","metallic_roughness_texture",)
FUNCTION = "process"
CATEGORY = "Trellis2Wrapper"
OUTPUT_NODE = True
def process(self,
pipeline,
front_image,
trimesh,
seed,
texture_steps,
texture_guidance_strength,
texture_guidance_rescale,
texture_rescale_t,
resolution,
texture_size,
texture_alpha_mode,
double_side_material,
texture_guidance_interval_start,
texture_guidance_interval_end,
bake_on_vertices,
use_custom_normals,
mesh_cluster_threshold_cone_half_angle_rad,
front_axis,
blend_temperature,
back_image = None,
left_image = None,
right_image = None):
reset_cuda()
# Convert front image tensor to PIL
front_pil = tensor2pil(front_image)
# Convert optional view image tensors to PIL
back_pil = tensor2pil(back_image) if back_image is not None else None
left_pil = tensor2pil(left_image) if left_image is not None else None
right_pil = tensor2pil(right_image) if right_image is not None else None
texture_guidance_interval = [texture_guidance_interval_start,texture_guidance_interval_end]
tex_slat_sampler_params = {"steps":texture_steps,"guidance_strength":texture_guidance_strength,"guidance_rescale":texture_guidance_rescale,"guidance_interval":texture_guidance_interval,"rescale_t":texture_rescale_t}
textured_mesh, baseColorTexture_np, metallicRoughnessTexture_np = pipeline.texture_mesh_multiview(mesh=trimesh,
front=front_pil,
back=back_pil,
left=left_pil,
right=right_pil,
seed=seed,
tex_slat_sampler_params = tex_slat_sampler_params,
resolution = resolution,
texture_size = texture_size,
texture_alpha_mode = texture_alpha_mode,
double_side_material = double_side_material,
bake_on_vertices = bake_on_vertices,
use_custom_normals = use_custom_normals,
mesh_cluster_threshold_cone_half_angle_rad = mesh_cluster_threshold_cone_half_angle_rad,
front_axis = front_axis,
blend_temperature = blend_temperature
)
baseColorTexture = pil2tensor(baseColorTexture_np)
metallicRoughnessTexture = pil2tensor(metallicRoughnessTexture_np)
return (textured_mesh, baseColorTexture, metallicRoughnessTexture, )
class Trellis2LoadMesh:
@classmethod
def INPUT_TYPES(s):
@@ -2680,6 +2913,8 @@ NODE_CLASS_MAPPINGS = {
"Trellis2SmoothNormals": Trellis2SmoothNormals,
"Trellis2RemeshWithQuad": Trellis2RemeshWithQuad,
"Trellis2BatchSimplifyMeshAndExport": Trellis2BatchSimplifyMeshAndExport,
"Trellis2MeshWithVoxelMultiViewGenerator": Trellis2MeshWithVoxelMultiViewGenerator,
"Trellis2MeshTexturingMultiView": Trellis2MeshTexturingMultiView,
}
@@ -2711,4 +2946,6 @@ NODE_DISPLAY_NAME_MAPPINGS = {
"Trellis2SmoothNormals": "Trellis2 - Smooth Normals",
"Trellis2RemeshWithQuad": "Trellis2 - Remesh With Quad",
"Trellis2BatchSimplifyMeshAndExport": "Trellis2 - Batch Simplify Mesh And Export",
"Trellis2MeshWithVoxelMultiViewGenerator": "Trellis2 - Mesh With Voxel Multi-View Generator",
"Trellis2MeshTexturingMultiView": "Trellis2 - Mesh Texturing Multi-View",
}
pyproject.toml
+1 -1
View File
@@ -1,7 +1,7 @@
[project]
name = "trellis2"
description = "ComfyUI Wrapper for Microsoft Trellis.2 - Native and Compact Structured Latents for 3D Generation"
version = "1.0.7"
version = "1.0.8"
license = {file = "LICENSE"}
# classifiers = [
# # For OS-independent nodes (works on all operating systems)
trellis2/pipelines/samplers/__init__.py
+2
View File
@@ -3,4 +3,6 @@ from .flow_euler import (
FlowEulerSampler,
FlowEulerCfgSampler,
FlowEulerGuidanceIntervalSampler,
FlowEulerMultiViewSampler,
FlowEulerMultiViewGuidanceIntervalSampler,
)
trellis2/pipelines/samplers/flow_euler.py
+212
View File
@@ -206,3 +206,215 @@ class FlowEulerGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, ClassifierF
- 'pred_x_0': a list of prediction of x_0.
"""
return super().sample(model, noise, cond, steps, rescale_t, verbose, neg_cond=neg_cond, guidance_strength=guidance_strength, guidance_interval=guidance_interval, **kwargs)
class FlowEulerMultiViewSampler(FlowEulerSampler):
"""
Generate samples from a flow-matching model using Euler sampling with multi-view blending.
"""
def __init__(self, sigma_min: float, resolution: int):
super().__init__(sigma_min)
self.resolution = resolution
def _compute_view_weights_sparse(self, coords, views, front_axis='z', blend_temperature=2.0) -> torch.Tensor:
"""
Compute blending weights for sparse voxels.
"""
# Normalize coords to [-1, 1] range (roughly)
z = (coords[:, 1].float() / self.resolution) * 2 - 1.0
x = (coords[:, 3].float() / self.resolution) * 2 - 1.0
if front_axis == 'z':
# Front (+Z), Back (-Z), Right (+X), Left (-X)
view_vectors = {
'front': torch.stack([torch.zeros_like(z), z], dim=1), # (0, z)
'back': torch.stack([torch.zeros_like(z), -z], dim=1),
'right': torch.stack([x, torch.zeros_like(x)], dim=1),
'left': torch.stack([-x, torch.zeros_like(x)], dim=1),
}
else: # front_axis == 'x' (swapped)
# Front (+X), Back (-X), Right (+Z), Left (-Z)
view_vectors = {
'front': torch.stack([x, torch.zeros_like(x)], dim=1),
'back': torch.stack([-x, torch.zeros_like(x)], dim=1),
'right': torch.stack([torch.zeros_like(z), z], dim=1),
'left': torch.stack([torch.zeros_like(z), -z], dim=1),
}
scores = []
for view in views:
if view in view_vectors:
v_vec = view_vectors[view]
score = v_vec.sum(dim=1)
scores.append(score)
else:
scores.append(torch.full_like(z, -10.0))
scores = torch.stack(scores, dim=1) # (N, num_views)
weights = torch.softmax(scores * blend_temperature, dim=1)
return weights
def _compute_view_weights_dense(self, shape, device, views, front_axis='z', blend_temperature=2.0) -> torch.Tensor:
"""
Compute blending weights for dense grid (B, C, D, H, W).
Returns weights of shape (1, 1, D, H, W, NumViews) for easy broadcasting (actually we want (1, 1, D, H, W) per view)
"""
# shape is (B, C, D, H, W)
D, H, W = shape[2], shape[3], shape[4]
# Create meshgrid in [-1, 1]
# We assume D is Z axis, W is X axis (usually D, H, W = Z, Y, X in 3D tensors?)
# Let's verify standard: (Batch, Channel, Depth, Height, Width) -> (B, C, Z, Y, X)
dz = torch.linspace(-1, 1, D, device=device)
dy = torch.linspace(-1, 1, H, device=device)
dx = torch.linspace(-1, 1, W, device=device)
# meshgrid 'ij' indexing: (D, H, W) order
grid_z, grid_y, grid_x = torch.meshgrid(dz, dy, dx, indexing='ij')
# Flatten for vector calc? Or keep structural. Keep structural.
if front_axis == 'z':
# Front (+Z), Back (-Z), Right (+X), Left (-X)
# Vectors are scalar fields here
view_scores = {
'front': grid_z,
'back': -grid_z,
'right': grid_x,
'left': -grid_x,
}
else:
view_scores = {
'front': grid_x,
'back': -grid_x,
'right': grid_z,
'left': -grid_z,
}
scores = []
for view in views:
if view in view_scores:
scores.append(view_scores[view])
else:
scores.append(torch.full_like(grid_z, -10.0))
# Stack: (NumViews, D, H, W)
scores = torch.stack(scores, dim=0)
# Softmax over views dimension (0)
weights = torch.softmax(scores * blend_temperature, dim=0)
# Reshape for broadcasting: (NumViews, 1, 1, D, H, W) -> No wait, loop is over views.
# We want to return something we can index like weights[i] -> (1, 1, D, H, W)
# Current shape: (NumViews, D, H, W)
return weights
@torch.no_grad()
def sample_once(
self,
model,
x_t,
t: float,
t_prev: float,
conds: Dict[str, Any], # Changed: expects dict of {view: cond}
views: List[str], # Changed: list of view keys corresponding to conds
front_axis: str = 'z',
blend_temperature: float = 2.0,
**kwargs
):
"""
Sample with multi-view blending.
"""
is_sparse = hasattr(x_t, 'coords')
if is_sparse:
# 1. Compute per-voxel weights based on current sparse coords
weights = self._compute_view_weights_sparse(x_t.coords, views, front_axis, blend_temperature)
# weights: (N, NumViews)
else:
# Dense tensor (B, C, D, H, W)
weights = self._compute_view_weights_dense(x_t.shape, x_t.device, views, front_axis, blend_temperature)
# weights: (NumViews, D, H, W)
# 2. Run model for each view and blend predictions
pred_v_accum = 0
for i, view in enumerate(views):
cond = conds[view]
# Use _inference_model to support mixins (CFG, etc)
# If cond is a dict containing 'cond' and 'neg_cond' (from pipeline.get_cond), unpack it
if isinstance(cond, dict) and 'cond' in cond and 'neg_cond' in cond:
pred_v_view = self._inference_model(model, x_t, t, cond=cond['cond'], neg_cond=cond['neg_cond'], **kwargs)
else:
pred_v_view = self._inference_model(model, x_t, t, cond=cond, **kwargs)
# Weighted accumulation
if is_sparse:
# weights[:, i] is (N,), pred_v_view might be SparseTensor or Tensor (N, C)
w = weights[:, i].unsqueeze(1)
v_feats = pred_v_view.feats if hasattr(pred_v_view, 'feats') else pred_v_view
pred_v_accum += v_feats * w
else:
# Dense
# weights[i] is (D, H, W). pred_v_view is (B, C, D, H, W)
w = weights[i].unsqueeze(0).unsqueeze(0) # (1, 1, D, H, W)
pred_v_accum += pred_v_view * w
if is_sparse:
# Re-wrap accumulated features into a SparseTensor matching x_t
# pred_v_accum is (N, C) tensor now
pred_v = x_t.replace(feats=pred_v_accum)
else:
pred_v = pred_v_accum
pred_x_0, pred_eps = self._v_to_xstart_eps(x_t=x_t, t=t, v=pred_v)
pred_x_prev = x_t - (t - t_prev) * pred_v
return edict({"pred_x_prev": pred_x_prev, "pred_x_0": pred_x_0})
@torch.no_grad()
def sample(
self,
model,
noise,
conds: Dict[str, Any], # {view: cond}
views: List[str], # ['front', 'back', ...]
steps: int = 50,
rescale_t: float = 1.0,
verbose: bool = True,
tqdm_desc: str = "Sampling MultiView",
front_axis: str = 'z',
blend_temperature: float = 2.0,
**kwargs
):
sample = noise
t_seq = np.linspace(1, 0, steps + 1)
t_seq = rescale_t * t_seq / (1 + (rescale_t - 1) * t_seq)
t_seq = t_seq.tolist()
t_pairs = list((t_seq[i], t_seq[i + 1]) for i in range(steps))
ret = edict({"samples": None, "pred_x_t": [], "pred_x_0": []})
for t, t_prev in tqdm(t_pairs, desc=tqdm_desc, disable=not verbose):
out = self.sample_once(
model, sample, t, t_prev,
conds=conds,
views=views,
front_axis=front_axis,
blend_temperature=blend_temperature,
**kwargs
)
sample = out.pred_x_prev
ret.pred_x_t.append(out.pred_x_prev)
ret.pred_x_0.append(out.pred_x_0)
ret.samples = sample
return ret
class FlowEulerMultiViewGuidanceIntervalSampler(GuidanceIntervalSamplerMixin, ClassifierFreeGuidanceSamplerMixin, FlowEulerMultiViewSampler):
"""
Generate samples from a flow-matching model using Euler sampling with multi-view blending, CFG, and guidance interval.
"""
pass
trellis2/pipelines/trellis2_image_to_3d.py
+645 -2
View File
@@ -1130,6 +1130,559 @@ class Trellis2ImageTo3DPipeline(Pipeline):
else:
return out_mesh
@torch.no_grad()
def run_multiview(
self,
front: Image.Image,
back: Image.Image = None,
left: Image.Image = None,
right: Image.Image = None,
seed: int = 42,
pipeline_type: str = None,
sparse_structure_sampler_params: dict = {},
shape_slat_sampler_params: dict = {},
tex_slat_sampler_params: dict = {},
max_num_tokens: int = 49152,
sparse_structure_resolution: int = 32,
generate_texture_slat: bool = True,
use_tiled: bool = True,
return_latent: bool = False,
pbar: ProgressBar = None,
front_axis: str = 'z',
blend_temperature: float = 2.0,
) -> List[MeshWithVoxel]:
"""
Run the pipeline with named multi-view images and spatial blending.
"""
if pipeline_type is None:
pipeline_type = self.default_pipeline_type
torch.manual_seed(seed)
# Collect views
views_dict = {'front': front}
if back is not None: views_dict['back'] = back
if left is not None: views_dict['left'] = left
if right is not None: views_dict['right'] = right
views_list = list(views_dict.keys())
# 1. Conditioning
# Calculate conditioning per view
conds = {} # 1024 or None (if 512)
lr_conds = {} # 512 (for cascade)
conds_512 = {} # Explicit 512 storage for structure sampling
conds_1024 = {}
self.load_image_cond_model()
if pipeline_type == '512':
for v, img in views_dict.items():
c = self.get_cond([img], 512)
conds[v] = c
conds_512[v] = c
elif pipeline_type == '1024':
for v, img in views_dict.items():
c1024 = self.get_cond([img], 1024)
conds[v] = c1024
conds_1024[v] = c1024
# Does 1024 pipeline use 512 for structure?
# run() says: cond_512 = get_cond(..., 512). So yes.
conds_512[v] = self.get_cond([img], 512)
elif 'cascade' in pipeline_type:
# 1024_cascade or 1536_cascade
for v, img in views_dict.items():
c512 = self.get_cond([img], 512)
c1024 = self.get_cond([img], 1024)
lr_conds[v] = c512
conds[v] = c1024
conds_512[v] = c512
conds_1024[v] = c1024
if not self.keep_models_loaded:
self.unload_image_cond_model()
if pbar is not None:
pbar.update(1)
self.load_sparse_structure_model()
coords = self.sample_sparse_structure_multiview(
conds_512,
views_list,
sparse_structure_resolution,
sampler_params=sparse_structure_sampler_params,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
if not self.keep_models_loaded:
self.unload_sparse_structure_model()
if pbar is not None:
pbar.update(1)
# 3. Shape Slat MultiView
shape_slat = None
res = 0
if pipeline_type == '1024_cascade':
self.load_shape_slat_flow_model_512()
self.load_shape_slat_flow_model_1024()
shape_slat = self.sample_shape_slat_cascade_multiview(
lr_conds, conds, views_list,
self.models['shape_slat_flow_model_512'], self.models['shape_slat_flow_model_1024'],
512, 1024,
coords, shape_slat_sampler_params,
max_num_tokens,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
res = 1024
if not self.keep_models_loaded:
self.unload_shape_slat_flow_model_512()
self.unload_shape_slat_flow_model_1024()
elif pipeline_type == '1536_cascade':
self.load_shape_slat_flow_model_512()
self.load_shape_slat_flow_model_1024()
shape_slat = self.sample_shape_slat_cascade_multiview(
lr_conds, conds, views_list,
self.models['shape_slat_flow_model_512'], self.models['shape_slat_flow_model_1024'],
512, 1536,
coords, shape_slat_sampler_params,
max_num_tokens,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
res = 1536
if not self.keep_models_loaded:
self.unload_shape_slat_flow_model_512()
self.unload_shape_slat_flow_model_1024()
elif pipeline_type == '512': # Single stage
self.load_shape_slat_flow_model_512()
shape_slat = self.sample_shape_slat_multiview(
conds, views_list,
self.models['shape_slat_flow_model_512'],
coords, shape_slat_sampler_params,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
res = 512
if not self.keep_models_loaded:
self.unload_shape_slat_flow_model_512()
elif pipeline_type == '1024': # Single stage
self.load_shape_slat_flow_model_1024()
shape_slat = self.sample_shape_slat_multiview(
conds, views_list,
self.models['shape_slat_flow_model_1024'],
coords, shape_slat_sampler_params,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
res = 1024
if not self.keep_models_loaded:
self.unload_shape_slat_flow_model_1024()
if pbar is not None:
pbar.update(1)
# Texture Slat MultiView
tex_slat = None
if generate_texture_slat:
tex_model_key = 'tex_slat_flow_model_1024'
if pipeline_type == '512':
tex_model_key = 'tex_slat_flow_model_512'
self.load_tex_slat_flow_model_512()
flow_model = self.models['tex_slat_flow_model_512']
tex_conds = conds_512
else:
self.load_tex_slat_flow_model_1024()
flow_model = self.models['tex_slat_flow_model_1024']
tex_conds = conds_1024
tex_slat = self.sample_tex_slat_multiview(
tex_conds, views_list,
shape_slat=shape_slat,
flow_model=flow_model,
sampler_params=tex_slat_sampler_params,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
if not self.keep_models_loaded:
if pipeline_type == '512':
self.unload_tex_slat_flow_model_512()
else:
self.unload_tex_slat_flow_model_1024()
if pbar is not None:
pbar.update(1)
torch.cuda.empty_cache()
if generate_texture_slat:
out_mesh = self.decode_latent(shape_slat, tex_slat, res, use_tiled=use_tiled)
else:
out_mesh = self.decode_latent(shape_slat, None, res, use_tiled=use_tiled)
torch.cuda.empty_cache()
if pbar is not None:
pbar.update(1)
if return_latent:
return out_mesh, (shape_slat, tex_slat, res)
else:
return out_mesh
def sample_sparse_structure_multiview(
self,
conds: dict,
views: list,
resolution: int,
num_samples: int = 1,
sampler_params: dict = {},
front_axis: str = 'z',
blend_temperature: float = 2.0,
) -> torch.Tensor:
"""
Sample sparse structures with multi-view blending.
"""
if self.low_vram:
for v in conds:
conds[v] = self._cond_to(conds[v], self.device)
# Sample sparse structure latent
flow_model = self.models['sparse_structure_flow_model']
reso = flow_model.resolution
in_channels = flow_model.in_channels
noise = torch.randn(num_samples, in_channels, reso, reso, reso).to(self.device)
sampler = samplers.FlowEulerMultiViewGuidanceIntervalSampler(
sigma_min=1e-5,
resolution=flow_model.resolution
)
sampler_params = {**self.sparse_structure_sampler_params, **sampler_params}
if self.low_vram:
flow_model.to(self.device)
z_s = sampler.sample(
flow_model,
noise,
conds=conds,
views=views,
front_axis=front_axis,
blend_temperature=blend_temperature,
**sampler_params,
verbose=True,
tqdm_desc="Sampling sparse structure (MultiView)",
).samples
if self.low_vram:
flow_model.cpu()
self._cleanup_cuda()
# Decode sparse structure latent
decoder = self.models['sparse_structure_decoder']
if self.low_vram:
decoder.to(self.device)
# Standard decoding logic from sample_sparse_structure
decoded = decoder(z_s) > 0
if self.low_vram:
decoder.cpu()
self._cleanup_cuda()
if resolution != decoded.shape[2]:
ratio = decoded.shape[2] // resolution
decoded = torch.nn.functional.max_pool3d(decoded.float(), ratio, ratio, 0) > 0.5
# Extract coordinates (N, 4) -> (b, d, h, w)
# argwhere returns (b, c, d, h, w), so we want [0, 2, 3, 4]
coords = torch.argwhere(decoded)[:, [0, 2, 3, 4]].int()
coords = coords.cpu()
del decoded
del z_s
if self.low_vram:
for v in conds:
conds[v] = self._cond_cpu(conds[v])
self._cleanup_cuda()
return coords
def sample_shape_slat_multiview(
self,
conds: dict,
views: list,
flow_model,
coords: torch.Tensor,
sampler_params: dict = {},
front_axis: str = 'z',
blend_temperature: float = 2.0,
) -> SparseTensor:
if self.low_vram:
for v in conds:
conds[v] = self._cond_to(conds[v], self.device)
coords_dev = coords.to(self.device)
noise = SparseTensor(
feats=torch.randn(coords.shape[0], flow_model.in_channels, device=self.device),
coords=coords_dev,
)
sampler = samplers.FlowEulerMultiViewGuidanceIntervalSampler(
sigma_min=1e-5,
resolution=flow_model.resolution,
)
sampler_params = {**self.shape_slat_sampler_params, **sampler_params}
if self.low_vram:
flow_model.to(self.device)
slat = sampler.sample(
flow_model,
noise,
conds=conds,
views=views,
front_axis=front_axis,
blend_temperature=blend_temperature,
**sampler_params,
verbose=True,
tqdm_desc="Sampling shape SLat (MultiView)",
).samples
if self.low_vram:
flow_model.cpu()
self._cleanup_cuda()
std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
slat = slat * std + mean
del coords_dev
if self.low_vram:
for v in conds:
conds[v] = self._cond_cpu(conds[v])
self._cleanup_cuda()
return slat
def sample_shape_slat_cascade_multiview(
self,
lr_conds: dict,
conds: dict,
views: list,
flow_model_lr,
flow_model,
lr_resolution: int,
resolution: int,
coords: torch.Tensor,
sampler_params: dict = {},
max_num_tokens: int = 49152,
front_axis: str = 'z',
blend_temperature: float = 2.0,
) -> SparseTensor:
# LR
if self.low_vram:
for v in lr_conds:
lr_conds[v] = self._cond_to(lr_conds[v], self.device)
for v in conds:
conds[v] = self._cond_to(conds[v], self.device)
coords_dev = coords.to(self.device)
noise = SparseTensor(
feats=torch.randn(coords.shape[0], flow_model_lr.in_channels, device=self.device),
coords=coords_dev,
)
sampler_lr = samplers.FlowEulerMultiViewGuidanceIntervalSampler(
sigma_min=1e-5,
resolution=flow_model_lr.resolution,
)
sampler_params_combined = {**self.shape_slat_sampler_params, **sampler_params}
if self.low_vram:
flow_model_lr.to(self.device)
slat = sampler_lr.sample(
flow_model_lr,
noise,
conds=lr_conds,
views=views,
front_axis=front_axis,
blend_temperature=blend_temperature,
**sampler_params_combined,
verbose=True,
tqdm_desc="Sampling shape SLat (MultiView LR)",
).samples
if self.low_vram:
flow_model_lr.cpu()
self._cleanup_cuda()
std = torch.tensor(self.shape_slat_normalization['std'])[None].to(slat.device)
mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(slat.device)
slat = slat * std + mean
del coords_dev
# Upsample logic
self.load_shape_slat_decoder()
if self.low_vram:
self.models['shape_slat_decoder'].to(self.device)
self.models['shape_slat_decoder'].low_vram = True
hr_coords = self.models['shape_slat_decoder'].upsample(slat, upsample_times=4)
if self.low_vram:
self.models['shape_slat_decoder'].cpu()
self.models['shape_slat_decoder'].low_vram = False
hr_resolution = resolution
while True:
quant_coords = torch.cat([
hr_coords[:, :1],
((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(),
], dim=1)
coords = quant_coords.unique(dim=0)
num_tokens = coords.shape[0]
if num_tokens < max_num_tokens:
if hr_resolution != resolution:
print(f"Due to the limited number of tokens, the resolution is reduced to {hr_resolution}.")
break
hr_resolution -= 128
if hr_resolution < 1024 and resolution >= 1024:
hr_resolution = 1024
break
if hr_resolution < 512:
hr_resolution = 512
break
# HR
sampler_hr = samplers.FlowEulerMultiViewGuidanceIntervalSampler(
sigma_min=1e-5,
resolution=flow_model.resolution,
)
coords_dev = coords.to(self.device).contiguous()
noise = SparseTensor(
feats=torch.randn(coords_dev.shape[0], flow_model.in_channels, device=self.device),
coords=coords_dev,
)
if self.low_vram:
flow_model.to(self.device)
d_slat = sampler_hr.sample(
flow_model,
noise,
conds=conds,
views=views,
front_axis=front_axis,
blend_temperature=blend_temperature,
**sampler_params_combined,
verbose=True,
tqdm_desc="Sampling shape SLat (MultiView HR)",
).samples
if self.low_vram:
flow_model.cpu()
self._cleanup_cuda()
slat = d_slat * std + mean
if self.low_vram:
for v in lr_conds:
lr_conds[v] = self._cond_cpu(lr_conds[v])
for v in conds:
conds[v] = self._cond_cpu(conds[v])
self._cleanup_cuda()
return slat
def sample_tex_slat_multiview(
self,
conds: dict,
views: list,
shape_slat: SparseTensor,
flow_model,
sampler_params: dict = {},
front_axis: str = 'z',
blend_temperature: float = 2.0,
) -> SparseTensor:
"""
Sample structured latent for texture with multi-view blending.
"""
if self.low_vram:
for v in conds:
conds[v] = self._cond_to(conds[v], self.device)
# Normalize shape slat for conditioning
std = torch.tensor(self.shape_slat_normalization['std'])[None].to(shape_slat.device)
mean = torch.tensor(self.shape_slat_normalization['mean'])[None].to(shape_slat.device)
shape_slat_normalized = (shape_slat - mean) / std
#coords = shape_slat.coords
#coords_dev = coords.to(self.device)
# Calculate noise channels: total input - concat cond channels
in_channels = flow_model.in_channels if isinstance(flow_model, nn.Module) else flow_model[0].in_channels
noise_channels = in_channels - shape_slat.feats.shape[1]
# noise = SparseTensor(
# feats=torch.randn(coords.shape[0], noise_channels, device=self.device),
# coords=coords_dev,
# )
noise = shape_slat.replace(feats=torch.randn(shape_slat.coords.shape[0], in_channels - shape_slat.feats.shape[1]).to(self.device))
sampler_params = {**self.tex_slat_sampler_params, **sampler_params}
sampler = samplers.FlowEulerMultiViewGuidanceIntervalSampler(
sigma_min=1e-5,
resolution=flow_model.resolution,
)
if self.low_vram:
flow_model.to(self.device)
slat = sampler.sample(
flow_model,
noise,
conds=conds,
views=views,
front_axis=front_axis,
blend_temperature=blend_temperature,
concat_cond=shape_slat_normalized,
**sampler_params,
verbose=True,
tqdm_desc="Sampling texture SLat (MultiView)",
).samples
if self.low_vram:
flow_model.cpu()
self._cleanup_cuda()
std = torch.tensor(self.tex_slat_normalization['std'])[None].to(slat.device)
mean = torch.tensor(self.tex_slat_normalization['mean'])[None].to(slat.device)
slat = slat * std + mean
#del coords_dev
if self.low_vram:
for v in conds:
conds[v] = self._cond_cpu(conds[v])
self._cleanup_cuda()
return slat
def preprocess_mesh(self, mesh: trimesh.Trimesh) -> trimesh.Trimesh:
"""
Preprocess the input mesh.
@@ -1442,8 +1995,6 @@ class Trellis2ImageTo3DPipeline(Pipeline):
images = list(image)
else:
images = [image]
torch.manual_seed(seed)
self.load_image_cond_model()
cond = self.get_cond(images, resolution, max_views = max_views)
@@ -1484,6 +2035,98 @@ class Trellis2ImageTo3DPipeline(Pipeline):
out_mesh, baseColorTexture, metallicRoughnessTexture = self.postprocess_mesh(mesh, pbr_voxel, resolution, texture_size, texture_alpha_mode, double_side_material, bake_on_vertices, use_custom_normals, mesh_cluster_threshold_cone_half_angle_rad)
return out_mesh, baseColorTexture, metallicRoughnessTexture
@torch.no_grad()
def texture_mesh_multiview(
self,
mesh: trimesh.Trimesh,
front: Image.Image,
back: Image.Image,
left: Image.Image,
right: Image.Image,
seed: int = 42,
tex_slat_sampler_params: dict = {},
resolution: int = 1024,
texture_size: int = 2048,
texture_alpha_mode = 'OPAQUE',
double_side_material = True,
bake_on_vertices = False,
use_custom_normals = False,
mesh_cluster_threshold_cone_half_angle_rad=60.0,
front_axis: str = 'z',
blend_temperature: float = 2.0,
):
mesh = self.preprocess_mesh(mesh)
torch.manual_seed(seed)
self.load_image_cond_model()
# Collect views
views_dict = {'front': front}
if back is not None: views_dict['back'] = back
if left is not None: views_dict['left'] = left
if right is not None: views_dict['right'] = right
views_list = list(views_dict.keys())
# 1. Conditioning
# Calculate conditioning per view
conds = {}
self.load_image_cond_model()
if resolution == 512:
for v, img in views_dict.items():
c = self.get_cond([img], 512)
conds[v] = c
else:
for v, img in views_dict.items():
c = self.get_cond([img], 1024)
conds[v] = c
if not self.keep_models_loaded:
self.unload_image_cond_model()
shape_slat = self.encode_shape_slat(mesh, resolution)
if resolution==512:
self.unload_tex_slat_flow_model_1024()
self.load_tex_slat_flow_model_512()
tex_model = self.models['tex_slat_flow_model_512']
tex_slat = self.sample_tex_slat_multiview(
conds, views_list,
shape_slat=shape_slat,
flow_model=tex_model,
sampler_params=tex_slat_sampler_params,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
if not self.keep_models_loaded:
self.unload_tex_slat_flow_model_512()
else:
self.unload_tex_slat_flow_model_512()
self.load_tex_slat_flow_model_1024()
tex_model = self.models['tex_slat_flow_model_1024']
tex_slat = self.sample_tex_slat_multiview(
conds, views_list,
shape_slat=shape_slat,
flow_model=tex_model,
sampler_params=tex_slat_sampler_params,
front_axis=front_axis,
blend_temperature=blend_temperature,
)
if not self.keep_models_loaded:
self.unload_shape_slat_flow_model_1024()
torch.cuda.empty_cache()
pbr_voxel = self.decode_tex_slat(tex_slat)
torch.cuda.empty_cache()
out_mesh, baseColorTexture, metallicRoughnessTexture = self.postprocess_mesh(mesh, pbr_voxel, resolution, texture_size, texture_alpha_mode, double_side_material, bake_on_vertices, use_custom_normals, mesh_cluster_threshold_cone_half_angle_rad)
return out_mesh, baseColorTexture, metallicRoughnessTexture
def get_coords_from_trimesh(self, mesh, resolution):
vertices = torch.from_numpy(mesh.vertices).float()