diff --git a/comfy/ldm/trellis2/model.py b/comfy/ldm/trellis2/model.py index 76dbacc93..f61c50629 100644 --- a/comfy/ldm/trellis2/model.py +++ b/comfy/ldm/trellis2/model.py @@ -829,21 +829,6 @@ class Trellis2(nn.Module): 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"]: @@ -878,16 +863,6 @@ class Trellis2(nn.Module): coord_batches.append(coords_rep) 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( feats=torch.cat(feat_batches, dim=0), coords=torch.cat(coord_batches, dim=0).to(torch.int32), @@ -972,16 +947,6 @@ class Trellis2(nn.Module): 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 @@ -1060,20 +1025,6 @@ class Trellis2(nn.Module): 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: @@ -1088,10 +1039,6 @@ class Trellis2(nn.Module): 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 @@ -1111,23 +1058,8 @@ class Trellis2(nn.Module): 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: @@ -1146,28 +1078,10 @@ class Trellis2(nn.Module): 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: 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 3967fba1b..7251aa799 100644 --- a/comfy/sample.py +++ b/comfy/sample.py @@ -10,18 +10,37 @@ 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()): + if noise_inds is None: + 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.set_state(base_state.clone()) + local_generator.set_state(index_states[int(noise_index)].clone()) 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, layout=latent_image.layout, generator=local_generator, 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) if noise_inds is None: diff --git a/comfy_extras/nodes_trellis2.py b/comfy_extras/nodes_trellis2.py index 26cb135e7..621cc9586 100644 --- a/comfy_extras/nodes_trellis2.py +++ b/comfy_extras/nodes_trellis2.py @@ -148,18 +148,6 @@ def split_batched_sparse_latent(samples, coords, coord_counts): 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. @@ -410,14 +398,6 @@ class Trellis2UpsampleCascade(IO.ComfyNode): ) 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) @@ -435,27 +415,18 @@ class Trellis2UpsampleCascade(IO.ComfyNode): 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 = [] + final_coords_list = [] output_resolutions = [] + output_coord_counts = [] for batch_index, (feats_i, coords_i) in enumerate(items): feats_i = feats_i.to(device) 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[:, 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)) + output_coord_counts.append(int(final_coords_i.shape[0])) - 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),) + return IO.NodeOutput({ + "coords": torch.cat(final_coords_list, dim=0), + "coord_counts": torch.tensor(output_coord_counts, dtype=torch.int64), + "resolutions": torch.tensor(output_resolutions, dtype=torch.int64), + },) 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) @@ -568,7 +534,6 @@ class Trellis2Conditioning(IO.ComfyNode): cond_512_list = [] cond_1024_list = [] - composite_trace = [] for b in range(batch_size): item_image = image[b] @@ -623,14 +588,6 @@ 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) @@ -642,19 +599,6 @@ 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}]] @@ -680,12 +624,20 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode): def execute(cls, structure_or_coords, model): # to accept the upscaled coords is_512_pass = False + coord_counts = None + coord_resolutions = None if hasattr(structure_or_coords, "data") and structure_or_coords.data.ndim == 4: decoded = structure_or_coords.data.unsqueeze(1) coords = torch.argwhere(decoded.bool())[:, [0, 2, 3, 4]].int() 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: coords = structure_or_coords.int() is_512_pass = False @@ -693,7 +645,15 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode): else: raise ValueError(f"Invalid input to EmptyShapeLatent: {type(structure_or_coords)}") 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: coord_counts = None latent = torch.randn(1, in_channels, coords.shape[0], 1) @@ -706,17 +666,6 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode): 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() @@ -729,6 +678,8 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode): model.model_options["transformer_options"]["coords"] = coords if coord_counts is not None: 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: model.model_options["transformer_options"]["generation_mode"] = "shape_generation_512" else: @@ -736,6 +687,8 @@ class EmptyShapeLatentTrellis2(IO.ComfyNode): output = {"samples": latent, "coords": coords, "type": "trellis2"} if coord_counts is not None: output["coord_counts"] = coord_counts + if coord_resolutions is not None: + output["coord_resolutions"] = coord_resolutions output["batch_index"] = [0] * batch_size return IO.NodeOutput(output, model) @@ -759,15 +712,28 @@ class EmptyTextureLatentTrellis2(IO.ComfyNode): @classmethod def execute(cls, structure_or_coords, shape_latent, model): channels = 32 + coord_counts = None if hasattr(structure_or_coords, "data") and structure_or_coords.data.ndim == 4: decoded = structure_or_coords.data.unsqueeze(1) 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: coords = structure_or_coords.int() 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.shape[0] != batch_size: raise ValueError( @@ -791,19 +757,6 @@ class EmptyTextureLatentTrellis2(IO.ComfyNode): 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() @@ -842,11 +795,7 @@ class EmptyStructureLatentTrellis2(IO.ComfyNode): def execute(cls, batch_size): in_channels = 8 resolution = 16 - 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}) + latent = torch.randn(1, in_channels, resolution, resolution, resolution).repeat(batch_size, 1, 1, 1, 1) output = {"samples": latent, "type": "trellis2"} if batch_size > 1: output["batch_index"] = [0] * batch_size