From 20125bf9394b1bc98ef3228277a31a3a52c72fc2 Mon Sep 17 00:00:00 2001 From: matt3o Date: Wed, 8 May 2024 16:10:20 +0200 Subject: [PATCH] add PuLID compatibility --- CrossAttentionPatch.py | 187 +++++++++++++++++++++++++++++++++++++++++ IPAdapterPlus.py | 35 +++++--- 2 files changed, 211 insertions(+), 11 deletions(-) diff --git a/CrossAttentionPatch.py b/CrossAttentionPatch.py index ac2d0fe..d502521 100644 --- a/CrossAttentionPatch.py +++ b/CrossAttentionPatch.py @@ -4,6 +4,192 @@ import torch.nn.functional as F from comfy.ldm.modules.attention import optimized_attention from .utils import tensor_to_size +class Attn2Replace: + def __init__(self, callback=None, **kwargs): + self.callback = [callback] + self.kwargs = [kwargs] + + def add(self, callback, **kwargs): + self.callback.append(callback) + self.kwargs.append(kwargs) + + for key, value in kwargs.items(): + setattr(self, key, value) + + def __call__(self, q, k, v, extra_options): + dtype = q.dtype + out = optimized_attention(q, k, v, extra_options["n_heads"]) + sigma = extra_options["sigmas"].detach().cpu()[0].item() if 'sigmas' in extra_options else 999999999.9 + + for i, callback in enumerate(self.callback): + if sigma <= self.kwargs[i]["sigma_start"] and sigma >= self.kwargs[i]["sigma_end"]: + out = out + callback(out, q, k, v, extra_options, **self.kwargs[i]) + + return out.to(dtype=dtype) + +def ipadapter_attention(out, q, k, v, extra_options, module_key='', ipadapter=None, weight=1.0, cond=None, cond_alt=None, uncond=None, weight_type="linear", mask=None, sigma_start=0.0, sigma_end=1.0, unfold_batch=False, embeds_scaling='V only', **kwargs): + dtype = q.dtype + cond_or_uncond = extra_options["cond_or_uncond"] + block_type = extra_options["block"][0] + #block_id = extra_options["block"][1] + t_idx = extra_options["transformer_index"] + layers = 11 if '101_to_k_ip' in ipadapter.ip_layers.to_kvs else 16 + k_key = module_key + "_to_k_ip" + v_key = module_key + "_to_v_ip" + + # extra options for AnimateDiff + ad_params = extra_options['ad_params'] if "ad_params" in extra_options else None + + b = q.shape[0] + seq_len = q.shape[1] + batch_prompt = b // len(cond_or_uncond) + _, _, oh, ow = extra_options["original_shape"] + + if weight_type == 'ease in': + weight = weight * (0.05 + 0.95 * (1 - t_idx / layers)) + elif weight_type == 'ease out': + weight = weight * (0.05 + 0.95 * (t_idx / layers)) + elif weight_type == 'ease in-out': + weight = weight * (0.05 + 0.95 * (1 - abs(t_idx - (layers/2)) / (layers/2))) + elif weight_type == 'reverse in-out': + weight = weight * (0.05 + 0.95 * (abs(t_idx - (layers/2)) / (layers/2))) + elif weight_type == 'weak input' and block_type == 'input': + weight = weight * 0.2 + elif weight_type == 'weak middle' and block_type == 'middle': + weight = weight * 0.2 + elif weight_type == 'weak output' and block_type == 'output': + weight = weight * 0.2 + elif weight_type == 'strong middle' and (block_type == 'input' or block_type == 'output'): + weight = weight * 0.2 + elif isinstance(weight, dict): + if t_idx not in weight: + return 0 + + weight = weight[t_idx] + + if cond_alt is not None and t_idx in cond_alt: + cond = cond_alt[t_idx] + del cond_alt + + if unfold_batch: + # Check AnimateDiff context window + if ad_params is not None and ad_params["sub_idxs"] is not None: + if isinstance(weight, torch.Tensor): + weight = tensor_to_size(weight, ad_params["full_length"]) + weight = torch.Tensor(weight[ad_params["sub_idxs"]]) + if torch.all(weight == 0): + return 0 + weight = weight.repeat(len(cond_or_uncond), 1, 1) # repeat for cond and uncond + elif weight == 0: + return 0 + + # if image length matches or exceeds full_length get sub_idx images + if cond.shape[0] >= ad_params["full_length"]: + cond = torch.Tensor(cond[ad_params["sub_idxs"]]) + uncond = torch.Tensor(uncond[ad_params["sub_idxs"]]) + # otherwise get sub_idxs images + else: + cond = tensor_to_size(cond, ad_params["full_length"]) + uncond = tensor_to_size(uncond, ad_params["full_length"]) + cond = cond[ad_params["sub_idxs"]] + uncond = uncond[ad_params["sub_idxs"]] + else: + if isinstance(weight, torch.Tensor): + weight = tensor_to_size(weight, batch_prompt) + if torch.all(weight == 0): + return 0 + weight = weight.repeat(len(cond_or_uncond), 1, 1) # repeat for cond and uncond + elif weight == 0: + return 0 + + cond = tensor_to_size(cond, batch_prompt) + uncond = tensor_to_size(uncond, batch_prompt) + + k_cond = ipadapter.ip_layers.to_kvs[k_key](cond) + k_uncond = ipadapter.ip_layers.to_kvs[k_key](uncond) + v_cond = ipadapter.ip_layers.to_kvs[v_key](cond) + v_uncond = ipadapter.ip_layers.to_kvs[v_key](uncond) + else: + # TODO: should we always convert the weights to a tensor? + if isinstance(weight, torch.Tensor): + weight = tensor_to_size(weight, batch_prompt) + if torch.all(weight == 0): + return 0 + weight = weight.repeat(len(cond_or_uncond), 1, 1) # repeat for cond and uncond + elif weight == 0: + return 0 + + k_cond = ipadapter.ip_layers.to_kvs[k_key](cond).repeat(batch_prompt, 1, 1) + k_uncond = ipadapter.ip_layers.to_kvs[k_key](uncond).repeat(batch_prompt, 1, 1) + v_cond = ipadapter.ip_layers.to_kvs[v_key](cond).repeat(batch_prompt, 1, 1) + v_uncond = ipadapter.ip_layers.to_kvs[v_key](uncond).repeat(batch_prompt, 1, 1) + + ip_k = torch.cat([(k_cond, k_uncond)[i] for i in cond_or_uncond], dim=0) + ip_v = torch.cat([(v_cond, v_uncond)[i] for i in cond_or_uncond], dim=0) + + if embeds_scaling == 'K+mean(V) w/ C penalty': + scaling = float(ip_k.shape[2]) / 1280.0 + weight = weight * scaling + ip_k = ip_k * weight + ip_v_mean = torch.mean(ip_v, dim=1, keepdim=True) + ip_v = (ip_v - ip_v_mean) + ip_v_mean * weight + out_ip = optimized_attention(q, ip_k, ip_v, extra_options["n_heads"]) + del ip_v_mean + elif embeds_scaling == 'K+V w/ C penalty': + scaling = float(ip_k.shape[2]) / 1280.0 + weight = weight * scaling + ip_k = ip_k * weight + ip_v = ip_v * weight + out_ip = optimized_attention(q, ip_k, ip_v, extra_options["n_heads"]) + elif embeds_scaling == 'K+V': + ip_k = ip_k * weight + ip_v = ip_v * weight + out_ip = optimized_attention(q, ip_k, ip_v, extra_options["n_heads"]) + else: + #ip_v = ip_v * weight + out_ip = optimized_attention(q, ip_k, ip_v, extra_options["n_heads"]) + out_ip = out_ip * weight # I'm doing this to get the same results as before + + if mask is not None: + mask_h = oh / math.sqrt(oh * ow / seq_len) + mask_h = int(mask_h) + int((seq_len % int(mask_h)) != 0) + mask_w = seq_len // mask_h + + # check if using AnimateDiff and sliding context window + if (mask.shape[0] > 1 and ad_params is not None and ad_params["sub_idxs"] is not None): + # if mask length matches or exceeds full_length, get sub_idx masks + if mask.shape[0] >= ad_params["full_length"]: + mask = torch.Tensor(mask[ad_params["sub_idxs"]]) + mask = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="bilinear").squeeze(1) + else: + mask = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="bilinear").squeeze(1) + mask = tensor_to_size(mask, ad_params["full_length"]) + mask = mask[ad_params["sub_idxs"]] + else: + mask = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="bilinear").squeeze(1) + mask = tensor_to_size(mask, batch_prompt) + + mask = mask.repeat(len(cond_or_uncond), 1, 1) + mask = mask.view(mask.shape[0], -1, 1).repeat(1, 1, out.shape[2]) + + # covers cases where extreme aspect ratios can cause the mask to have a wrong size + mask_len = mask_h * mask_w + if mask_len < seq_len: + pad_len = seq_len - mask_len + pad1 = pad_len // 2 + pad2 = pad_len - pad1 + mask = F.pad(mask, (0, 0, pad1, pad2), value=0.0) + elif mask_len > seq_len: + crop_start = (mask_len - seq_len) // 2 + mask = mask[:, crop_start:crop_start+seq_len, :] + + out_ip = out_ip * mask + + #out = out + out_ip + + return out_ip.to(dtype=dtype) + +""" class CrossAttentionPatch: # forward for patching def __init__(self, ipadapter=None, number=0, weight=1.0, cond=None, cond_alt=None, uncond=None, weight_type="linear", mask=None, sigma_start=0.0, sigma_end=1.0, unfold_batch=False, embeds_scaling='V only'): @@ -199,3 +385,4 @@ class CrossAttentionPatch: out = out + out_ip return out.to(dtype=dtype) +""" \ No newline at end of file diff --git a/IPAdapterPlus.py b/IPAdapterPlus.py index 527a992..28bed51 100644 --- a/IPAdapterPlus.py +++ b/IPAdapterPlus.py @@ -17,7 +17,7 @@ except ImportError: import torchvision.transforms as T from .image_proj_models import MLPProjModel, MLPProjModelFaceId, ProjModelFaceIdPlus, Resampler, ImageProjModel -from .CrossAttentionPatch import CrossAttentionPatch +from .CrossAttentionPatch import Attn2Replace, ipadapter_attention from .utils import ( encode_image_masked, tensor_to_size, @@ -135,15 +135,22 @@ class To_KV(nn.Module): self.to_kvs[key.replace(".weight", "").replace(".", "_")].weight.data = value def set_model_patch_replace(model, patch_kwargs, key): - to = model.model_options["transformer_options"] + to = model.model_options["transformer_options"].copy() if "patches_replace" not in to: to["patches_replace"] = {} + else: + to["patches_replace"] = to["patches_replace"].copy() + if "attn2" not in to["patches_replace"]: to["patches_replace"]["attn2"] = {} - if key not in to["patches_replace"]["attn2"]: - to["patches_replace"]["attn2"][key] = CrossAttentionPatch(**patch_kwargs) else: - to["patches_replace"]["attn2"][key].set_new_condition(**patch_kwargs) + to["patches_replace"]["attn2"] = to["patches_replace"]["attn2"].copy() + + if key not in to["patches_replace"]["attn2"]: + to["patches_replace"]["attn2"][key] = Attn2Replace(ipadapter_attention, **patch_kwargs) + model.model_options["transformer_options"] = to + else: + to["patches_replace"]["attn2"][key].add(ipadapter_attention, **patch_kwargs) def ipadapter_execute(model, ipadapter, @@ -365,7 +372,6 @@ def ipadapter_execute(model, patch_kwargs = { "ipadapter": ipa, - "number": 0, "weight": weight, "cond": cond, "cond_alt": cond_alt, @@ -378,28 +384,35 @@ def ipadapter_execute(model, "embeds_scaling": embeds_scaling, } + number = 0 if not is_sdxl: for id in [1,2,4,5,7,8]: # id of input_blocks that have cross attention + patch_kwargs["module_key"] = str(number*2+1) set_model_patch_replace(model, patch_kwargs, ("input", id)) - patch_kwargs["number"] += 1 + number += 1 for id in [3,4,5,6,7,8,9,10,11]: # id of output_blocks that have cross attention + patch_kwargs["module_key"] = str(number*2+1) set_model_patch_replace(model, patch_kwargs, ("output", id)) - patch_kwargs["number"] += 1 + number += 1 + patch_kwargs["module_key"] = str(number*2+1) set_model_patch_replace(model, patch_kwargs, ("middle", 0)) else: for id in [4,5,7,8]: # id of input_blocks that have cross attention block_indices = range(2) if id in [4, 5] else range(10) # transformer_depth for index in block_indices: + patch_kwargs["module_key"] = str(number*2+1) set_model_patch_replace(model, patch_kwargs, ("input", id, index)) - patch_kwargs["number"] += 1 + number += 1 for id in range(6): # id of output_blocks that have cross attention block_indices = range(2) if id in [3, 4, 5] else range(10) # transformer_depth for index in block_indices: + patch_kwargs["module_key"] = str(number*2+1) set_model_patch_replace(model, patch_kwargs, ("output", id, index)) - patch_kwargs["number"] += 1 + number += 1 for index in range(10): + patch_kwargs["module_key"] = str(number*2+1) set_model_patch_replace(model, patch_kwargs, ("middle", 0, index)) - patch_kwargs["number"] += 1 + number += 1 return (model, image)