From b443f423b430429ad9e6c98c9318a7658eb2ad8d Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sun, 19 Apr 2026 21:26:48 -0500 Subject: [PATCH] Trellis2: slice cond half of x symmetrically under shape_rule pruning --- comfy/ldm/trellis2/model.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/trellis2/model.py b/comfy/ldm/trellis2/model.py index 3081d5919..fc8df3c2b 100644 --- a/comfy/ldm/trellis2/model.py +++ b/comfy/ldm/trellis2/model.py @@ -863,11 +863,12 @@ class Trellis2(nn.Module): else: # structure orig_bsz = x.shape[0] if shape_rule: - x = x[1].unsqueeze(0) - timestep = timestep[1].unsqueeze(0) + half = orig_bsz // 2 + x = x[half:] + timestep = timestep[half:] if timestep.shape[0] > 1 else timestep out = self.structure_model(x, timestep, context if not shape_rule else cond) if shape_rule: - out = out.repeat(orig_bsz, 1, 1, 1, 1) + out = out.repeat(2, 1, 1, 1, 1) if not_struct_mode: out = out.feats