From 1e05d1a8cc77dd873ebcf8a3161fdbccbca05472 Mon Sep 17 00:00:00 2001 From: asagi4 <130366179+asagi4@users.noreply.github.com> Date: Fri, 30 May 2025 15:25:30 +0300 Subject: [PATCH] Refactor 2: new cond amount can be calculated from num_conds --- prompt_control/attention_couple_ppm.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/prompt_control/attention_couple_ppm.py b/prompt_control/attention_couple_ppm.py index 6d1b439..965d916 100644 --- a/prompt_control/attention_couple_ppm.py +++ b/prompt_control/attention_couple_ppm.py @@ -17,7 +17,6 @@ log = logging.getLogger("comfyui-prompt-control") COND = 0 UNCOND = 1 -COND_UNCOND_COUPLE = "cond_or_uncond_couple" DEBUG_KEYS = {} @@ -155,26 +154,24 @@ class AttentionCoupleHook(TransformerOptionsHook): qs.append(q_target.repeat(self.num_conds, 1, 1)) ks.append(torch.cat([k_target * self.base_strength, conds_k_tensor], dim=0)) vs.append(torch.cat([v_target * self.base_strength, conds_v_tensor], dim=0)) - cond_or_uncond_couple.extend(itertools.repeat(COND, self.num_conds)) qs = torch.cat(qs, dim=0) ks = torch.cat(ks, dim=0) vs = torch.cat(vs, dim=0) - extra_options[COND_UNCOND_COUPLE] = cond_or_uncond_couple - return qs, ks, vs return q, k, v def attn2_output_patch(self, out, extra_options): - cond_or_uncond = extra_options[COND_UNCOND_COUPLE] + cond_or_uncond = extra_options["cond_or_uncond"] bs = self.batch_size mask_downsample = get_mask(self.mask, self.batch_size, out.shape[1], extra_options) outputs = [] cond_outputs = [] i_cond = 0 - for i, cond_type in enumerate(cond_or_uncond): + # out has been extended to shape [num_conds*batch_size, TOKENS, N] + for i in range(len(cond_or_uncond) * self.num_conds): pos, next_pos = i * bs, (i + 1) * bs pos_cond, next_pos_cond = i_cond * bs, (i_cond + 1) * bs @@ -182,6 +179,7 @@ class AttentionCoupleHook(TransformerOptionsHook): cond_outputs.append(masked_output) i_cond += 1 + # cond_outputs is [num_conds, bs, tokens, N] + # output needs to be [bs, tokens, N] cond_output = torch.stack(cond_outputs).sum(0) - outputs.append(cond_output) - return torch.cat(outputs, dim=0) + return cond_output