wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-05-16 03:57:27 +08:00
Code Issues Actions 11 Packages Projects Releases Wiki Activity

shape_structure and tooltip

Browse Source
This commit is contained in:
Yousef Rafat 2026-05-14 20:37:58 +03:00
parent 0743ffb1be
commit 3d5f9aead7
1 changed files with 27 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

40
comfy_extras/nodes_trellis2.py
View File

@ -648,7 +648,14 @@ class EmptyTrellis2ShapeLatent(IO.ComfyNode):
category="latent/3d", category="latent/3d",
inputs=[ inputs=[
IO.Model.Input("model"), IO.Model.Input("model"),
IO.AnyType.Input("structure_or_coords"), IO.MultiType.Input(
"shape_structure",
types=[IO.Voxel, IO.AnyType],
tooltip=(
"Shape structure input. Accepts either a voxel structure "
"or upsampled coordinates from a previous cascade stage."
)
)
], ],
outputs=[ outputs=[
IO.Model.Output(), IO.Model.Output(),
@ -657,20 +664,20 @@ class EmptyTrellis2ShapeLatent(IO.ComfyNode):
) )
@classmethod @classmethod
def execute(cls, model, structure_or_coords): def execute(cls, model, shape_structure):
# to accept the upscaled coords # to accept the upscaled coords
is_512_pass = False is_512_pass = False
if hasattr(structure_or_coords, "data") and structure_or_coords.data.ndim == 4: if hasattr(shape_structure, "data") and shape_structure.data.ndim == 4:
decoded = structure_or_coords.data.unsqueeze(1) decoded = shape_structure.data.unsqueeze(1)
coords = torch.argwhere(decoded.bool())[:, [0, 2, 3, 4]].int() coords = torch.argwhere(decoded.bool())[:, [0, 2, 3, 4]].int()
is_512_pass = True is_512_pass = True
elif isinstance(structure_or_coords, torch.Tensor) and structure_or_coords.ndim == 2: elif isinstance(shape_structure, torch.Tensor) and shape_structure.ndim == 2:
coords = structure_or_coords.int() coords = shape_structure.int()
is_512_pass = False is_512_pass = False
else: else:
raise ValueError(f"Invalid input to EmptyTrellis2ShapeLatent: {type(structure_or_coords)}") raise ValueError(f"Invalid input to EmptyTrellis2ShapeLatent: {type(shape_structure)}")
batch_size, counts, max_tokens = infer_batched_coord_layout(coords) batch_size, counts, max_tokens = infer_batched_coord_layout(coords)
in_channels = 32 in_channels = 32
@ -698,7 +705,14 @@ class EmptyTrellis2LatentTexture(IO.ComfyNode):
category="latent/3d", category="latent/3d",
inputs=[ inputs=[
IO.Model.Input("model"), IO.Model.Input("model"),
IO.Voxel.Input("structure_or_coords"), IO.MultiType.Input(
"shape_structure",
types=[IO.Voxel, IO.AnyType],
tooltip=(
"Shape structure input. Accepts either a voxel structure "
"or upsampled coordinates from a previous cascade stage."
)
),
IO.Latent.Input("shape_latent"), IO.Latent.Input("shape_latent"),
], ],
outputs=[ outputs=[
@ -708,14 +722,14 @@ class EmptyTrellis2LatentTexture(IO.ComfyNode):
) )
@classmethod @classmethod
def execute(cls, model, structure_or_coords, shape_latent): def execute(cls, model, shape_structure, shape_latent):
channels = 32 channels = 32
if hasattr(structure_or_coords, "data") and structure_or_coords.data.ndim == 4: if hasattr(shape_structure, "data") and shape_structure.data.ndim == 4:
decoded = structure_or_coords.data.unsqueeze(1) decoded = shape_structure.data.unsqueeze(1)
coords = torch.argwhere(decoded.bool())[:, [0, 2, 3, 4]].int() coords = torch.argwhere(decoded.bool())[:, [0, 2, 3, 4]].int()
elif isinstance(structure_or_coords, torch.Tensor) and structure_or_coords.ndim == 2: elif isinstance(shape_structure, torch.Tensor) and shape_structure.ndim == 2:
coords = structure_or_coords.int() coords = shape_structure.int()
batch_size, counts, max_tokens = infer_batched_coord_layout(coords) batch_size, counts, max_tokens = infer_batched_coord_layout(coords)