From 57f1122cbc855079e2d2d9c04515276a9282b9b9 Mon Sep 17 00:00:00 2001 From: matt3o Date: Sun, 3 Sep 2023 20:11:39 +0200 Subject: [PATCH] initial work on ipadapter plus model --- IPAdapterPlus.py | 231 ++++++++++++++++++++++++++++++++++++----------- README.md | 2 + resampler.py | 121 +++++++++++++++++++++++++ 3 files changed, 302 insertions(+), 52 deletions(-) create mode 100644 resampler.py diff --git a/IPAdapterPlus.py b/IPAdapterPlus.py index c1956ad..6437225 100644 --- a/IPAdapterPlus.py +++ b/IPAdapterPlus.py @@ -5,8 +5,11 @@ import comfy.utils import comfy.model_management from torch import nn +import torch.nn.functional as F from transformers import CLIPVisionModelWithProjection +from .resampler import Resampler + MODELS_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models") # attention_channels @@ -36,46 +39,65 @@ class To_KV(nn.Module): super().__init__() channels = SD_XL_CHANNELS if cross_attention_dim == 2048 else SD_V12_CHANNELS - self.to_kvs = nn.ModuleList([torch.nn.Linear(cross_attention_dim, channel, bias=False) for channel in channels]) + self.to_kvs = nn.ModuleList([nn.Linear(cross_attention_dim, channel, bias=False) for channel in channels]) def load_state_dict(self, state_dict): - # input -> output -> middle for i, key in enumerate(state_dict.keys()): self.to_kvs[i].weight.data = state_dict[key] -''' -code from I'm using the code from: -https://github.com/comfyanonymous/ComfyUI/blob/a74c5dbf3764fa598b58da8c88da823aaf8364fa/comfy/model_patcher.py#L64 -''' -def set_model_patch_replace(model, patch, block_name, number, index, name="attn2"): +def set_model_patch_replace(model, patch_kwargs, key): to = model.model_options["transformer_options"] if "patches_replace" not in to: to["patches_replace"] = {} - if name not in to["patches_replace"]: - to["patches_replace"][name] = {} - to["patches_replace"][name][(block_name, number, index)] = patch + if "attn2" not in to["patches_replace"]: + to["patches_replace"]["attn2"] = {} + if key not in to["patches_replace"]["attn2"]: + patch = CrossAttentionPatch(**patch_kwargs) + to["patches_replace"]["attn2"][key] = patch + else: + to["patches_replace"]["attn2"][key].set_new_condition(**patch_kwargs) + +def attention(q, k, v, extra_options): + if not hasattr(F, "multi_head_attention_forward"): + q, k, v = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h=extra_options["n_heads"]), (q, k, v)) + sim = torch.einsum('b i d, b j d -> b i j', q, k) * (extra_options["dim_head"] ** -0.5) + sim = F.softmax(sim, dim=-1) + out = torch.einsum('b i j, b j d -> b i d', sim, v) + out = rearrange(out, '(b h) n d -> b n (h d)', h=extra_options["n_heads"]) + else: + b, _, _ = q.shape + q, k, v = map( + lambda t: t.view(b, -1, extra_options["n_heads"], extra_options["dim_head"]).transpose(1, 2), + (q, k, v), + ) + out = F.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False) + out = out.transpose(1, 2).reshape(b, -1, extra_options["n_heads"] * extra_options["dim_head"]) + return out class IPAdapter(nn.Module): - def __init__(self, ipadapter_model, clip_embeddings_dim, dtype=torch.float16, device="cuda", num_tokens=4): + def __init__(self, ipadapter_model, clip_embeddings, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4): super().__init__() self.ipadapter_model = ipadapter_model - self.device = device - #self.num_tokens = num_tokens - self.dtype = dtype + self.clip_embeddings = clip_embeddings + self.clip_embeddings_dim = clip_embeddings_dim self.cross_attention_dim = self.ipadapter_model["ip_adapter"]["1.to_k_ip.weight"].shape[1] - self.clip_extra_context_tokens = num_tokens - - self.image_proj_model = ImageProjModel( - cross_attention_dim=self.cross_attention_dim, - clip_embeddings_dim=clip_embeddings_dim, - clip_extra_context_tokens=self.clip_extra_context_tokens - ) + self.clip_extra_context_tokens = clip_extra_context_tokens + + self.image_proj_model = self.init_proj() self.image_proj_model.load_state_dict(self.ipadapter_model["image_proj"]) - self.ip_layers = To_KV(self.cross_attention_dim) + self.ip_layers = To_KV(cross_attention_dim) self.ip_layers.load_state_dict(self.ipadapter_model["ip_adapter"]) + + def init_proj(self): + image_proj_model = ImageProjModel( + cross_attention_dim=self.cross_attention_dim, + clip_embeddings_dim=self.clip_embeddings_dim, + clip_extra_context_tokens=self.clip_extra_context_tokens + ) + return image_proj_model @torch.inference_mode() def get_image_embeds(self, clip_image_embeds): @@ -83,6 +105,76 @@ class IPAdapter(nn.Module): uncond_image_prompt_embeds = self.image_proj_model(torch.zeros_like(clip_image_embeds)) return image_prompt_embeds, uncond_image_prompt_embeds +class IPAdapterPlus(IPAdapter): + def init_proj(self): + image_proj_model = Resampler( + dim=self.cross_attention_dim, + depth=4, + dim_head=64, + heads=12, + num_queries=self.clip_extra_context_tokens, + embedding_dim=self.clip_embeddings_dim, + output_dim=self.cross_attention_dim, + ff_mult=4 + ) + return image_proj_model + + @torch.inference_mode() + def get_image_embeds(self, clip_image_embeds): + image_prompt_embeds = self.image_proj_model(clip_image_embeds) + uncond_clip_image_embeds = self.clip_vision(torch.zeros_like(self.clip_embeddings), output_hidden_states=True).hidden_states[-2] + uncond_image_prompt_embeds = self.image_proj_model(uncond_clip_image_embeds) + return image_prompt_embeds, uncond_image_prompt_embeds + +class CrossAttentionPatch: + # forward for patching + def __init__(self, weight, ipadapter, dtype, number, cond, uncond, mask=None): + self.weights = [weight] + self.ipadapters = [ipadapter] + self.conds = [cond] + self.unconds = [uncond] + self.dtype = dtype + self.number = number + self.masks = [mask] + + def set_new_condition(self, weight, ipadapter, cond, uncond, dtype, number, mask=None): + self.weights.append(weight) + self.ipadapters.append(ipadapter) + self.conds.append(cond) + self.unconds.append(uncond) + self.masks.append(mask) + self.dtype = dtype + + def __call__(self, n, context_attn2, value_attn2, extra_options): + org_dtype = n.dtype + with torch.autocast("cuda", dtype=self.dtype): + q = n + k = context_attn2 + v = value_attn2 + b, _, _ = q.shape + + out = attention(q, k, v, extra_options) + + for weight, cond, uncond, ipadapter, mask in zip(self.weights, self.conds, self.unconds, self.ipadapters, self.masks): + uncond_cond = torch.cat([uncond.repeat(b//2, 1, 1), cond.repeat(b//2, 1, 1)], dim=0) + + # k, v for ip_adapter + ip_k = ipadapter.ip_layers.to_kvs[self.number*2](uncond_cond) + ip_v = ipadapter.ip_layers.to_kvs[self.number*2+1](uncond_cond) + + ip_out = attention(q, ip_k, ip_v, extra_options) + + if mask is not None: + mask_size = mask.shape[0] * mask.shape[1] + down_sample_rate = int((mask_size // 64 // out.shape[1]) ** (1/2)) + mask_downsample = torch.nn.functional.interpolate(mask.unsqueeze(0).unsqueeze(0), scale_factor= 1/8/down_sample_rate, mode="nearest").squeeze(0) + mask_downsample = mask_downsample.view(1,-1, 1).repeat(out.shape[0], 1, out.shape[2]) + ip_out = ip_out * mask_downsample + + out = out + ip_out * weight + + return out.to(dtype=org_dtype) + class IPAdapterModelLoader: @classmethod def INPUT_TYPES(s): @@ -107,31 +199,58 @@ class IPAdapterModelLoader: class IPAdapterApply: @classmethod def INPUT_TYPES(s): - return {"required": {"model": ("MODEL", ), - "ipadapter": ("IPADAPTER", ), - "clip_vision_output": ("CLIP_VISION_OUTPUT", ), - "weight": ("FLOAT", { "default": 1.0, "min": -1, "max": 3, "step": 0.1 }), - }} + return { + "required": { + "model": ("MODEL", ), + "ipadapter": ("IPADAPTER", ), + "clip_vision_output": ("CLIP_VISION_OUTPUT", ), + "weight": ("FLOAT", { "default": 1.0, "min": -1, "max": 3, "step": 0.1 }), + }, + "optional": { + "mask": ("MASK",), + } + } + RETURN_TYPES = ("MODEL",) FUNCTION = "apply_ipadapter" CATEGORY = "conditioning/ipadapter" - def apply_ipadapter(self, model, ipadapter, clip_vision_output, weight): + def apply_ipadapter(self, model, ipadapter, clip_vision_output, weight, mask=None): self.dtype = model.model.diffusion_model.dtype self.device = comfy.model_management.get_torch_device() self.weight = weight - + print(self.dtype) clip_embeddings = clip_vision_output.image_embeds.to(self.device, dtype=self.dtype) - clip_embeddings_dim = clip_embeddings.shape[1] + clip_embeddings_dim = clip_embeddings.shape[-1] - self.sdxl = clip_embeddings_dim == 1280 # is there a better way? + self.is_plus = "latents" in ipadapter["image_proj"] + cross_attention_dim = ipadapter["ip_adapter"]["1.to_k_ip.weight"].shape[1] - self.ipadapter = IPAdapter( - ipadapter, - clip_embeddings_dim, - dtype=self.dtype, - device=self.device - ).to(self.device, dtype=self.dtype) + if self.is_plus: + clip_extra_context_tokens = ipadapter["image_proj"]["latents"].shape[1] + else: + clip_extra_context_tokens = 4 # ipadapter["image_proj"]["proj.weight"].shape[0] <-- doesn't work? + + self.is_sdxl = cross_attention_dim == 2048 + + if self.is_plus: + self.ipadapter = IPAdapterPlus( + ipadapter, + clip_embeddings, + cross_attention_dim=cross_attention_dim, + clip_embeddings_dim=clip_embeddings_dim, + clip_extra_context_tokens=clip_extra_context_tokens + ) + else: + self.ipadapter = IPAdapter( + ipadapter, + clip_embeddings, + cross_attention_dim=cross_attention_dim, + clip_embeddings_dim=clip_embeddings_dim, + clip_extra_context_tokens=clip_extra_context_tokens + ) + + self.ipadapter.to(self.device, dtype=self.dtype) self.image_prompt_embeds, self.uncond_image_prompt_embeds = self.ipadapter.get_image_embeds(clip_embeddings) self.image_prompt_embeds = self.image_prompt_embeds.to(self.device, dtype=self.dtype) @@ -140,32 +259,40 @@ class IPAdapterApply: work_model = model.clone() - if not self.sdxl: - number = 0 # index of to_kvs + patch_kwargs = { + "number": 0, + "weight": self.weight, + "ipadapter": self.ipadapter, + "dtype": self.dtype, + "cond": self.image_prompt_embeds, + "uncond": self.uncond_image_prompt_embeds, + "mask": mask if mask is None else mask.to(self.device) + } + + if not self.is_sdxl: for id in [1,2,4,5,7,8]: # id of input_blocks that have cross attention - work_model.set_model_attn2_replace(self.patch_forward(number), "input", id) - number += 1 + set_model_patch_replace(work_model, patch_kwargs, ("input", id)) + patch_kwargs["number"] += 1 for id in [3,4,5,6,7,8,9,10,11]: # id of output_blocks that have cross attention - work_model.set_model_attn2_replace(self.patch_forward(number), "output", id) - number += 1 - work_model.set_model_attn2_replace(self.patch_forward(number), "middle", 0) + set_model_patch_replace(work_model, patch_kwargs, ("output", id)) + patch_kwargs["number"] += 1 + set_model_patch_replace(work_model, patch_kwargs, ("middle", 0)) else: - number = 0 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: - set_model_patch_replace(work_model, self.patch_forward(number), "input", id, index) - number += 1 + set_model_patch_replace(work_model, patch_kwargs, ("input", id, index)) + patch_kwargs["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: - set_model_patch_replace(work_model, self.patch_forward(number), "output", id, index) - number += 1 + set_model_patch_replace(work_model, patch_kwargs, ("output", id, index)) + patch_kwargs["number"] += 1 for index in range(10): - set_model_patch_replace(work_model, self.patch_forward(number), "middle", 0, index) - number += 1 + set_model_patch_replace(work_model, patch_kwargs, ("midlle", 0, index)) + patch_kwargs["number"] += 1 - return (work_model, ) + return (work_model,) def patch_forward(self, number): device = self.device.type diff --git a/README.md b/README.md index f3c94cc..a7e9a83 100644 --- a/README.md +++ b/README.md @@ -1,3 +1,5 @@ +**ALPHA STAGE! Not production ready!** + # ComfyUI_IPAdapter_Plus ComfyUI reference implementation for [IPAdapter](https://github.com/tencent-ailab/IP-Adapter/tree/6fb9d3554a5c774f41e187e8fdbc7b9a1db8c2e3) models. diff --git a/resampler.py b/resampler.py new file mode 100644 index 0000000..4521c8c --- /dev/null +++ b/resampler.py @@ -0,0 +1,121 @@ +# modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py +import math + +import torch +import torch.nn as nn + + +# FFN +def FeedForward(dim, mult=4): + inner_dim = int(dim * mult) + return nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, inner_dim, bias=False), + nn.GELU(), + nn.Linear(inner_dim, dim, bias=False), + ) + + +def reshape_tensor(x, heads): + bs, length, width = x.shape + #(bs, length, width) --> (bs, length, n_heads, dim_per_head) + x = x.view(bs, length, heads, -1) + # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head) + x = x.transpose(1, 2) + # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head) + x = x.reshape(bs, heads, length, -1) + return x + + +class PerceiverAttention(nn.Module): + def __init__(self, *, dim, dim_head=64, heads=8): + super().__init__() + self.scale = dim_head**-0.5 + self.dim_head = dim_head + self.heads = heads + inner_dim = dim_head * heads + + self.norm1 = nn.LayerNorm(dim) + self.norm2 = nn.LayerNorm(dim) + + self.to_q = nn.Linear(dim, inner_dim, bias=False) + self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False) + self.to_out = nn.Linear(inner_dim, dim, bias=False) + + + def forward(self, x, latents): + """ + Args: + x (torch.Tensor): image features + shape (b, n1, D) + latent (torch.Tensor): latent features + shape (b, n2, D) + """ + x = self.norm1(x) + latents = self.norm2(latents) + + b, l, _ = latents.shape + + q = self.to_q(latents) + kv_input = torch.cat((x, latents), dim=-2) + k, v = self.to_kv(kv_input).chunk(2, dim=-1) + + q = reshape_tensor(q, self.heads) + k = reshape_tensor(k, self.heads) + v = reshape_tensor(v, self.heads) + + # attention + scale = 1 / math.sqrt(math.sqrt(self.dim_head)) + weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards + weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype) + out = weight @ v + + out = out.permute(0, 2, 1, 3).reshape(b, l, -1) + + return self.to_out(out) + + +class Resampler(nn.Module): + def __init__( + self, + dim=1024, + depth=8, + dim_head=64, + heads=16, + num_queries=8, + embedding_dim=768, + output_dim=1024, + ff_mult=4, + ): + super().__init__() + + self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5) + + self.proj_in = nn.Linear(embedding_dim, dim) + + self.proj_out = nn.Linear(dim, output_dim) + self.norm_out = nn.LayerNorm(output_dim) + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, x): + + latents = self.latents.repeat(x.size(0), 1, 1) + + x = self.proj_in(x) + + for attn, ff in self.layers: + latents = attn(x, latents) + latents + latents = ff(latents) + latents + + latents = self.proj_out(latents) + return self.norm_out(latents) \ No newline at end of file