From 70511a9a91d89efeca0cc3d11e7e4041c5e2438c Mon Sep 17 00:00:00 2001 From: John Pollock Date: Sun, 19 Apr 2026 21:38:45 -0500 Subject: [PATCH] Trellis2: guard structure shape_rule pruning to CFG batches --- comfy/ldm/trellis2/model.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/comfy/ldm/trellis2/model.py b/comfy/ldm/trellis2/model.py index fc8df3c2b..1c5d6c3ec 100644 --- a/comfy/ldm/trellis2/model.py +++ b/comfy/ldm/trellis2/model.py @@ -862,12 +862,12 @@ class Trellis2(nn.Module): out = self.shape2txt(x_st, t_eval, c_eval) else: # structure orig_bsz = x.shape[0] - if shape_rule: + 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, context if not shape_rule else cond) - if shape_rule: + 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) if not_struct_mode: