diff --git a/README.md b/README.md index 0544b99..4967440 100644 --- a/README.md +++ b/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`. > diff --git a/nodes.py b/nodes.py index 88855d5..79cb732 100644 --- a/nodes.py +++ b/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, diff --git a/utils/cond_functions.py b/utils/cond_functions.py index 0251a54..efcd38c 100644 --- a/utils/cond_functions.py +++ b/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 diff --git a/workflows/workflow_simple.png b/workflows/workflow_simple.png index 6de266a..ef1a724 100644 Binary files a/workflows/workflow_simple.png and b/workflows/workflow_simple.png differ