212 lines
8.9 KiB
Python
212 lines
8.9 KiB
Python
import torch
|
|
import os
|
|
|
|
import comfy.utils
|
|
import comfy.model_management
|
|
|
|
from torch import nn
|
|
from transformers import CLIPVisionModelWithProjection
|
|
|
|
MODELS_DIR = os.path.join(os.path.dirname(os.path.realpath(__file__)), "models")
|
|
|
|
# attention_channels
|
|
SD_V12_CHANNELS = [320] * 4 + [640] * 4 + [1280] * 4 + [1280] * 6 + [640] * 6 + [320] * 6 + [1280] * 2
|
|
SD_XL_CHANNELS = [640] * 8 + [1280] * 40 + [1280] * 60 + [640] * 12 + [1280] * 20
|
|
|
|
def get_filename_list(path):
|
|
return [f for f in os.listdir(path) if f.endswith('.bin')]
|
|
|
|
class ImageProjModel(nn.Module):
|
|
def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4):
|
|
super().__init__()
|
|
|
|
self.cross_attention_dim = cross_attention_dim
|
|
self.clip_extra_context_tokens = clip_extra_context_tokens
|
|
self.proj = nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim)
|
|
self.norm = nn.LayerNorm(cross_attention_dim)
|
|
|
|
def forward(self, image_embeds):
|
|
embeds = image_embeds
|
|
clip_extra_context_tokens = self.proj(embeds).reshape(-1, self.clip_extra_context_tokens, self.cross_attention_dim)
|
|
clip_extra_context_tokens = self.norm(clip_extra_context_tokens)
|
|
return clip_extra_context_tokens
|
|
|
|
class To_KV(nn.Module):
|
|
def __init__(self, cross_attention_dim):
|
|
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])
|
|
|
|
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"):
|
|
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
|
|
|
|
class IPAdapter(nn.Module):
|
|
def __init__(self, ipadapter_model, clip_embeddings_dim, dtype=torch.float16, device="cuda", num_tokens=4):
|
|
super().__init__()
|
|
|
|
self.ipadapter_model = ipadapter_model
|
|
self.device = device
|
|
#self.num_tokens = num_tokens
|
|
self.dtype = dtype
|
|
|
|
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.image_proj_model.load_state_dict(self.ipadapter_model["image_proj"])
|
|
self.ip_layers = To_KV(self.cross_attention_dim)
|
|
self.ip_layers.load_state_dict(self.ipadapter_model["ip_adapter"])
|
|
|
|
@torch.inference_mode()
|
|
def get_image_embeds(self, clip_image_embeds):
|
|
image_prompt_embeds = self.image_proj_model(clip_image_embeds)
|
|
uncond_image_prompt_embeds = self.image_proj_model(torch.zeros_like(clip_image_embeds))
|
|
return image_prompt_embeds, uncond_image_prompt_embeds
|
|
|
|
class IPAdapterModelLoader:
|
|
@classmethod
|
|
def INPUT_TYPES(s):
|
|
return {"required": { "ipadapter_file": (get_filename_list(MODELS_DIR), )}}
|
|
|
|
RETURN_TYPES = ("IPADAPTER",)
|
|
FUNCTION = "load_ipadapter_model"
|
|
|
|
CATEGORY = "loaders"
|
|
|
|
def load_ipadapter_model(self, ipadapter_file):
|
|
ckpt_path = os.path.join(MODELS_DIR, ipadapter_file)
|
|
|
|
model = comfy.utils.load_torch_file(ckpt_path, safe_load=True)
|
|
keys = model.keys()
|
|
|
|
if not "ip_adapter" in keys:
|
|
raise Exception("invalid IPAdapter model {}".format(ckpt_path))
|
|
|
|
return (model,)
|
|
|
|
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_TYPES = ("MODEL",)
|
|
FUNCTION = "apply_ipadapter"
|
|
CATEGORY = "conditioning/ipadapter"
|
|
|
|
def apply_ipadapter(self, model, ipadapter, clip_vision_output, weight):
|
|
self.dtype = model.model.diffusion_model.dtype
|
|
self.device = comfy.model_management.get_torch_device()
|
|
self.weight = weight
|
|
|
|
clip_embeddings = clip_vision_output.image_embeds.to(self.device, dtype=self.dtype)
|
|
clip_embeddings_dim = clip_embeddings.shape[1]
|
|
|
|
self.sdxl = clip_embeddings_dim == 1280 # is there a better way?
|
|
|
|
self.ipadapter = IPAdapter(
|
|
ipadapter,
|
|
clip_embeddings_dim,
|
|
dtype=self.dtype,
|
|
device=self.device
|
|
).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)
|
|
self.uncond_image_prompt_embeds = self.uncond_image_prompt_embeds.to(self.device, dtype=self.dtype)
|
|
self.cond_uncond_image_emb = None
|
|
|
|
work_model = model.clone()
|
|
|
|
if not self.sdxl:
|
|
number = 0 # index of to_kvs
|
|
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
|
|
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)
|
|
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
|
|
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
|
|
for index in range(10):
|
|
set_model_patch_replace(work_model, self.patch_forward(number), "middle", 0, index)
|
|
number += 1
|
|
|
|
return (work_model, )
|
|
|
|
def patch_forward(self, number):
|
|
device = self.device.type
|
|
def forward(n, context_attn2, value_attn2, extra_options):
|
|
org_dtype = n.dtype
|
|
with torch.autocast(device, dtype=self.dtype):
|
|
q = n
|
|
k = context_attn2
|
|
v = value_attn2
|
|
b, _, _ = q.shape
|
|
|
|
if self.cond_uncond_image_emb is None or self.cond_uncond_image_emb.shape[0] != b:
|
|
self.cond_uncond_image_emb = torch.cat([self.uncond_image_prompt_embeds.repeat(b//2, 1, 1), self.image_prompt_embeds.repeat(b//2, 1, 1)], dim=0)
|
|
|
|
# k, v for ip_adapter
|
|
ip_k = self.ipadapter.ip_layers.to_kvs[number*2](self.cond_uncond_image_emb)
|
|
ip_v = self.ipadapter.ip_layers.to_kvs[number*2+1](self.cond_uncond_image_emb)
|
|
|
|
q, k, v, ip_k, ip_v = map(
|
|
lambda t: t.view(b, -1, extra_options["n_heads"], extra_options["dim_head"]).transpose(1, 2),
|
|
(q, k, v, ip_k, ip_v),
|
|
)
|
|
|
|
out = nn.functional.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"])
|
|
|
|
# output of ip_adapter
|
|
ip_out = nn.functional.scaled_dot_product_attention(q, ip_k, ip_v, attn_mask=None, dropout_p=0.0, is_causal=False)
|
|
ip_out = ip_out.transpose(1, 2).reshape(b, -1, extra_options["n_heads"] * extra_options["dim_head"])
|
|
|
|
out = out + ip_out * self.weight
|
|
|
|
return out.to(dtype=org_dtype)
|
|
return forward
|
|
|
|
NODE_CLASS_MAPPINGS = {
|
|
"IPAdapterModelLoader": IPAdapterModelLoader,
|
|
"IPAdapterApply": IPAdapterApply,
|
|
}
|
|
|
|
NODE_DISPLAY_NAME_MAPPINGS = {
|
|
"IPAdapterModelLoader": "Load IPAdapter Model",
|
|
"IPAdapterApply": "Apply IPAdapter",
|
|
} |