Conditional padding

README.md

@@ -25,9 +25,6 @@ This has currently only been tested with 1.5 based models.
 - `middle_block`: Focuses attention on the middle layers.
 - `output_blocks`: Focuses attention on the decoder layers.
 
-> [!NOTE]
-> `positive` and `reference_cond` must have equal amount of chunks. Mismatch will cause `RuntimeError: Sizes of tensors must match except in dimension 0.`
-
 > [!TIP]
 > In order to get the best results (matching legacy functionality), you must use the same prompt for `positive` and `reference_cond`.
 >

nodes.py

@@ -2,7 +2,7 @@ import comfy
 import torch
 
 from .utils.attention_functions import VisualStyleProcessor
-from .utils.cond_functions import copy_cond
+from .utils.cond_functions import cat_cond
 
 
 class ApplyVisualStyle:
@@ -11,6 +11,7 @@ class ApplyVisualStyle:
         return {
             "required": {
                 "model": ("MODEL",),
+                "clip": ("CLIP", ),
                 "vae": ("VAE", ),
                 "reference_image": ("IMAGE",),
                 "reference_cond": ("CONDITIONING",),
@@ -54,6 +55,7 @@ class ApplyVisualStyle:
     def apply_visual_style_prompt(
         self,
         model: comfy.model_patcher.ModelPatcher,
+        clip,
         vae,
         reference_image,
         reference_cond,
@@ -66,10 +68,6 @@ class ApplyVisualStyle:
         output_blocks,
         init_image = None
     ):
-        positive = copy_cond(positive)
-        reference_cond = copy_cond(reference_cond)
-        negative = copy_cond(negative)
-
         reference_latent = vae.encode(reference_image[:,:,:,:3])
 
         block_choices = self.get_block_choices(input_blocks, middle_block, output_blocks)
@@ -84,8 +82,8 @@ class ApplyVisualStyle:
                     processor = VisualStyleProcessor(m, enabled=is_enabled)
                     setattr(m, 'forward', processor)
 
-        positive[0][0] = torch.cat([reference_cond[0][0], positive[0][0]])
-        negative[0][0] = torch.cat([negative[0][0]] * 2)
+        positive_cat = cat_cond(clip, reference_cond, positive)
+        negative_cat = cat_cond(clip, negative, negative)
 
         latents = torch.zeros_like(reference_latent)
         latents = torch.cat([latents] * 2)
@@ -99,7 +97,7 @@ class ApplyVisualStyle:
 
         denoise_mask[::2] = 0.
 
-        return (model, positive, negative, {"samples": latents, "noise_mask": denoise_mask})
+        return (model, positive_cat, negative_cat, {"samples": latents, "noise_mask": denoise_mask})
 
 NODE_CLASS_MAPPINGS = {
     "ApplyVisualStyle": ApplyVisualStyle,

utils/cond_functions.py

@@ -1,7 +1,9 @@
 import torch
 
+EMPTY_PROMPT = ""
 
-def copy_cond(cond: list[list[torch.Tensor | dict]], device="cpu"):
+
+def copy_cond(cond: list[tuple[torch.Tensor, dict]], device="cpu") -> list[list[torch.Tensor | dict]]:
     cond_copy = [[]]
     t_copy = cond_copy[0]
     for t in cond:
@@ -11,3 +13,25 @@ def copy_cond(cond: list[list[torch.Tensor | dict]], device="cpu"):
             if "pooled_output" in t[1]:
                 t_copy[1]["pooled_output"] = t[1]["pooled_output"].detach().clone().to(device=device)
     return cond_copy
+
+
+def cat_cond(clip, *conds: list[list[torch.Tensor | dict]], device="cpu") -> list[list[torch.Tensor | dict]]:
+    output = copy_cond(conds[0], device=device)
+    cond_tensors = list(map(lambda c: c[0][0], conds))
+
+    chunk_size = 77
+    chunks = list(map(lambda c: c.shape[1] // chunk_size, cond_tensors))
+    chunk_min, chunk_max = min(chunks), max(chunks)
+
+    if chunk_min != chunk_max:
+        empty_tokens = clip.tokenize(EMPTY_PROMPT)
+        empty_cond: torch.Tensor = clip.encode_from_tokens(empty_tokens, return_pooled=False)
+
+        for i, cond in enumerate(cond_tensors):
+            chunk_count = chunks[i]
+            empty_pad = empty_cond.repeat(1, chunk_max - chunk_count, 1)
+            cond_tensors[i] = torch.cat([cond.to(device=device), empty_pad.to(device=device)], dim=1)
+
+    output[0][0] = torch.cat(cond_tensors).to(device=device)
+
+    return output