From c81ddf23498d27f82293272494fd66f31dacb7fc Mon Sep 17 00:00:00 2001 From: John Pollock Date: Mon, 20 Apr 2026 11:06:04 -0500 Subject: [PATCH] Fix Trellis2 batched shape and texture semantics --- comfy/ldm/trellis2/model.py | 347 ++++++++++++++++++++++++++++++--- comfy/sample.py | 17 ++ comfy_extras/nodes_trellis2.py | 332 +++++++++++++++++++++++++++---- 3 files changed, 635 insertions(+), 61 deletions(-) diff --git a/comfy/ldm/trellis2/model.py b/comfy/ldm/trellis2/model.py index 1c5d6c3ec..76dbacc93 100644 --- a/comfy/ldm/trellis2/model.py +++ b/comfy/ldm/trellis2/model.py @@ -786,6 +786,7 @@ class Trellis2(nn.Module): # 32 -> 512px path, 64 -> 1024px path. uses_1024_conditioning = self.img2shape.resolution == 64 coords = transformer_options.get("coords", None) + coord_counts = transformer_options.get("coord_counts") mode = transformer_options.get("generation_mode", "structure_generation") is_512_run = False timestep = timestep.to(self.dtype) @@ -811,40 +812,205 @@ class Trellis2(nn.Module): cond = context shape_rule = sigmas < self.guidance_interval[0] or sigmas > self.guidance_interval[1] txt_rule = sigmas < self.guidance_interval_txt[0] or sigmas > self.guidance_interval_txt[1] + dense_out = None if not_struct_mode: orig_bsz = x.shape[0] rule = txt_rule if mode == "texture_generation" else shape_rule - if rule and orig_bsz > 1: - x_eval = x[1].unsqueeze(0) - t_eval = timestep[1].unsqueeze(0) if timestep.shape[0] > 1 else timestep + logical_batch = coord_counts.shape[0] if coord_counts is not None else 1 + if rule and orig_bsz > logical_batch: + half = orig_bsz // 2 + x_eval = x[half:] + t_eval = timestep[half:] if timestep.shape[0] > 1 else timestep c_eval = cond else: x_eval = x t_eval = timestep 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 if mode in ["shape_generation", "texture_generation"]: - feats_flat = x_eval.reshape(-1, C) + if coord_counts is not None: + logical_batch = coord_counts.shape[0] + if B % logical_batch != 0: + raise ValueError( + f"Trellis2 coord_counts batch {logical_batch} doesn't divide latent batch {B}" + ) + repeat_factor = B // logical_batch + sparse_outs = [] + active_coord_counts = [] + if mode == "shape_generation" and repeat_factor > 1: + grouped_outs = [] + grouped_counts = [] + for i in range(logical_batch): + count = int(coord_counts[i].item()) + coords_i = coords[coords[:, 0] == i].clone() + if coords_i.shape[0] != count: + raise ValueError( + f"Trellis2 coords rows for batch {i} expected {count}, got {coords_i.shape[0]}" + ) - # inflate coords [N, 4] -> [B*N, 4] - coords_list = [] - for i in range(B): - c = coords.clone() - c[:, 0] = i - coords_list.append(c) + feat_batches = [] + coord_batches = [] + index_batch = [] + for rep in range(repeat_factor): + out_index = rep * logical_batch + i + feat_batches.append(x_eval[out_index, :count]) + coords_rep = coords_i.clone() + coords_rep[:, 0] = rep + coord_batches.append(coords_rep) + index_batch.append(out_index) - batched_coords = torch.cat(coords_list, dim=0) + 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( + feats=torch.cat(feat_batches, dim=0), + coords=torch.cat(coord_batches, dim=0).to(torch.int32), + ) + index_tensor = torch.tensor(index_batch, device=x_eval.device, dtype=torch.long) + if t_eval.shape[0] > 1: + t_i = t_eval.index_select(0, index_tensor) + else: + t_i = t_eval + if c_eval.shape[0] > 1: + c_i = c_eval.index_select(0, index_tensor) + else: + c_i = c_eval + + if is_512_run: + sparse_out = self.img2shape_512(x_st_i, t_i, c_i) + else: + sparse_out = self.img2shape(x_st_i, t_i, c_i) + + feats_group, coords_group = sparse_out.to_tensor_list() + if len(feats_group) != repeat_factor: + raise ValueError( + f"Trellis2 expected {repeat_factor} sparse output groups for batch {i}, got {len(feats_group)}" + ) + for rep, (feats_rep, coords_rep) in enumerate(zip(feats_group, coords_group)): + if feats_rep.shape[0] != count: + raise ValueError( + f"Trellis2 sparse output rows for batch {i} rep {rep} expected {count}, got {feats_rep.shape[0]}" + ) + if coords_rep.shape[0] != count: + raise ValueError( + f"Trellis2 sparse output coords for batch {i} rep {rep} expected {count}, got {coords_rep.shape[0]}" + ) + grouped_outs.append(feats_group) + grouped_counts.append(count) + + for rep in range(repeat_factor): + for i in range(logical_batch): + sparse_outs.append(grouped_outs[i][rep]) + active_coord_counts.append(grouped_counts[i]) + else: + for rep in range(repeat_factor): + for i in range(logical_batch): + out_index = rep * logical_batch + i + count = int(coord_counts[i].item()) + coords_i = coords[coords[:, 0] == i].clone() + if coords_i.shape[0] != count: + raise ValueError( + f"Trellis2 coords rows for batch {i} expected {count}, got {coords_i.shape[0]}" + ) + coords_i[:, 0] = 0 + feats_i = x_eval[out_index, :count] + x_st_i = SparseTensor(feats=feats_i, coords=coords_i.to(torch.int32)) + t_i = t_eval[out_index].unsqueeze(0) if t_eval.shape[0] > 1 else t_eval + c_i = c_eval[out_index].unsqueeze(0) if c_eval.shape[0] > 1 else c_eval + + if mode == "shape_generation": + if is_512_run: + sparse_out = self.img2shape_512(x_st_i, t_i, c_i) + else: + sparse_out = self.img2shape(x_st_i, t_i, c_i) + else: + slat = transformer_options.get("shape_slat") + if slat is None: + raise ValueError("shape_slat can't be None") + if slat.ndim == 3: + if slat.shape[0] != logical_batch: + raise ValueError( + f"shape_slat batch {slat.shape[0]} doesn't match coord_counts batch {logical_batch}" + ) + if slat.shape[1] < count: + raise ValueError( + f"shape_slat tokens {slat.shape[1]} can't cover coord count {count} for batch {i}" + ) + slat_feats = slat[i, :count].to(x_st_i.device) + else: + slat_feats = slat[:count].to(x_st_i.device) + x_st_i = x_st_i.replace(feats=torch.cat([x_st_i.feats, slat_feats], dim=-1)) + sparse_out = self.shape2txt(x_st_i, t_i, c_i) + + sparse_outs.append(sparse_out.feats) + active_coord_counts.append(count) + + 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)) + for out_index, (count, feats_i) in enumerate(zip(active_coord_counts, sparse_outs)): + padded[out_index, :count] = feats_i + dense_out = padded.transpose(1, 2).unsqueeze(-1) + elif coords.shape[0] == N: + feats_flat = x_eval.reshape(-1, C) + coords_list = [] + for i in range(B): + c = coords.clone() + c[:, 0] = i + coords_list.append(c) + batched_coords = torch.cat(coords_list, dim=0) + elif coords.shape[0] == B * N: + feats_flat = x_eval.reshape(-1, C) + batched_coords = coords + else: + raise ValueError( + f"Trellis2 expected coords rows {N} or {B * N}, got {coords.shape[0]}" + ) else: batched_coords = coords feats_flat = x_eval - x_st = SparseTensor(feats=feats_flat, coords=batched_coords.to(torch.int32)) + if dense_out is None: + x_st = SparseTensor(feats=feats_flat, coords=batched_coords.to(torch.int32)) - if mode == "shape_generation": + if dense_out is not None: + out = dense_out + elif mode == "shape_generation": if is_512_run: out = self.img2shape_512(x_st, t_eval, c_eval) else: @@ -856,23 +1022,152 @@ class Trellis2(nn.Module): if slat is None: raise ValueError("shape_slat can't be None") - base_slat_feats = slat[:N] - slat_feats_batched = base_slat_feats.repeat(B, 1).to(x_st.device) + if slat.ndim == 3: + if coord_counts is not None: + logical_batch = coord_counts.shape[0] + if slat.shape[0] != logical_batch: + raise ValueError( + f"shape_slat batch {slat.shape[0]} doesn't match coord_counts batch {logical_batch}" + ) + if B % logical_batch != 0: + raise ValueError( + f"Trellis2 coord_counts batch {logical_batch} doesn't divide latent batch {B}" + ) + repeat_factor = B // logical_batch + slat_list = [] + for _ in range(repeat_factor): + for i in range(logical_batch): + count = int(coord_counts[i].item()) + if slat.shape[1] < count: + raise ValueError( + f"shape_slat tokens {slat.shape[1]} can't cover coord count {count} for batch {i}" + ) + slat_list.append(slat[i, :count]) + slat_feats_batched = torch.cat(slat_list, dim=0).to(x_st.device) + else: + if slat.shape[0] != B: + raise ValueError(f"shape_slat batch {slat.shape[0]} doesn't match latent batch {B}") + if slat.shape[1] != N: + raise ValueError(f"shape_slat tokens {slat.shape[1]} doesn't match latent tokens {N}") + slat_feats_batched = slat.reshape(B * N, -1).to(x_st.device) + else: + base_slat_feats = slat[:N] + slat_feats_batched = base_slat_feats.repeat(B, 1).to(x_st.device) x_st = x_st.replace(feats=torch.cat([x_st.feats, slat_feats_batched], dim=-1)) out = self.shape2txt(x_st, t_eval, c_eval) else: # structure orig_bsz = x.shape[0] - if shape_rule and orig_bsz > 1: - half = orig_bsz // 2 - x = x[half:] - timestep = timestep[half:] if timestep.shape[0] > 1 else timestep - out = self.structure_model(x, timestep, cond if shape_rule and orig_bsz > 1 else context) - if shape_rule and orig_bsz > 1: - out = out.repeat(2, 1, 1, 1, 1) + 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 + 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: + x_groups = x.reshape(batch_groups, logical_batch, *x.shape[1:]) + if timestep.shape[0] > 1: + t_groups = timestep.reshape(batch_groups, logical_batch, *timestep.shape[1:]) + else: + t_groups = timestep + c_groups = context.reshape(batch_groups, logical_batch, *context.shape[1:]) + + if shape_rule and batch_groups > 1: + selected_group_indices = [batch_groups - 1] + else: + selected_group_indices = list(range(batch_groups)) + + out_groups = [] + selected_x_norms = [] + selected_x_sums = [] + selected_c_norms = [] + selected_c_sums = [] + for sample_index in range(logical_batch): + if shape_rule and batch_groups > 1: + half = orig_bsz // 2 + x_i = x[half + sample_index].unsqueeze(0) + if timestep.shape[0] > 1: + t_i = timestep[half + sample_index].unsqueeze(0) + else: + t_i = timestep + if cond.shape[0] > 1: + c_i = cond[sample_index].unsqueeze(0) + else: + c_i = cond + else: + x_i = x_groups[selected_group_indices, sample_index] + if timestep.shape[0] > 1: + t_i = t_groups[selected_group_indices, sample_index] + else: + t_i = timestep + 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)) + + 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:])) + for sample_index, out_sample in enumerate(out_groups): + if shape_rule and batch_groups > 1: + repeated = out_sample[0] + for group_index in range(batch_groups): + out[group_index * logical_batch + sample_index] = repeated + else: + for local_group_index, group_index in enumerate(selected_group_indices): + out[group_index * logical_batch + sample_index] = out_sample[local_group_index] + else: + if shape_rule and orig_bsz > 1: + half = orig_bsz // 2 + x = x[half:] + timestep = timestep[half:] if timestep.shape[0] > 1 else timestep + out = self.structure_model(x, timestep, cond if shape_rule and orig_bsz > 1 else context) + if shape_rule and orig_bsz > 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: - out = out.feats - out = out.view(B, N, -1).transpose(1, 2).unsqueeze(-1) - if rule and orig_bsz > 1: - out = out.repeat(orig_bsz, 1, 1, 1) + if dense_out is None: + out = out.feats + out = out.view(B, N, -1).transpose(1, 2).unsqueeze(-1) + if rule and orig_bsz > B: + 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 diff --git a/comfy/sample.py b/comfy/sample.py index 653829582..3967fba1b 100644 --- a/comfy/sample.py +++ b/comfy/sample.py @@ -7,6 +7,23 @@ import logging import comfy.nested_tensor def prepare_noise_inner(latent_image, generator, noise_inds=None): + coord_counts = getattr(latent_image, "trellis_coord_counts", None) + if coord_counts is not None: + noise = torch.zeros(latent_image.size(), dtype=torch.float32, layout=latent_image.layout, device="cpu") + base_state = generator.get_state() + for i, count in enumerate(coord_counts.tolist()): + local_generator = torch.Generator(device="cpu") + local_generator.set_state(base_state.clone()) + sample_noise = torch.randn( + [1, latent_image.size(1), int(count), latent_image.size(3)], + dtype=torch.float32, + layout=latent_image.layout, + generator=local_generator, + device="cpu", + ) + noise[i:i + 1, :, :int(count), :] = sample_noise + return noise.to(dtype=latent_image.dtype) + if noise_inds is None: return torch.randn(latent_image.size(), dtype=torch.float32, layout=latent_image.layout, generator=generator, device="cpu").to(dtype=latent_image.dtype) diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py index 8121e261b..26cb135e7 100644 --- a/comfy_extras/nodes_trellis2.py +++ b/comfy_extras/nodes_trellis2.py @@ -96,6 +96,70 @@ def shape_norm(shape_latent, coords): samples = samples * std + mean return samples + +def infer_batched_coord_layout(coords): + if coords.ndim != 2 or coords.shape[1] != 4: + raise ValueError(f"Expected Trellis2 coords with shape [N, 4], got {tuple(coords.shape)}") + + if coords.shape[0] == 0: + raise ValueError("Trellis2 coords can't be empty") + + batch_ids = coords[:, 0].to(torch.int64) + batch_size = int(batch_ids.max().item()) + 1 + counts = torch.bincount(batch_ids, minlength=batch_size) + + if (counts == 0).any(): + raise ValueError(f"Non-contiguous Trellis2 batch ids in coords: {batch_ids.unique(sorted=True).tolist()}") + + max_tokens = int(counts.max().item()) + return batch_size, counts, max_tokens + + +def flatten_batched_sparse_latent(samples, coords, coord_counts): + samples = samples.squeeze(-1).transpose(1, 2) + if coord_counts is None: + return samples.reshape(-1, samples.shape[-1]), coords + + feat_list = [] + coord_list = [] + for i in range(coord_counts.shape[0]): + count = int(coord_counts[i].item()) + coords_i = coords[coords[:, 0] == i] + if coords_i.shape[0] != count: + raise ValueError(f"Trellis2 coords rows for batch {i} expected {count}, got {coords_i.shape[0]}") + feat_list.append(samples[i, :count]) + coord_list.append(coords_i) + + return torch.cat(feat_list, dim=0), torch.cat(coord_list, dim=0) + + +def split_batched_sparse_latent(samples, coords, coord_counts): + samples = samples.squeeze(-1).transpose(1, 2) + if coord_counts is None: + return [(samples.reshape(-1, samples.shape[-1]), coords)] + + items = [] + for i in range(coord_counts.shape[0]): + count = int(coord_counts[i].item()) + coords_i = coords[coords[:, 0] == i] + if coords_i.shape[0] != count: + raise ValueError(f"Trellis2 coords rows for batch {i} expected {count}, got {coords_i.shape[0]}") + items.append((samples[i, :count], coords_i)) + 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): """ Generic function to paint a mesh using nearest-neighbor colors from a sparse voxel field. @@ -169,12 +233,32 @@ class VaeDecodeShapeTrellis(IO.ComfyNode): vae = vae.first_stage_model coords = samples["coords"] + coord_counts = samples.get("coord_counts") samples = samples["samples"] - samples = samples.squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device) - samples = shape_norm(samples, coords) + if coord_counts is None: + samples, coords = flatten_batched_sparse_latent(samples, coords, coord_counts) + samples = shape_norm(samples.to(device), coords.to(device)) + mesh, subs = vae.decode_shape_slat(samples, resolution) + else: + split_items = split_batched_sparse_latent(samples, coords, coord_counts) + mesh = [] + subs_per_sample = [] + for feats_i, coords_i in split_items: + coords_i = coords_i.to(device).clone() + coords_i[:, 0] = 0 + sample_i = shape_norm(feats_i.to(device), coords_i) + mesh_i, subs_i = vae.decode_shape_slat(sample_i, resolution) + mesh.append(mesh_i[0]) + subs_per_sample.append(subs_i) + + subs = [] + for stage_index in range(len(subs_per_sample[0])): + stage_tensors = [sample_subs[stage_index] for sample_subs in subs_per_sample] + feats_list = [stage_tensor.feats for stage_tensor in stage_tensors] + coords_list = [stage_tensor.coords for stage_tensor in stage_tensors] + subs.append(SparseTensor.from_tensor_list(feats_list, coords_list)) - mesh, subs = vae.decode_shape_slat(samples, resolution) face_list = [m.faces for m in mesh] vert_list = [m.vertices for m in mesh] if all(v.shape == vert_list[0].shape for v in vert_list) and all(f.shape == face_list[0].shape for f in face_list): @@ -210,12 +294,14 @@ class VaeDecodeTextureTrellis(IO.ComfyNode): vae = vae.first_stage_model coords = samples["coords"] + coord_counts = samples.get("coord_counts") samples = samples["samples"] - samples = samples.squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device) + samples, coords = flatten_batched_sparse_latent(samples, coords, coord_counts) + samples = samples.to(device) std = tex_slat_normalization["std"].to(samples) mean = tex_slat_normalization["mean"].to(samples) - samples = SparseTensor(feats = samples, coords=coords) + samples = SparseTensor(feats = samples, coords=coords.to(device)) samples = samples * std + mean voxel = vae.decode_tex_slat(samples, shape_subs) @@ -273,7 +359,13 @@ class VaeDecodeStructureTrellis2(IO.ComfyNode): decoder = decoder.to(load_device) samples = samples["samples"] samples = samples.to(load_device) - decoded = decoder(samples)>0 + if samples.shape[0] > 1: + decoded_items = [] + for i in range(samples.shape[0]): + decoded_items.append(decoder(samples[i:i + 1]) > 0) + decoded = torch.cat(decoded_items, dim=0) + else: + decoded = decoder(samples) > 0 decoder.to(offload_device) current_res = decoded.shape[2] @@ -305,32 +397,102 @@ class Trellis2UpsampleCascade(IO.ComfyNode): device = comfy.model_management.get_torch_device() comfy.model_management.load_model_gpu(vae.patcher) - feats = shape_latent_512["samples"].squeeze(-1).transpose(1, 2).reshape(-1, 32).to(device) - coords_512 = shape_latent_512["coords"].to(device) - - slat = shape_norm(feats, coords_512) - + coord_counts = shape_latent_512.get("coord_counts") decoder = vae.first_stage_model.shape_dec - - slat.feats = slat.feats.to(next(decoder.parameters()).dtype) - hr_coords = decoder.upsample(slat, upsample_times=4) - lr_resolution = 512 - hr_resolution = int(target_resolution) + target_resolution = int(target_resolution) - while True: - quant_coords = torch.cat([ - hr_coords[:, :1], - ((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(), - ], dim=1) - final_coords = quant_coords.unique(dim=0) - num_tokens = final_coords.shape[0] + if coord_counts is None: + feats, coords_512 = flatten_batched_sparse_latent( + shape_latent_512["samples"], + shape_latent_512["coords"], + coord_counts, + ) + feats = feats.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.feats = slat.feats.to(next(decoder.parameters()).dtype) + hr_coords = decoder.upsample(slat, upsample_times=4) - if num_tokens < max_tokens or hr_resolution <= 1024: - break - hr_resolution -= 128 + hr_resolution = target_resolution + while True: + quant_coords = torch.cat([ + hr_coords[:, :1], + ((hr_coords[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(), + ], dim=1) + final_coords = quant_coords.unique(dim=0) + num_tokens = final_coords.shape[0] - return IO.NodeOutput(final_coords,) + if num_tokens < max_tokens or hr_resolution <= 1024: + break + 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,) + + final_coords_list = [] + items = split_batched_sparse_latent( + shape_latent_512["samples"], + shape_latent_512["coords"], + coord_counts, + ) + log_sparse_batch_trace("TRELLIS2_UPSAMPLE_INPUT_TRACE", items) + decoder_dtype = next(decoder.parameters()).dtype + + output_coord_rows = [] + output_resolutions = [] + for batch_index, (feats_i, coords_i) in enumerate(items): + feats_i = feats_i.to(device) + coords_i = coords_i.to(device).clone() + coords_i[:, 0] = 0 + slat_i = shape_norm(feats_i, coords_i) + slat_i.feats = slat_i.feats.to(decoder_dtype) + hr_coords_i = decoder.upsample(slat_i, upsample_times=4) + + hr_resolution = target_resolution + while True: + quant_coords_i = torch.cat([ + hr_coords_i[:, :1], + ((hr_coords_i[:, 1:] + 0.5) / lr_resolution * (hr_resolution // 16)).int(), + ], dim=1) + final_coords_i = quant_coords_i.unique(dim=0) + num_tokens = final_coords_i.shape[0] + + if num_tokens < max_tokens or hr_resolution <= 1024: + break + hr_resolution -= 128 + + final_coords_i = final_coords_i.clone() + final_coords_i[:, 0] = batch_index + final_coords_list.append(final_coords_i) + output_coord_rows.append(int(final_coords_i.shape[0])) + output_resolutions.append(int(hr_resolution)) + + print( + "TRELLIS2_UPSAMPLE_OUTPUT_TRACE", + { + "batch_size": len(final_coords_list), + "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_std = torch.tensor([0.229, 0.224, 0.225]).view(1, 3, 1, 1) @@ -406,6 +568,7 @@ class Trellis2Conditioning(IO.ComfyNode): cond_512_list = [] cond_1024_list = [] + composite_trace = [] for b in range(batch_size): item_image = image[b] @@ -460,6 +623,14 @@ class Trellis2Conditioning(IO.ComfyNode): # to match trellis2 code (quantize -> dequantize) 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) @@ -471,6 +642,19 @@ class Trellis2Conditioning(IO.ComfyNode): cond_1024_batched = torch.cat(cond_1024_list, dim=0) neg_cond_batched = torch.zeros_like(cond_512_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}]] negative = [[neg_cond_batched, {"embeds": neg_embeds_batched}]] @@ -509,8 +693,32 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode): else: raise ValueError(f"Invalid input to EmptyShapeLatent: {type(structure_or_coords)}") in_channels = 32 - # image like format - latent = torch.randn(1, in_channels, coords.shape[0], 1) + batch_size, coord_counts, max_tokens = infer_batched_coord_layout(coords) + if batch_size == 1: + coord_counts = None + latent = torch.randn(1, in_channels, coords.shape[0], 1) + else: + latent = torch.zeros(batch_size, in_channels, max_tokens, 1) + base_state = torch.random.get_rng_state() + for i in range(batch_size): + count = int(coord_counts[i].item()) + generator = torch.Generator(device="cpu") + generator.set_state(base_state.clone()) + latent_i = torch.randn(1, in_channels, count, 1, generator=generator) + 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: + latent.trellis_coord_counts = coord_counts.clone() model = model.clone() model.model_options = model.model_options.copy() if "transformer_options" in model.model_options: @@ -519,11 +727,17 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode): model.model_options["transformer_options"] = {} model.model_options["transformer_options"]["coords"] = coords + if coord_counts is not None: + model.model_options["transformer_options"]["coord_counts"] = coord_counts if is_512_pass: model.model_options["transformer_options"]["generation_mode"] = "shape_generation_512" else: model.model_options["transformer_options"]["generation_mode"] = "shape_generation" - return IO.NodeOutput({"samples": latent, "coords": coords, "type": "trellis2"}, model) + output = {"samples": latent, "coords": coords, "type": "trellis2"} + if coord_counts is not None: + output["coord_counts"] = coord_counts + output["batch_index"] = [0] * batch_size + return IO.NodeOutput(output, model) class EmptyTextureLatentTrellis2(IO.ComfyNode): @classmethod @@ -553,10 +767,45 @@ class EmptyTextureLatentTrellis2(IO.ComfyNode): coords = structure_or_coords.int() shape_latent = shape_latent["samples"] + batch_size, coord_counts, max_tokens = infer_batched_coord_layout(coords) if shape_latent.ndim == 4: - shape_latent = shape_latent.squeeze(-1).transpose(1, 2).reshape(-1, channels) + if shape_latent.shape[0] != batch_size: + raise ValueError( + f"shape_latent batch {shape_latent.shape[0]} doesn't match coords batch {batch_size}" + ) + shape_latent = shape_latent.squeeze(-1).transpose(1, 2) + if shape_latent.shape[1] < max_tokens: + raise ValueError( + f"shape_latent tokens {shape_latent.shape[1]} can't cover coords max tokens {max_tokens}" + ) - latent = torch.randn(1, channels, coords.shape[0], 1) + if batch_size == 1: + coord_counts = None + latent = torch.randn(1, channels, coords.shape[0], 1) + else: + latent = torch.zeros(batch_size, channels, max_tokens, 1) + base_state = torch.random.get_rng_state() + for i in range(batch_size): + count = int(coord_counts[i].item()) + generator = torch.Generator(device="cpu") + generator.set_state(base_state.clone()) + latent_i = torch.randn(1, channels, count, 1, generator=generator) + 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: + latent.trellis_coord_counts = coord_counts.clone() model = model.clone() model.model_options = model.model_options.copy() if "transformer_options" in model.model_options: @@ -565,9 +814,15 @@ class EmptyTextureLatentTrellis2(IO.ComfyNode): model.model_options["transformer_options"] = {} model.model_options["transformer_options"]["coords"] = coords + if coord_counts is not None: + model.model_options["transformer_options"]["coord_counts"] = coord_counts model.model_options["transformer_options"]["generation_mode"] = "texture_generation" model.model_options["transformer_options"]["shape_slat"] = shape_latent - return IO.NodeOutput({"samples": latent, "coords": coords, "type": "trellis2"}, model) + output = {"samples": latent, "coords": coords, "type": "trellis2"} + if coord_counts is not None: + output["coord_counts"] = coord_counts + output["batch_index"] = [0] * batch_size + return IO.NodeOutput(output, model) class EmptyStructureLatentTrellis2(IO.ComfyNode): @@ -587,8 +842,15 @@ class EmptyStructureLatentTrellis2(IO.ComfyNode): def execute(cls, batch_size): in_channels = 8 resolution = 16 - latent = torch.randn(batch_size, in_channels, resolution, resolution, resolution) - return IO.NodeOutput({"samples": latent, "type": "trellis2"}) + generator = torch.Generator(device="cpu") + 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"} + if batch_size > 1: + output["batch_index"] = [0] * batch_size + return IO.NodeOutput(output) def simplify_fn(vertices, faces, colors=None, target=100000): if vertices.ndim == 3: