wangbo/ComfyUI
1
0
Fork 0
mirror of https://github.com/comfyanonymous/ComfyUI.git synced 2026-04-27 19:02:31 +08:00
Code Issues Actions 13 Packages Projects Releases Wiki Activity

Fix Trellis PR review regressions

Browse Source
This commit is contained in:
John Pollock 2026-04-20 12:15:49 -05:00
parent c81ddf2349
commit 49c1adeed6
3 changed files with 68 additions and 186 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

86
comfy/ldm/trellis2/model.py
View File

@ -829,21 +829,6 @@ class Trellis2(nn.Module):
t_eval = timestep t_eval = timestep
c_eval = context c_eval = context
x_eval_norms = [float(v) for v in x_eval.square().sum(dim=(1, 2)).detach().cpu().tolist()]
c_eval_norms = [float(v) for v in c_eval.square().sum(dim=(1, 2)).detach().cpu().tolist()]
print(
"TRELLIS2_NOT_STRUCT_INPUT_TRACE",
{
"mode": mode,
"orig_bsz": int(orig_bsz),
"logical_batch": int(logical_batch),
"rule": bool(rule),
"coord_counts": coord_counts.tolist() if coord_counts is not None else None,
"x_eval_norms": x_eval_norms,
"c_eval_norms": c_eval_norms,
},
)
B, N, C = x_eval.shape B, N, C = x_eval.shape
if mode in ["shape_generation", "texture_generation"]: if mode in ["shape_generation", "texture_generation"]:
@ -878,16 +863,6 @@ class Trellis2(nn.Module):
coord_batches.append(coords_rep) coord_batches.append(coords_rep)
index_batch.append(out_index) index_batch.append(out_index)
print(
"TRELLIS2_GROUPED_INPUT_TRACE",
{
"mode": mode,
"sample_index": int(i),
"coord_count": int(count),
"feat_norms": [float(v.square().sum().detach().cpu().item()) for v in feat_batches],
},
)
x_st_i = SparseTensor( x_st_i = SparseTensor(
feats=torch.cat(feat_batches, dim=0), feats=torch.cat(feat_batches, dim=0),
coords=torch.cat(coord_batches, dim=0).to(torch.int32), coords=torch.cat(coord_batches, dim=0).to(torch.int32),
@ -972,16 +947,6 @@ class Trellis2(nn.Module):
active_coord_counts.append(count) active_coord_counts.append(count)
out_channels = sparse_outs[0].shape[-1] out_channels = sparse_outs[0].shape[-1]
sparse_out_norms = [float(feats.square().sum().detach().cpu().item()) for feats in sparse_outs]
print(
"TRELLIS2_SPARSE_OUT_TRACE",
{
"mode": mode,
"coords_rows": int(coords.shape[0]),
"active_coord_counts": active_coord_counts,
"sparse_out_norms": sparse_out_norms,
},
)
padded = sparse_outs[0].new_zeros((B, N, out_channels)) padded = sparse_outs[0].new_zeros((B, N, out_channels))
for out_index, (count, feats_i) in enumerate(zip(active_coord_counts, sparse_outs)): for out_index, (count, feats_i) in enumerate(zip(active_coord_counts, sparse_outs)):
padded[out_index, :count] = feats_i padded[out_index, :count] = feats_i
@ -1060,20 +1025,6 @@ class Trellis2(nn.Module):
cond_or_uncond = transformer_options.get("cond_or_uncond") or [] cond_or_uncond = transformer_options.get("cond_or_uncond") or []
batch_groups = len(cond_or_uncond) if len(cond_or_uncond) > 0 and orig_bsz % len(cond_or_uncond) == 0 else 1 batch_groups = len(cond_or_uncond) if len(cond_or_uncond) > 0 and orig_bsz % len(cond_or_uncond) == 0 else 1
logical_batch = orig_bsz // batch_groups logical_batch = orig_bsz // batch_groups
print(
"TRELLIS2_STRUCTURE_INPUT_TRACE",
{
"orig_bsz": int(orig_bsz),
"batch_groups": int(batch_groups),
"logical_batch": int(logical_batch),
"cond_or_uncond": cond_or_uncond,
"x_norms": [float(v) for v in x.square().sum(dim=(1, 2, 3, 4)).detach().cpu().tolist()],
"x_sums": [float(v) for v in x.sum(dim=(1, 2, 3, 4)).detach().cpu().tolist()],
"c_norms": [float(v) for v in context.square().sum(dim=(1, 2)).detach().cpu().tolist()],
"c_sums": [float(v) for v in context.sum(dim=(1, 2)).detach().cpu().tolist()],
},
)
if logical_batch > 1: if logical_batch > 1:
x_groups = x.reshape(batch_groups, logical_batch, *x.shape[1:]) x_groups = x.reshape(batch_groups, logical_batch, *x.shape[1:])
if timestep.shape[0] > 1: if timestep.shape[0] > 1:
@ -1088,10 +1039,6 @@ class Trellis2(nn.Module):
selected_group_indices = list(range(batch_groups)) selected_group_indices = list(range(batch_groups))
out_groups = [] out_groups = []
selected_x_norms = []
selected_x_sums = []
selected_c_norms = []
selected_c_sums = []
for sample_index in range(logical_batch): for sample_index in range(logical_batch):
if shape_rule and batch_groups > 1: if shape_rule and batch_groups > 1:
half = orig_bsz // 2 half = orig_bsz // 2
@ -1111,23 +1058,8 @@ class Trellis2(nn.Module):
else: else:
t_i = timestep t_i = timestep
c_i = c_groups[selected_group_indices, sample_index] c_i = c_groups[selected_group_indices, sample_index]
selected_x_norms.extend(float(v) for v in x_i.square().sum(dim=(1, 2, 3, 4)).detach().cpu().tolist())
selected_x_sums.extend(float(v) for v in x_i.sum(dim=(1, 2, 3, 4)).detach().cpu().tolist())
selected_c_norms.extend(float(v) for v in c_i.square().sum(dim=(1, 2)).detach().cpu().tolist())
selected_c_sums.extend(float(v) for v in c_i.sum(dim=(1, 2)).detach().cpu().tolist())
out_groups.append(self.structure_model(x_i, t_i, c_i)) out_groups.append(self.structure_model(x_i, t_i, c_i))
print(
"TRELLIS2_STRUCTURE_SELECTED_TRACE",
{
"selected_group_indices": selected_group_indices,
"selected_x_norms": selected_x_norms,
"selected_x_sums": selected_x_sums,
"selected_c_norms": selected_c_norms,
"selected_c_sums": selected_c_sums,
},
)
out = out_groups[0].new_zeros((orig_bsz, *out_groups[0].shape[1:])) out = out_groups[0].new_zeros((orig_bsz, *out_groups[0].shape[1:]))
for sample_index, out_sample in enumerate(out_groups): for sample_index, out_sample in enumerate(out_groups):
if shape_rule and batch_groups > 1: if shape_rule and batch_groups > 1:
@ -1146,28 +1078,10 @@ class Trellis2(nn.Module):
if shape_rule and orig_bsz > 1: if shape_rule and orig_bsz > 1:
out = out.repeat(2, 1, 1, 1, 1) out = out.repeat(2, 1, 1, 1, 1)
print(
"TRELLIS2_STRUCTURE_OUTPUT_TRACE",
{
"out_norms": [float(v) for v in out.square().sum(dim=(1, 2, 3, 4)).detach().cpu().tolist()],
"out_sums": [float(v) for v in out.sum(dim=(1, 2, 3, 4)).detach().cpu().tolist()],
},
)
if not_struct_mode: if not_struct_mode:
if dense_out is None: if dense_out is None:
out = out.feats out = out.feats
out = out.view(B, N, -1).transpose(1, 2).unsqueeze(-1) out = out.view(B, N, -1).transpose(1, 2).unsqueeze(-1)
if rule and orig_bsz > B: if rule and orig_bsz > B:
out = out.repeat(orig_bsz // B, 1, 1, 1) out = out.repeat(orig_bsz // B, 1, 1, 1)
print(
"TRELLIS2_DENSE_OUT_TRACE",
{
"mode": mode,
"coords_rows": int(coords.shape[0]) if coords is not None else None,
"output_shape": list(out.shape),
"output_norms": [float(v) for v in out.squeeze(-1).square().sum(dim=(1, 2)).detach().cpu().tolist()],
"coord_counts": coord_counts.tolist() if coord_counts is not None else None,
},
)
return out return out

29
comfy/sample.py
View File

@ -10,18 +10,37 @@ def prepare_noise_inner(latent_image, generator, noise_inds=None):
coord_counts = getattr(latent_image, "trellis_coord_counts", None) coord_counts = getattr(latent_image, "trellis_coord_counts", None)
if coord_counts is not None: if coord_counts is not None:
noise = torch.zeros(latent_image.size(), dtype=torch.float32, layout=latent_image.layout, device="cpu") noise = torch.zeros(latent_image.size(), dtype=torch.float32, layout=latent_image.layout, device="cpu")
base_state = generator.get_state() if noise_inds is None:
for i, count in enumerate(coord_counts.tolist()): noise_inds = np.arange(latent_image.size(0), dtype=np.int64)
else:
noise_inds = np.asarray(noise_inds, dtype=np.int64)
unique_inds = np.unique(noise_inds)
first_indices = {int(unique_index): int(np.flatnonzero(noise_inds == unique_index)[0]) for unique_index in unique_inds.tolist()}
index_states = {}
for unique_index in sorted(first_indices):
index_states[unique_index] = generator.get_state().clone()
count = int(coord_counts[first_indices[unique_index]].item())
torch.randn(
[1, latent_image.size(1), count, latent_image.size(3)],
dtype=torch.float32,
layout=latent_image.layout,
generator=generator,
device="cpu",
)
for batch_index, noise_index in enumerate(noise_inds.tolist()):
count = int(coord_counts[batch_index].item())
local_generator = torch.Generator(device="cpu") local_generator = torch.Generator(device="cpu")
local_generator.set_state(base_state.clone()) local_generator.set_state(index_states[int(noise_index)].clone())
sample_noise = torch.randn( sample_noise = torch.randn(
[1, latent_image.size(1), int(count), latent_image.size(3)], [1, latent_image.size(1), count, latent_image.size(3)],
dtype=torch.float32, dtype=torch.float32,
layout=latent_image.layout, layout=latent_image.layout,
generator=local_generator, generator=local_generator,
device="cpu", device="cpu",
) )
noise[i:i + 1, :, :int(count), :] = sample_noise noise[batch_index:batch_index + 1, :, :count, :] = sample_noise
return noise.to(dtype=latent_image.dtype) return noise.to(dtype=latent_image.dtype)
if noise_inds is None: if noise_inds is None:

139
comfy_extras/nodes_trellis2.py
View File

@ -148,18 +148,6 @@ def split_batched_sparse_latent(samples, coords, coord_counts):
return items return items
def log_sparse_batch_trace(tag, items):
feat_norms = [float(feats.square().sum().detach().cpu().item()) for feats, _ in items]
coord_rows = [int(coords_i.shape[0]) for _, coords_i in items]
print(
tag,
{
"batch_size": len(items),
"coord_rows": coord_rows,
"feat_norms": feat_norms,
},
)
def paint_mesh_with_voxels(mesh, voxel_coords, voxel_colors, resolution): def paint_mesh_with_voxels(mesh, voxel_coords, voxel_colors, resolution):
""" """
Generic function to paint a mesh using nearest-neighbor colors from a sparse voxel field. Generic function to paint a mesh using nearest-neighbor colors from a sparse voxel field.
@ -410,14 +398,6 @@ class Trellis2UpsampleCascade(IO.ComfyNode):
) )
feats = feats.to(device) feats = feats.to(device)
coords_512 = coords_512.to(device) coords_512 = coords_512.to(device)
print(
"TRELLIS2_UPSAMPLE_INPUT_TRACE",
{
"batch_size": 1,
"coord_rows": [int(coords_512.shape[0])],
"feat_norms": [float(feats.square().sum().detach().cpu().item())],
},
)
slat = shape_norm(feats, coords_512) slat = shape_norm(feats, coords_512)
slat.feats = slat.feats.to(next(decoder.parameters()).dtype) slat.feats = slat.feats.to(next(decoder.parameters()).dtype)
hr_coords = decoder.upsample(slat, upsample_times=4) hr_coords = decoder.upsample(slat, upsample_times=4)
@ -435,27 +415,18 @@ class Trellis2UpsampleCascade(IO.ComfyNode):
break break
hr_resolution -= 128 hr_resolution -= 128
print(
"TRELLIS2_UPSAMPLE_OUTPUT_TRACE",
{
"batch_size": 1,
"coord_rows": [int(final_coords.shape[0])],
"hr_resolution": int(hr_resolution),
},
)
return IO.NodeOutput(final_coords,) return IO.NodeOutput(final_coords,)
final_coords_list = []
items = split_batched_sparse_latent( items = split_batched_sparse_latent(
shape_latent_512["samples"], shape_latent_512["samples"],
shape_latent_512["coords"], shape_latent_512["coords"],
coord_counts, coord_counts,
) )
log_sparse_batch_trace("TRELLIS2_UPSAMPLE_INPUT_TRACE", items)
decoder_dtype = next(decoder.parameters()).dtype decoder_dtype = next(decoder.parameters()).dtype
output_coord_rows = [] final_coords_list = []
output_resolutions = [] output_resolutions = []
output_coord_counts = []
for batch_index, (feats_i, coords_i) in enumerate(items): for batch_index, (feats_i, coords_i) in enumerate(items):
feats_i = feats_i.to(device) feats_i = feats_i.to(device)
coords_i = coords_i.to(device).clone() coords_i = coords_i.to(device).clone()
@ -480,19 +451,14 @@ class Trellis2UpsampleCascade(IO.ComfyNode):
final_coords_i = final_coords_i.clone() final_coords_i = final_coords_i.clone()
final_coords_i[:, 0] = batch_index final_coords_i[:, 0] = batch_index
final_coords_list.append(final_coords_i) final_coords_list.append(final_coords_i)
output_coord_rows.append(int(final_coords_i.shape[0]))
output_resolutions.append(int(hr_resolution)) output_resolutions.append(int(hr_resolution))
output_coord_counts.append(int(final_coords_i.shape[0]))
print( return IO.NodeOutput({
"TRELLIS2_UPSAMPLE_OUTPUT_TRACE", "coords": torch.cat(final_coords_list, dim=0),
{ "coord_counts": torch.tensor(output_coord_counts, dtype=torch.int64),
"batch_size": len(final_coords_list), "resolutions": torch.tensor(output_resolutions, dtype=torch.int64),
"coord_rows": output_coord_rows, },)
"hr_resolution": output_resolutions,
},
)
return IO.NodeOutput(torch.cat(final_coords_list, dim=0),)
dino_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1) dino_mean = torch.tensor([0.485, 0.456, 0.406]).view(1, 3, 1, 1)
dino_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) dino_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1)
@ -568,7 +534,6 @@ class Trellis2Conditioning(IO.ComfyNode):
cond_512_list = [] cond_512_list = []
cond_1024_list = [] cond_1024_list = []
composite_trace = []
for b in range(batch_size): for b in range(batch_size):
item_image = image[b] item_image = image[b]
@ -623,14 +588,6 @@ class Trellis2Conditioning(IO.ComfyNode):
# to match trellis2 code (quantize -> dequantize) # to match trellis2 code (quantize -> dequantize)
composite_uint8 = (composite_np * 255.0).round().clip(0, 255).astype(np.uint8) composite_uint8 = (composite_np * 255.0).round().clip(0, 255).astype(np.uint8)
composite_trace.append(
{
"sample_index": int(b),
"shape": list(composite_uint8.shape),
"sum": int(composite_uint8.sum(dtype=np.int64)),
"prefix": composite_uint8[:2, :2, :].reshape(-1).tolist(),
}
)
cropped_pil = Image.fromarray(composite_uint8) cropped_pil = Image.fromarray(composite_uint8)
@ -642,19 +599,6 @@ class Trellis2Conditioning(IO.ComfyNode):
cond_1024_batched = torch.cat(cond_1024_list, dim=0) cond_1024_batched = torch.cat(cond_1024_list, dim=0)
neg_cond_batched = torch.zeros_like(cond_512_batched) neg_cond_batched = torch.zeros_like(cond_512_batched)
neg_embeds_batched = torch.zeros_like(cond_1024_batched) neg_embeds_batched = torch.zeros_like(cond_1024_batched)
print(
"TRELLIS2_CONDITIONING_TRACE",
{
"batch_size": int(batch_size),
"cond_512_norms": [float(v) for v in cond_512_batched.square().sum(dim=(1, 2)).detach().cpu().tolist()],
"cond_512_sums": [float(v) for v in cond_512_batched.sum(dim=(1, 2)).detach().cpu().tolist()],
"cond_512_prefix": cond_512_batched[:, 0, :8].detach().cpu().tolist(),
"cond_1024_norms": [float(v) for v in cond_1024_batched.square().sum(dim=(1, 2)).detach().cpu().tolist()],
"cond_1024_sums": [float(v) for v in cond_1024_batched.sum(dim=(1, 2)).detach().cpu().tolist()],
"cond_1024_prefix": cond_1024_batched[:, 0, :8].detach().cpu().tolist(),
"composite_trace": composite_trace,
},
)
positive = [[cond_512_batched, {"embeds": cond_1024_batched}]] positive = [[cond_512_batched, {"embeds": cond_1024_batched}]]
negative = [[neg_cond_batched, {"embeds": neg_embeds_batched}]] negative = [[neg_cond_batched, {"embeds": neg_embeds_batched}]]
@ -680,12 +624,20 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode):
def execute(cls, structure_or_coords, model): def execute(cls, structure_or_coords, model):
# to accept the upscaled coords # to accept the upscaled coords
is_512_pass = False is_512_pass = False
coord_counts = None
coord_resolutions = None
if hasattr(structure_or_coords, "data") and structure_or_coords.data.ndim == 4: if hasattr(structure_or_coords, "data") and structure_or_coords.data.ndim == 4:
decoded = structure_or_coords.data.unsqueeze(1) decoded = structure_or_coords.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, dict):
coords = structure_or_coords["coords"].int()
coord_counts = structure_or_coords.get("coord_counts")
coord_resolutions = structure_or_coords.get("resolutions")
is_512_pass = False
elif isinstance(structure_or_coords, torch.Tensor) and structure_or_coords.ndim == 2: elif isinstance(structure_or_coords, torch.Tensor) and structure_or_coords.ndim == 2:
coords = structure_or_coords.int() coords = structure_or_coords.int()
is_512_pass = False is_512_pass = False
@ -693,7 +645,15 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode):
else: else:
raise ValueError(f"Invalid input to EmptyShapeLatent: {type(structure_or_coords)}") raise ValueError(f"Invalid input to EmptyShapeLatent: {type(structure_or_coords)}")
in_channels = 32 in_channels = 32
batch_size, coord_counts, max_tokens = infer_batched_coord_layout(coords) batch_size, inferred_coord_counts, max_tokens = infer_batched_coord_layout(coords)
if coord_counts is not None:
coord_counts = coord_counts.to(dtype=torch.int64, device=coords.device)
if coord_counts.shape != inferred_coord_counts.shape or not torch.equal(coord_counts, inferred_coord_counts):
raise ValueError(
f"Trellis2 coord_counts metadata {coord_counts.tolist()} does not match coords layout {inferred_coord_counts.tolist()}"
)
else:
coord_counts = inferred_coord_counts
if batch_size == 1: if batch_size == 1:
coord_counts = None coord_counts = None
latent = torch.randn(1, in_channels, coords.shape[0], 1) latent = torch.randn(1, in_channels, coords.shape[0], 1)
@ -706,17 +666,6 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode):
generator.set_state(base_state.clone()) generator.set_state(base_state.clone())
latent_i = torch.randn(1, in_channels, count, 1, generator=generator) latent_i = torch.randn(1, in_channels, count, 1, generator=generator)
latent[i, :, :count] = latent_i[0] latent[i, :, :count] = latent_i[0]
if coords.shape[0] > 1000:
norms = [float(v) for v in latent.squeeze(-1).square().sum(dim=(1, 2)).detach().cpu().tolist()]
print(
"TRELLIS2_EMPTY_SHAPE_TRACE",
{
"coords_rows": int(coords.shape[0]),
"batch_size": int(batch_size),
"coord_counts": coord_counts.tolist() if coord_counts is not None else None,
"latent_norms": norms,
},
)
if coord_counts is not None: if coord_counts is not None:
latent.trellis_coord_counts = coord_counts.clone() latent.trellis_coord_counts = coord_counts.clone()
model = model.clone() model = model.clone()
@ -729,6 +678,8 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode):
model.model_options["transformer_options"]["coords"] = coords model.model_options["transformer_options"]["coords"] = coords
if coord_counts is not None: if coord_counts is not None:
model.model_options["transformer_options"]["coord_counts"] = coord_counts model.model_options["transformer_options"]["coord_counts"] = coord_counts
if coord_resolutions is not None:
model.model_options["transformer_options"]["coord_resolutions"] = coord_resolutions
if is_512_pass: if is_512_pass:
model.model_options["transformer_options"]["generation_mode"] = "shape_generation_512" model.model_options["transformer_options"]["generation_mode"] = "shape_generation_512"
else: else:
@ -736,6 +687,8 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode):
output = {"samples": latent, "coords": coords, "type": "trellis2"} output = {"samples": latent, "coords": coords, "type": "trellis2"}
if coord_counts is not None: if coord_counts is not None:
output["coord_counts"] = coord_counts output["coord_counts"] = coord_counts
if coord_resolutions is not None:
output["coord_resolutions"] = coord_resolutions
output["batch_index"] = [0] * batch_size output["batch_index"] = [0] * batch_size
return IO.NodeOutput(output, model) return IO.NodeOutput(output, model)
@ -759,15 +712,28 @@ class EmptyTextureLatentTrellis2(IO.ComfyNode):
@classmethod @classmethod
def execute(cls, structure_or_coords, shape_latent, model): def execute(cls, structure_or_coords, shape_latent, model):
channels = 32 channels = 32
coord_counts = None
if hasattr(structure_or_coords, "data") and structure_or_coords.data.ndim == 4: if hasattr(structure_or_coords, "data") and structure_or_coords.data.ndim == 4:
decoded = structure_or_coords.data.unsqueeze(1) decoded = structure_or_coords.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, dict):
coords = structure_or_coords["coords"].int()
coord_counts = structure_or_coords.get("coord_counts")
elif isinstance(structure_or_coords, torch.Tensor) and structure_or_coords.ndim == 2: elif isinstance(structure_or_coords, torch.Tensor) and structure_or_coords.ndim == 2:
coords = structure_or_coords.int() coords = structure_or_coords.int()
shape_latent = shape_latent["samples"] shape_latent = shape_latent["samples"]
batch_size, coord_counts, max_tokens = infer_batched_coord_layout(coords) batch_size, inferred_coord_counts, max_tokens = infer_batched_coord_layout(coords)
if coord_counts is not None:
coord_counts = coord_counts.to(dtype=torch.int64, device=coords.device)
if coord_counts.shape != inferred_coord_counts.shape or not torch.equal(coord_counts, inferred_coord_counts):
raise ValueError(
f"Trellis2 coord_counts metadata {coord_counts.tolist()} does not match coords layout {inferred_coord_counts.tolist()}"
)
else:
coord_counts = inferred_coord_counts
if shape_latent.ndim == 4: if shape_latent.ndim == 4:
if shape_latent.shape[0] != batch_size: if shape_latent.shape[0] != batch_size:
raise ValueError( raise ValueError(
@ -791,19 +757,6 @@ class EmptyTextureLatentTrellis2(IO.ComfyNode):
generator.set_state(base_state.clone()) generator.set_state(base_state.clone())
latent_i = torch.randn(1, channels, count, 1, generator=generator) latent_i = torch.randn(1, channels, count, 1, generator=generator)
latent[i, :, :count] = latent_i[0] latent[i, :, :count] = latent_i[0]
if coords.shape[0] > 1000:
norms = [float(v) for v in latent.squeeze(-1).square().sum(dim=(1, 2)).detach().cpu().tolist()]
shape_norms = [float(v) for v in shape_latent.square().sum(dim=(1, 2)).detach().cpu().tolist()] if shape_latent.ndim == 3 else None
print(
"TRELLIS2_EMPTY_TEXTURE_TRACE",
{
"coords_rows": int(coords.shape[0]),
"batch_size": int(batch_size),
"coord_counts": coord_counts.tolist() if coord_counts is not None else None,
"latent_norms": norms,
"shape_latent_norms": shape_norms,
},
)
if coord_counts is not None: if coord_counts is not None:
latent.trellis_coord_counts = coord_counts.clone() latent.trellis_coord_counts = coord_counts.clone()
model = model.clone() model = model.clone()
@ -842,11 +795,7 @@ class EmptyStructureLatentTrellis2(IO.ComfyNode):
def execute(cls, batch_size): def execute(cls, batch_size):
in_channels = 8 in_channels = 8
resolution = 16 resolution = 16
generator = torch.Generator(device="cpu") latent = torch.randn(1, in_channels, resolution, resolution, resolution).repeat(batch_size, 1, 1, 1, 1)
generator.manual_seed(11426)
latent = torch.randn(1, in_channels, resolution, resolution, resolution, generator=generator).repeat(batch_size, 1, 1, 1, 1)
norms = [float(v) for v in latent.square().sum(dim=(1, 2, 3, 4)).detach().cpu().tolist()]
print("TRELLIS2_EMPTY_STRUCTURE_TRACE", {"batch_size": int(batch_size), "latent_norms": norms})
output = {"samples": latent, "type": "trellis2"} output = {"samples": latent, "type": "trellis2"}
if batch_size > 1: if batch_size > 1:
output["batch_index"] = [0] * batch_size output["batch_index"] = [0] * batch_size