Fix Trellis PR review regressions

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

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:

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