diff --git a/SteerableMotion.py b/SteerableMotion.py index 80774ca..bb77631 100644 --- a/SteerableMotion.py +++ b/SteerableMotion.py @@ -301,8 +301,13 @@ class BatchCreativeInterpolationNode: # ADD BUFFER TO KEYFRAME POSITIONS if buffer > 0: + # add front buffer keyframe_positions = [position + buffer - 1 for position in keyframe_positions] keyframe_positions.insert(0, 0) + # add end buffer + last_position_with_buffer = keyframe_positions[-1] + buffer - 1 + keyframe_positions.append(last_position_with_buffer) + # GET STRENGTH VALUES @@ -334,13 +339,11 @@ class BatchCreativeInterpolationNode: interpolation = "ease-in-out" # strength_from = strength_to = 1.0 - if i == 0: # buffer + if i == 0: # buffer - if buffer > 0: # First image with buffer - image = images[0] - strength_from = strength_to = strength_values[0][1] - else: - continue # Skip first image without buffer + image = images[0] + strength_from = strength_to = strength_values[0][1] + batch_index_from = 0 batch_index_to_excl = buffer weights, frame_numbers = calculate_weights(batch_index_from, batch_index_to_excl, strength_from, strength_to, interpolation, False, last_key_frame_position, i, len(keyframe_positions), buffer) @@ -360,7 +363,7 @@ class BatchCreativeInterpolationNode: weights, frame_numbers = calculate_weights(batch_index_from, batch_index_to_excl, mid_strength, end_strength, interpolation, False, last_key_frame_position, i, len(keyframe_positions), buffer) # interpolation = "ease-in" - elif i == len(images): # last image + elif i == len(keyframe_positions) - 2: # last image # GET IMAGE AND KEYFRAME INFLUENCE VALUES image = images[i-1] @@ -375,7 +378,14 @@ class BatchCreativeInterpolationNode: batch_index_to_excl = keyframe_position weights, frame_numbers = calculate_weights(batch_index_from, batch_index_to_excl, start_strength, mid_strength, interpolation, False, last_key_frame_position, i, len(keyframe_positions), buffer) - # interpolation = "ease-out" + # interpolation = "ease-out" + + elif i == len(keyframe_positions) - 1: + image = images[i-2] + strength_from = strength_to = strength_values[i-1][1] + batch_index_from = keyframe_positions[i-1] + batch_index_to_excl = last_key_frame_position + weights, frame_numbers = calculate_weights(batch_index_from, batch_index_to_excl, strength_from, strength_to, interpolation, False, last_key_frame_position, i, len(keyframe_positions), buffer) else: # middle images @@ -430,7 +440,7 @@ class BatchCreativeInterpolationNode: prepped_image = prep_image(image=image.unsqueeze(0), interpolation="LANCZOS", crop_position="pad", sharpening=0.0)[0] mask = create_mask_batch(last_key_frame_position, ipa_weights, ipa_frame_numbers) embed, = ipadapter_encoder.preprocess(clip_vision, prepped_image, True, 0.0, 1.0) - model, = ipadapter_application.apply_ipadapter(ipadapter=ipadapter, model=model, weight=1.0, image=None, weight_type="original", + model,_,_ = ipadapter_application.apply_ipadapter(ipadapter=ipadapter, model=model, weight=1.0, image=None, weight_type="original", noise=ipadapter_noise, embeds=embed, attn_mask=mask, start_at=ipadapter_start_at, end_at=ipadapter_end_at, unfold_batch=True) all_ipa_frame_numbers.append(ipa_frame_numbers) all_ipa_weights.append(ipa_weights) diff --git a/imports/IPAdapterPlus.py b/imports/IPAdapterPlus.py index 573bdc3..a9aab68 100644 --- a/imports/IPAdapterPlus.py +++ b/imports/IPAdapterPlus.py @@ -1,29 +1,69 @@ import torch -import contextlib import os import math import comfy.utils import comfy.model_management -from comfy.clip_vision import clip_preprocess +from comfy.clip_vision import clip_preprocess, Output from comfy.ldm.modules.attention import optimized_attention +from nodes import MAX_RESOLUTION import folder_paths -from torch import nn +import torch.nn as nn from PIL import Image import torch.nn.functional as F import torchvision.transforms as TT + + # set the models directory backward compatible GLOBAL_MODELS_DIR = os.path.join(folder_paths.models_dir, "ipadapter") MODELS_DIR = GLOBAL_MODELS_DIR if os.path.isdir(GLOBAL_MODELS_DIR) else os.path.join(os.path.dirname(os.path.realpath(__file__)), "models") if "ipadapter" not in folder_paths.folder_names_and_paths: - folder_paths.folder_names_and_paths["ipadapter"] = ([MODELS_DIR], folder_paths.supported_pt_extensions) + current_paths = [MODELS_DIR] else: - folder_paths.folder_names_and_paths["ipadapter"][1].update(folder_paths.supported_pt_extensions) + current_paths, _ = folder_paths.folder_names_and_paths["ipadapter"] +folder_paths.folder_names_and_paths["ipadapter"] = (current_paths, folder_paths.supported_pt_extensions) + +INSIGHTFACE_DIR = os.path.join(folder_paths.models_dir, "insightface") + +class FacePerceiverResamplerImport(torch.nn.Module): + def __init__( + self, + *, + dim=768, + depth=4, + dim_head=64, + heads=16, + embedding_dim=1280, + output_dim=768, + ff_mult=4, + ): + super().__init__() + + self.proj_in = torch.nn.Linear(embedding_dim, dim) + self.proj_out = torch.nn.Linear(dim, output_dim) + self.norm_out = torch.nn.LayerNorm(output_dim) + self.layers = torch.nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + torch.nn.ModuleList( + [ + PerceiverAttentionImport(dim=dim, dim_head=dim_head, heads=heads), + FeedForwardImport(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, latents, x): + 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) class MLPProjModelImport(torch.nn.Module): - """SD model with image prompt""" def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024): super().__init__() @@ -38,6 +78,59 @@ class MLPProjModelImport(torch.nn.Module): clip_extra_context_tokens = self.proj(image_embeds) return clip_extra_context_tokens +class MLPProjModelFaceIdImport(torch.nn.Module): + def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, num_tokens=4): + super().__init__() + + self.cross_attention_dim = cross_attention_dim + self.num_tokens = num_tokens + + self.proj = torch.nn.Sequential( + torch.nn.Linear(id_embeddings_dim, id_embeddings_dim*2), + torch.nn.GELU(), + torch.nn.Linear(id_embeddings_dim*2, cross_attention_dim*num_tokens), + ) + self.norm = torch.nn.LayerNorm(cross_attention_dim) + + def forward(self, id_embeds): + clip_extra_context_tokens = self.proj(id_embeds) + clip_extra_context_tokens = clip_extra_context_tokens.reshape(-1, self.num_tokens, self.cross_attention_dim) + clip_extra_context_tokens = self.norm(clip_extra_context_tokens) + return clip_extra_context_tokens + +class ProjModelFaceIdPlusImport(torch.nn.Module): + def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, clip_embeddings_dim=1280, num_tokens=4): + super().__init__() + + self.cross_attention_dim = cross_attention_dim + self.num_tokens = num_tokens + + self.proj = torch.nn.Sequential( + torch.nn.Linear(id_embeddings_dim, id_embeddings_dim*2), + torch.nn.GELU(), + torch.nn.Linear(id_embeddings_dim*2, cross_attention_dim*num_tokens), + ) + self.norm = torch.nn.LayerNorm(cross_attention_dim) + + self.perceiver_resampler = FacePerceiverResamplerImport( + dim=cross_attention_dim, + depth=4, + dim_head=64, + heads=cross_attention_dim // 64, + embedding_dim=clip_embeddings_dim, + output_dim=cross_attention_dim, + ff_mult=4, + ) + + def forward(self, id_embeds, clip_embeds, scale=1.0, shortcut=False): + x = self.proj(id_embeds) + x = x.reshape(-1, self.num_tokens, self.cross_attention_dim) + x = self.norm(x) + out = self.perceiver_resampler(x, clip_embeds) + if shortcut: + out = x + scale * out + return out + class ImageProjModelImport(nn.Module): def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4): super().__init__() @@ -62,73 +155,6 @@ class To_KVImport(nn.Module): self.to_kvs[key.replace(".weight", "").replace(".", "_")] = nn.Linear(value.shape[1], value.shape[0], bias=False) self.to_kvs[key.replace(".weight", "").replace(".", "_")].weight.data = value -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), - ) - - -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) - -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 - def set_model_patch_replace(model, patch_kwargs, key): to = model.model_options["transformer_options"] if "patches_replace" not in to: @@ -136,11 +162,86 @@ def set_model_patch_replace(model, patch_kwargs, key): if "attn2" not in to["patches_replace"]: to["patches_replace"]["attn2"] = {} if key not in to["patches_replace"]["attn2"]: - patch = CrossAttentionPatchImport(**patch_kwargs) - to["patches_replace"]["attn2"][key] = patch + to["patches_replace"]["attn2"][key] = CrossAttentionPatchImport(**patch_kwargs) else: to["patches_replace"]["attn2"][key].set_new_condition(**patch_kwargs) +def masked_tiling(image, short_side_tiles, weight=0.6, blur=0): + _, orig_height, orig_width, _ = image.shape + tile_size = 224 + + if orig_width < orig_height: + num_tiles_x = short_side_tiles + new_width = tile_size * num_tiles_x + new_height = orig_height * new_width // orig_width + num_tiles_y = new_height // tile_size + else: + num_tiles_y = short_side_tiles + new_height = tile_size * num_tiles_y + new_width = orig_width * new_height // orig_height + num_tiles_x = new_width // tile_size + + start_x = start_y = 0 + if (new_height % tile_size) >= tile_size//4: + num_tiles_y += 1 + else: + start_y = (new_height - tile_size*num_tiles_y) // 2 + + if (new_width % tile_size) >= tile_size//4: + num_tiles_x += 1 + else: + start_x = (new_width - tile_size*num_tiles_x) // 2 + + weight = 1.0 if num_tiles_x == 1 or num_tiles_y == 1 else weight + + ref_image = F.interpolate(image.permute([0,3,1,2]), size=(new_height, new_width), mode="bicubic").permute([0,2,3,1]) + + tiles = [] + attn_mask = [] + + for i in range(num_tiles_y): + for j in range(num_tiles_x): + start_height = i * tile_size + start_y + end_height = start_height + tile_size + start_y + start_width = j * tile_size + start_x + end_width = start_width + tile_size + start_x + + if end_height > new_height: + start_height = new_height - tile_size + end_height = new_height + if end_width > new_width: + start_width = new_width - tile_size + end_width = new_width + + tile = ref_image[:1, start_height:end_height, start_width:end_width, :] + tiles.append(tile.squeeze(0)) + + # create mask + mask = torch.zeros([new_height, new_width], dtype=image.dtype, device=image.device) + mask[start_height:end_height, start_width:end_width] = weight + + attn_mask.append(mask) + + image = torch.stack(tiles, dim=0) + attn_mask = torch.stack(attn_mask, dim=0) + + # If we have a lot of tiles we add bigger tiles with a higher weight to give clipvision a better idea of the overall image + if num_tiles_x != 1 and num_tiles_y != 1: + comp_tiles, comp_attn_mask = masked_tiling(ref_image, 1, 1.0, blur) + + if image.shape[1:3] != comp_tiles.shape[1:3]: + comp_tiles = F.interpolate(comp_tiles.permute([0,3,1,2]), size=(image.shape[1], image.shape[2]), mode="bicubic").permute([0,2,3,1]) + image = torch.cat([comp_tiles, image], dim=0) + + if attn_mask.shape[1:3] != comp_attn_mask.shape[1:3]: + comp_attn_mask = F.interpolate(comp_attn_mask.unsqueeze(1), size=(attn_mask.shape[1], attn_mask.shape[2]), mode="bicubic").squeeze(1) + attn_mask = torch.cat([comp_attn_mask, attn_mask], dim=0) + + if blur > 0: + attn_mask = TT.GaussianBlur(int(6*blur+1), blur)(attn_mask.unsqueeze(1)).permute([0,2,3,1]).squeeze(-1) + + return image, attn_mask + def image_add_noise(image, noise): image = image.permute([0,3,1,2]) torch.manual_seed(0) # use a fixed random for reproducible results @@ -159,19 +260,25 @@ def image_add_noise(image, noise): def zeroed_hidden_states(clip_vision, batch_size): image = torch.zeros([batch_size, 224, 224, 3]) comfy.model_management.load_model_gpu(clip_vision.patcher) - pixel_values = clip_preprocess(image.to(clip_vision.load_device)) - - if clip_vision.dtype != torch.float32: - precision_scope = torch.autocast - else: - precision_scope = lambda a, b: contextlib.nullcontext(a) - - with precision_scope(comfy.model_management.get_autocast_device(clip_vision.load_device), torch.float32): - outputs = clip_vision.model(pixel_values, intermediate_output=-2) + pixel_values = clip_preprocess(image.to(clip_vision.load_device)).float() + outputs = clip_vision.model(pixel_values=pixel_values, intermediate_output=-2) # we only need the penultimate hidden states - outputs = outputs[1].to(comfy.model_management.intermediate_device()) + return outputs[1].to(comfy.model_management.intermediate_device()) +def encode_image_masked(clip_vision, image, mask=None): + comfy.model_management.load_model_gpu(clip_vision.patcher) + pixel_values = clip_preprocess(image.to(clip_vision.load_device)).float() + + if mask is not None: + pixel_values = pixel_values * mask.to(clip_vision.load_device) + + out = clip_vision.model(pixel_values=pixel_values, intermediate_output=-2) + + outputs = Output() + outputs["last_hidden_state"] = out[0].to(comfy.model_management.intermediate_device()) + outputs["image_embeds"] = out[2].to(comfy.model_management.intermediate_device()) + outputs["penultimate_hidden_states"] = out[1].to(comfy.model_management.intermediate_device()) return outputs def min_(tensor_list): @@ -221,13 +328,27 @@ def contrast_adaptive_sharpening(image, amount): div = torch.reciprocal(1 + 4*w) output = ((b + d + f + h)*w + e) * div - output = output.clamp(0, 1) output = torch.nan_to_num(output) + output = output.clamp(0, 1) return (output) +def tensorToNP(image): + out = torch.clamp(255. * image.detach().cpu(), 0, 255).to(torch.uint8) + out = out[..., [2, 1, 0]] + out = out.numpy() + + return out + +def NPToTensor(image): + out = torch.from_numpy(image) + out = torch.clamp(out.to(torch.float)/255., 0.0, 1.0) + out = out[..., [2, 1, 0]] + + return out + class IPAdapterImport(nn.Module): - def __init__(self, ipadapter_model, cross_attention_dim=1024, output_cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4, is_sdxl=False, is_plus=False, is_full=False): + def __init__(self, ipadapter_model, cross_attention_dim=1024, output_cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4, is_sdxl=False, is_plus=False, is_full=False, is_faceid=False): super().__init__() self.clip_embeddings_dim = clip_embeddings_dim @@ -236,8 +357,15 @@ class IPAdapterImport(nn.Module): self.clip_extra_context_tokens = clip_extra_context_tokens self.is_sdxl = is_sdxl self.is_full = is_full + self.is_plus = is_plus + + if is_faceid: + self.image_proj_model = self.init_proj_faceid() + elif is_plus: + self.image_proj_model = self.init_proj_plus() + else: + self.image_proj_model = self.init_proj() - self.image_proj_model = self.init_proj() if not is_plus else self.init_proj_plus() self.image_proj_model.load_state_dict(ipadapter_model["image_proj"]) self.ip_layers = To_KVImport(ipadapter_model["ip_adapter"]) @@ -268,21 +396,40 @@ class IPAdapterImport(nn.Module): ) return image_proj_model + def init_proj_faceid(self): + if self.is_plus: + image_proj_model = ProjModelFaceIdPlusImport( + cross_attention_dim=self.cross_attention_dim, + id_embeddings_dim=512, + clip_embeddings_dim=1280, + num_tokens=4, + ) + else: + image_proj_model = MLPProjModelFaceIdImport( + cross_attention_dim=self.cross_attention_dim, + id_embeddings_dim=512, + num_tokens=self.clip_extra_context_tokens, + ) + return image_proj_model + @torch.inference_mode() def get_image_embeds(self, clip_embed, clip_embed_zeroed): image_prompt_embeds = self.image_proj_model(clip_embed) uncond_image_prompt_embeds = self.image_proj_model(clip_embed_zeroed) return image_prompt_embeds, uncond_image_prompt_embeds + @torch.inference_mode() + def get_image_embeds_faceid_plus(self, face_embed, clip_embed, s_scale, shortcut): + embeds = self.image_proj_model(face_embed, clip_embed, scale=s_scale, shortcut=shortcut) + return embeds + class CrossAttentionPatchImport: # forward for patching - def __init__(self, weight, ipadapter, device, dtype, number, cond, uncond, weight_type, mask=None, sigma_start=0.0, sigma_end=1.0, unfold_batch=False): + def __init__(self, weight, ipadapter, number, cond, uncond, weight_type="original", mask=None, sigma_start=0.0, sigma_end=1.0, unfold_batch=False): self.weights = [weight] self.ipadapters = [ipadapter] self.conds = [cond] self.unconds = [uncond] - self.device = 'cuda' if 'cuda' in device.type else 'cpu' - self.dtype = dtype if 'cuda' in self.device else torch.bfloat16 self.number = number self.weight_type = [weight_type] self.masks = [mask] @@ -293,14 +440,12 @@ class CrossAttentionPatchImport: self.k_key = str(self.number*2+1) + "_to_k_ip" self.v_key = str(self.number*2+1) + "_to_v_ip" - def set_new_condition(self, weight, ipadapter, device, dtype, number, cond, uncond, weight_type, mask=None, sigma_start=0.0, sigma_end=1.0, unfold_batch=False): + def set_new_condition(self, weight, ipadapter, number, cond, uncond, weight_type="original", mask=None, sigma_start=0.0, sigma_end=1.0, unfold_batch=False): self.weights.append(weight) self.ipadapters.append(ipadapter) self.conds.append(cond) self.unconds.append(uncond) self.masks.append(mask) - self.device = 'cuda' if 'cuda' in device.type else 'cpu' - self.dtype = dtype if 'cuda' in self.device else torch.bfloat16 self.weight_type.append(weight_type) self.sigma_start.append(sigma_start) self.sigma_end.append(sigma_end) @@ -309,131 +454,186 @@ class CrossAttentionPatchImport: def __call__(self, n, context_attn2, value_attn2, extra_options): org_dtype = n.dtype cond_or_uncond = extra_options["cond_or_uncond"] - sigma = extra_options["sigmas"][0].item() if 'sigmas' in extra_options else 999999999.9 + sigma = extra_options["sigmas"][0] if 'sigmas' in extra_options else None + sigma = sigma.item() if sigma is not None else 999999999.9 # extra options for AnimateDiff ad_params = extra_options['ad_params'] if "ad_params" in extra_options else None - with torch.autocast(device_type=self.device, dtype=self.dtype): - q = n - k = context_attn2 - v = value_attn2 - b = q.shape[0] - qs = q.shape[1] - batch_prompt = b // len(cond_or_uncond) - out = optimized_attention(q, k, v, extra_options["n_heads"]) - _, _, lh, lw = extra_options["original_shape"] - - for weight, cond, uncond, ipadapter, mask, weight_type, sigma_start, sigma_end, unfold_batch in zip(self.weights, self.conds, self.unconds, self.ipadapters, self.masks, self.weight_type, self.sigma_start, self.sigma_end, self.unfold_batch): - if sigma > sigma_start or sigma < sigma_end: - continue + q = n + k = context_attn2 + v = value_attn2 + b = q.shape[0] + qs = q.shape[1] + batch_prompt = b // len(cond_or_uncond) + out = optimized_attention(q, k, v, extra_options["n_heads"]) + _, _, lh, lw = extra_options["original_shape"] + + for weight, cond, uncond, ipadapter, mask, weight_type, sigma_start, sigma_end, unfold_batch in zip(self.weights, self.conds, self.unconds, self.ipadapters, self.masks, self.weight_type, self.sigma_start, self.sigma_end, self.unfold_batch): + if sigma > sigma_start or sigma < sigma_end: + continue - if unfold_batch and cond.shape[0] > 1: - # Check AnimateDiff context window - if ad_params is not None and ad_params["sub_idxs"] is not None: - # if images 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, need to do more to get proper sub_idxs masks - else: - # check if images length matches full_length - if not, make it match - if cond.shape[0] < ad_params["full_length"]: - cond = torch.cat((cond, cond[-1:].repeat((ad_params["full_length"]-cond.shape[0], 1, 1))), dim=0) - uncond = torch.cat((uncond, uncond[-1:].repeat((ad_params["full_length"]-uncond.shape[0], 1, 1))), dim=0) - # if we have too many remove the excess (should not happen, but just in case) - if cond.shape[0] > ad_params["full_length"]: - cond = cond[:ad_params["full_length"]] - uncond = uncond[:ad_params["full_length"]] - cond = cond[ad_params["sub_idxs"]] - uncond = uncond[ad_params["sub_idxs"]] - - # if we don't have enough reference images repeat the last one until we reach the right size - if cond.shape[0] < batch_prompt: - cond = torch.cat((cond, cond[-1:].repeat((batch_prompt-cond.shape[0], 1, 1))), dim=0) - uncond = torch.cat((uncond, uncond[-1:].repeat((batch_prompt-uncond.shape[0], 1, 1))), dim=0) - # if we have too many remove the exceeding - elif cond.shape[0] > batch_prompt: - cond = cond[:batch_prompt] - uncond = uncond[:batch_prompt] - - k_cond = ipadapter.ip_layers.to_kvs[self.k_key](cond) - k_uncond = ipadapter.ip_layers.to_kvs[self.k_key](uncond) - v_cond = ipadapter.ip_layers.to_kvs[self.v_key](cond) - v_uncond = ipadapter.ip_layers.to_kvs[self.v_key](uncond) - else: - k_cond = ipadapter.ip_layers.to_kvs[self.k_key](cond).repeat(batch_prompt, 1, 1) - k_uncond = ipadapter.ip_layers.to_kvs[self.k_key](uncond).repeat(batch_prompt, 1, 1) - v_cond = ipadapter.ip_layers.to_kvs[self.v_key](cond).repeat(batch_prompt, 1, 1) - v_uncond = ipadapter.ip_layers.to_kvs[self.v_key](uncond).repeat(batch_prompt, 1, 1) - - if weight_type.startswith("linear"): - ip_k = torch.cat([(k_cond, k_uncond)[i] for i in cond_or_uncond], dim=0) * weight - ip_v = torch.cat([(v_cond, v_uncond)[i] for i in cond_or_uncond], dim=0) * weight - else: - 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 weight_type.startswith("channel"): - # code by Lvmin Zhang at Stanford University as also seen on Fooocus IPAdapter implementation - # please read licensing notes https://github.com/lllyasviel/Fooocus/blob/main/fooocus_extras/ip_adapter.py#L225 - ip_v_mean = torch.mean(ip_v, dim=1, keepdim=True) - ip_v_offset = ip_v - ip_v_mean - _, _, C = ip_k.shape - channel_penalty = float(C) / 1280.0 - W = weight * channel_penalty - ip_k = ip_k * W - ip_v = ip_v_offset + ip_v_mean * W - - out_ip = optimized_attention(q, ip_k, ip_v, extra_options["n_heads"]) - if weight_type.startswith("original"): - out_ip = out_ip * weight - - if mask is not None: - # TODO: needs checking - mask_h = max(1, round(lh / math.sqrt(lh * lw / qs))) - mask_w = qs // 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, just get sub_idx masks, resize, and continue - if mask.shape[0] >= ad_params["full_length"]: - mask_downsample = torch.Tensor(mask[ad_params["sub_idxs"]]) - mask_downsample = F.interpolate(mask_downsample.unsqueeze(1), size=(mask_h, mask_w), mode="bicubic").squeeze(1) - # otherwise, need to do more to get proper sub_idxs masks - else: - # resize to needed attention size (to save on memory) - mask_downsample = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="bicubic").squeeze(1) - # check if mask length matches full_length - if not, make it match - if mask_downsample.shape[0] < ad_params["full_length"]: - mask_downsample = torch.cat((mask_downsample, mask_downsample[-1:].repeat((ad_params["full_length"]-mask_downsample.shape[0], 1, 1))), dim=0) - # if we have too many remove the excess (should not happen, but just in case) - if mask_downsample.shape[0] > ad_params["full_length"]: - mask_downsample = mask_downsample[:ad_params["full_length"]] - # now, select sub_idxs masks - mask_downsample = mask_downsample[ad_params["sub_idxs"]] - # otherwise, perform usual mask interpolation + if unfold_batch and cond.shape[0] > 1: + # Check AnimateDiff context window + if ad_params is not None and ad_params["sub_idxs"] is not None: + # if images 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, need to do more to get proper sub_idxs masks else: + # check if images length matches full_length - if not, make it match + if cond.shape[0] < ad_params["full_length"]: + cond = torch.cat((cond, cond[-1:].repeat((ad_params["full_length"]-cond.shape[0], 1, 1))), dim=0) + uncond = torch.cat((uncond, uncond[-1:].repeat((ad_params["full_length"]-uncond.shape[0], 1, 1))), dim=0) + # if we have too many remove the excess (should not happen, but just in case) + if cond.shape[0] > ad_params["full_length"]: + cond = cond[:ad_params["full_length"]] + uncond = uncond[:ad_params["full_length"]] + cond = cond[ad_params["sub_idxs"]] + uncond = uncond[ad_params["sub_idxs"]] + + # if we don't have enough reference images repeat the last one until we reach the right size + if cond.shape[0] < batch_prompt: + cond = torch.cat((cond, cond[-1:].repeat((batch_prompt-cond.shape[0], 1, 1))), dim=0) + uncond = torch.cat((uncond, uncond[-1:].repeat((batch_prompt-uncond.shape[0], 1, 1))), dim=0) + # if we have too many remove the exceeding + elif cond.shape[0] > batch_prompt: + cond = cond[:batch_prompt] + uncond = uncond[:batch_prompt] + + k_cond = ipadapter.ip_layers.to_kvs[self.k_key](cond) + k_uncond = ipadapter.ip_layers.to_kvs[self.k_key](uncond) + v_cond = ipadapter.ip_layers.to_kvs[self.v_key](cond) + v_uncond = ipadapter.ip_layers.to_kvs[self.v_key](uncond) + else: + k_cond = ipadapter.ip_layers.to_kvs[self.k_key](cond).repeat(batch_prompt, 1, 1) + k_uncond = ipadapter.ip_layers.to_kvs[self.k_key](uncond).repeat(batch_prompt, 1, 1) + v_cond = ipadapter.ip_layers.to_kvs[self.v_key](cond).repeat(batch_prompt, 1, 1) + v_uncond = ipadapter.ip_layers.to_kvs[self.v_key](uncond).repeat(batch_prompt, 1, 1) + + if weight_type.startswith("linear"): + ip_k = torch.cat([(k_cond, k_uncond)[i] for i in cond_or_uncond], dim=0) * weight + ip_v = torch.cat([(v_cond, v_uncond)[i] for i in cond_or_uncond], dim=0) * weight + else: + 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 weight_type.startswith("channel"): + # code by Lvmin Zhang at Stanford University as also seen on Fooocus IPAdapter implementation + ip_v_mean = torch.mean(ip_v, dim=1, keepdim=True) + ip_v_offset = ip_v - ip_v_mean + _, _, C = ip_k.shape + channel_penalty = float(C) / 1280.0 + W = weight * channel_penalty + ip_k = ip_k * W + ip_v = ip_v_offset + ip_v_mean * W + + out_ip = optimized_attention(q, ip_k, ip_v, extra_options["n_heads"]) + if weight_type.startswith("original"): + out_ip = out_ip * weight + + if mask is not None: + # TODO: needs checking + mask_h = lh / math.sqrt(lh * lw / qs) + mask_h = int(mask_h) + int((qs % int(mask_h)) != 0) + mask_w = qs // 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, just get sub_idx masks, resize, and continue + if mask.shape[0] >= ad_params["full_length"]: + mask_downsample = torch.Tensor(mask[ad_params["sub_idxs"]]) + mask_downsample = F.interpolate(mask_downsample.unsqueeze(1), size=(mask_h, mask_w), mode="bicubic").squeeze(1) + # otherwise, need to do more to get proper sub_idxs masks + else: + # resize to needed attention size (to save on memory) mask_downsample = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="bicubic").squeeze(1) + # check if mask length matches full_length - if not, make it match + if mask_downsample.shape[0] < ad_params["full_length"]: + mask_downsample = torch.cat((mask_downsample, mask_downsample[-1:].repeat((ad_params["full_length"]-mask_downsample.shape[0], 1, 1))), dim=0) + # if we have too many remove the excess (should not happen, but just in case) + if mask_downsample.shape[0] > ad_params["full_length"]: + mask_downsample = mask_downsample[:ad_params["full_length"]] + # now, select sub_idxs masks + mask_downsample = mask_downsample[ad_params["sub_idxs"]] + # otherwise, perform usual mask interpolation + else: + mask_downsample = F.interpolate(mask.unsqueeze(1), size=(mask_h, mask_w), mode="bicubic").squeeze(1) - # if we don't have enough masks repeat the last one until we reach the right size - if mask_downsample.shape[0] < batch_prompt: - mask_downsample = torch.cat((mask_downsample, mask_downsample[-1:, :, :].repeat((batch_prompt-mask_downsample.shape[0], 1, 1))), dim=0) - # if we have too many remove the exceeding - elif mask_downsample.shape[0] > batch_prompt: - mask_downsample = mask_downsample[:batch_prompt, :, :] - - # repeat the masks - mask_downsample = mask_downsample.repeat(len(cond_or_uncond), 1, 1) - mask_downsample = mask_downsample.view(mask_downsample.shape[0], -1, 1).repeat(1, 1, out.shape[2]) + # if we don't have enough masks repeat the last one until we reach the right size + if mask_downsample.shape[0] < batch_prompt: + mask_downsample = torch.cat((mask_downsample, mask_downsample[-1:, :, :].repeat((batch_prompt-mask_downsample.shape[0], 1, 1))), dim=0) + # if we have too many remove the exceeding + elif mask_downsample.shape[0] > batch_prompt: + mask_downsample = mask_downsample[:batch_prompt, :, :] + + # repeat the masks + mask_downsample = mask_downsample.repeat(len(cond_or_uncond), 1, 1) + mask_downsample = mask_downsample.view(mask_downsample.shape[0], -1, 1).repeat(1, 1, out.shape[2]) - out_ip = out_ip * mask_downsample + out_ip = out_ip * mask_downsample - out = out + out_ip + out = out + out_ip return out.to(dtype=org_dtype) +class IPAdapterModelLoaderImport: + @classmethod + def INPUT_TYPES(s): + return {"required": { "ipadapter_file": (folder_paths.get_filename_list("ipadapter"), )}} + RETURN_TYPES = ("IPADAPTER",) + FUNCTION = "load_ipadapter_model" + CATEGORY = "ipadapter" + + def load_ipadapter_model(self, ipadapter_file): + ckpt_path = folder_paths.get_full_path("ipadapter", ipadapter_file) + + model = comfy.utils.load_torch_file(ckpt_path, safe_load=True) + + if ckpt_path.lower().endswith(".safetensors"): + st_model = {"image_proj": {}, "ip_adapter": {}} + for key in model.keys(): + if key.startswith("image_proj."): + st_model["image_proj"][key.replace("image_proj.", "")] = model[key] + elif key.startswith("ip_adapter."): + st_model["ip_adapter"][key.replace("ip_adapter.", "")] = model[key] + model = st_model + + if not "ip_adapter" in model.keys() or not model["ip_adapter"]: + raise Exception("invalid IPAdapter model {}".format(ckpt_path)) + + return (model,) + +insightface_face_align = None +class InsightFaceLoaderImport: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "provider": (["CPU", "CUDA", "ROCM"], ), + }, + } + + RETURN_TYPES = ("INSIGHTFACE",) + FUNCTION = "load_insight_face" + CATEGORY = "ipadapter" + + def load_insight_face(self, provider): + try: + from insightface.app import FaceAnalysis + except ImportError as e: + raise Exception(e) + + from insightface.utils import face_align + global insightface_face_align + insightface_face_align = face_align + + model = FaceAnalysis(name="buffalo_l", root=INSIGHTFACE_DIR, providers=[provider + 'ExecutionProvider',]) + model.prepare(ctx_id=0, det_size=(640, 640)) + + return (model,) class IPAdapterApplyImport: @classmethod @@ -460,41 +660,109 @@ class IPAdapterApplyImport: FUNCTION = "apply_ipadapter" CATEGORY = "ipadapter" - def apply_ipadapter(self, ipadapter, model, weight, clip_vision=None, image=None, weight_type="original", noise=None, embeds=None, attn_mask=None, start_at=0.0, end_at=1.0, unfold_batch=False): - self.dtype = model.model.diffusion_model.dtype + def apply_ipadapter(self, + ipadapter, + model, + weight, + clip_vision=None, + image=None, + weight_type="original", + noise=None, + embeds=None, + attn_mask=None, + start_at=0.0, + end_at=1.0, + unfold_batch=False, + insightface=None, + faceid_v2=False, + weight_v2=False, + clip_vision_mask=None, + short_side_tiles=0, + tile_weight=0.0, + tile_blur=0, + ): + + self.dtype = torch.float16 if comfy.model_management.should_use_fp16() else torch.float32 self.device = comfy.model_management.get_torch_device() self.weight = weight - self.is_full = "proj.0.weight" in ipadapter["image_proj"] - self.is_plus = self.is_full or "latents" in ipadapter["image_proj"] + self.is_full = "proj.3.weight" in ipadapter["image_proj"] + self.is_portrait = "proj.2.weight" in ipadapter["image_proj"] and not "proj.3.weight" in ipadapter["image_proj"] and not "0.to_q_lora.down.weight" in ipadapter["ip_adapter"] + self.is_faceid = self.is_portrait or "0.to_q_lora.down.weight" in ipadapter["ip_adapter"] + self.is_plus = (self.is_full or "latents" in ipadapter["image_proj"] or "perceiver_resampler.proj_in.weight" in ipadapter["image_proj"]) + is_tiled = True if short_side_tiles > 0 else False + + if is_tiled: + image, attn_mask = masked_tiling(image, short_side_tiles, tile_weight, tile_blur) + + if self.is_faceid and not insightface: + raise Exception('InsightFace must be provided for FaceID models.') output_cross_attention_dim = ipadapter["ip_adapter"]["1.to_k_ip.weight"].shape[1] self.is_sdxl = output_cross_attention_dim == 2048 - cross_attention_dim = 1280 if self.is_plus and self.is_sdxl else output_cross_attention_dim - clip_extra_context_tokens = 16 if self.is_plus else 4 + cross_attention_dim = 1280 if self.is_plus and self.is_sdxl and not self.is_faceid else output_cross_attention_dim + clip_extra_context_tokens = 16 if self.is_plus or self.is_portrait else 4 if embeds is not None: embeds = torch.unbind(embeds) clip_embed = embeds[0].cpu() clip_embed_zeroed = embeds[1].cpu() else: - if image.shape[1] != image.shape[2]: - print("\033[33mINFO: the IPAdapter reference image is not a square, CLIPImageProcessor will resize and crop it at the center. If the main focus of the picture is not in the middle the result might not be what you are expecting.\033[0m") + if self.is_faceid: + insightface.det_model.input_size = (640,640) # reset the detection size + face_img = tensorToNP(image) + face_embed = [] + face_clipvision = [] - clip_embed = clip_vision.encode_image(image) - neg_image = image_add_noise(image, noise) if noise > 0 else None - - if self.is_plus: - clip_embed = clip_embed.penultimate_hidden_states - if noise > 0: - clip_embed_zeroed = clip_vision.encode_image(neg_image).penultimate_hidden_states - else: - clip_embed_zeroed = zeroed_hidden_states(clip_vision, image.shape[0]) - else: - clip_embed = clip_embed.image_embeds - if noise > 0: - clip_embed_zeroed = clip_vision.encode_image(neg_image).image_embeds + for i in range(face_img.shape[0]): + for size in [(size, size) for size in range(640, 128, -64)]: + insightface.det_model.input_size = size # TODO: hacky but seems to be working + face = insightface.get(face_img[i]) + if face: + face_embed.append(torch.from_numpy(face[0].normed_embedding).unsqueeze(0)) + face_clipvision.append(NPToTensor(insightface_face_align.norm_crop(face_img[i], landmark=face[0].kps, image_size=256))) + + if 640 not in size: + print(f"\033[33mINFO: InsightFace detection resolution lowered to {size}.\033[0m") + break + else: + raise Exception('InsightFace: No face detected.') + + face_embed = torch.stack(face_embed, dim=0) + image = torch.stack(face_clipvision, dim=0) + + neg_image = image_add_noise(image, noise) if noise > 0 else None + + if self.is_plus: + clip_embed = clip_vision.encode_image(image).penultimate_hidden_states + if noise > 0: + clip_embed_zeroed = clip_vision.encode_image(neg_image).penultimate_hidden_states + else: + clip_embed_zeroed = zeroed_hidden_states(clip_vision, image.shape[0]) + + # TODO: check noise to the uncods too + face_embed_zeroed = torch.zeros_like(face_embed) else: + clip_embed = face_embed clip_embed_zeroed = torch.zeros_like(clip_embed) + else: + if image.shape[1] != image.shape[2]: + print("\033[33mINFO: the IPAdapter reference image is not a square, CLIPImageProcessor will resize and crop it at the center. If the main focus of the picture is not in the middle the result might not be what you are expecting.\033[0m") + + clip_embed = encode_image_masked(clip_vision, image, clip_vision_mask) + neg_image = image_add_noise(image, noise) if noise > 0 else None + + if self.is_plus: + clip_embed = clip_embed.penultimate_hidden_states + if noise > 0: + clip_embed_zeroed = clip_vision.encode_image(neg_image).penultimate_hidden_states + else: + clip_embed_zeroed = zeroed_hidden_states(clip_vision, image.shape[0]) + else: + clip_embed = clip_embed.image_embeds + if noise > 0: + clip_embed_zeroed = clip_vision.encode_image(neg_image).image_embeds + else: + clip_embed_zeroed = torch.zeros_like(clip_embed) clip_embeddings_dim = clip_embed.shape[-1] @@ -507,63 +775,131 @@ class IPAdapterApplyImport: is_sdxl=self.is_sdxl, is_plus=self.is_plus, is_full=self.is_full, + is_faceid=self.is_faceid, ) self.ipadapter.to(self.device, dtype=self.dtype) - image_prompt_embeds, uncond_image_prompt_embeds = self.ipadapter.get_image_embeds(clip_embed.to(self.device, self.dtype), clip_embed_zeroed.to(self.device, self.dtype)) + if self.is_faceid and self.is_plus: + image_prompt_embeds = self.ipadapter.get_image_embeds_faceid_plus(face_embed.to(self.device, dtype=self.dtype), clip_embed.to(self.device, dtype=self.dtype), weight_v2, faceid_v2) + uncond_image_prompt_embeds = self.ipadapter.get_image_embeds_faceid_plus(face_embed_zeroed.to(self.device, dtype=self.dtype), clip_embed_zeroed.to(self.device, dtype=self.dtype), weight_v2, faceid_v2) + else: + image_prompt_embeds, uncond_image_prompt_embeds = self.ipadapter.get_image_embeds(clip_embed.to(self.device, dtype=self.dtype), clip_embed_zeroed.to(self.device, dtype=self.dtype)) + image_prompt_embeds = image_prompt_embeds.to(self.device, dtype=self.dtype) uncond_image_prompt_embeds = uncond_image_prompt_embeds.to(self.device, dtype=self.dtype) - work_model = model.clone() + self.work_model = model.clone() if attn_mask is not None: attn_mask = attn_mask.to(self.device) - sigma_start = model.model.model_sampling.percent_to_sigma(start_at) - sigma_end = model.model.model_sampling.percent_to_sigma(end_at) + sigma_start = self.work_model.model.model_sampling.percent_to_sigma(start_at) + sigma_end = self.work_model.model.model_sampling.percent_to_sigma(end_at) + + if is_tiled: + for i in range(image_prompt_embeds.shape[0]): #TODO: check if we can apply one mask per image in the attention patch phase + patch_kwargs = { + "number": 0, + "weight": self.weight, + "ipadapter": self.ipadapter, + "cond": image_prompt_embeds[i].unsqueeze(0), + "uncond": uncond_image_prompt_embeds[i].unsqueeze(0), + "weight_type": weight_type, + "mask": attn_mask[i].unsqueeze(0), + "sigma_start": sigma_start, + "sigma_end": sigma_end, + "unfold_batch": unfold_batch, + } + self.apply_patch(patch_kwargs) + else: + patch_kwargs = { + "number": 0, + "weight": self.weight, + "ipadapter": self.ipadapter, + "cond": image_prompt_embeds, + "uncond": uncond_image_prompt_embeds, + "weight_type": weight_type, + "mask": attn_mask, + "sigma_start": sigma_start, + "sigma_end": sigma_end, + "unfold_batch": unfold_batch, + } + self.apply_patch(patch_kwargs) - patch_kwargs = { - "number": 0, - "weight": self.weight, - "ipadapter": self.ipadapter, - "device": self.device, - "dtype": self.dtype, - "cond": image_prompt_embeds, - "uncond": uncond_image_prompt_embeds, - "weight_type": weight_type, - "mask": attn_mask, - "sigma_start": sigma_start, - "sigma_end": sigma_end, - "unfold_batch": unfold_batch, - } + return (self.work_model, attn_mask, image,) + def apply_patch(self, patch_kwargs): if not self.is_sdxl: for id in [1,2,4,5,7,8]: # id of input_blocks that have cross attention - set_model_patch_replace(work_model, patch_kwargs, ("input", id)) + set_model_patch_replace(self.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 - set_model_patch_replace(work_model, patch_kwargs, ("output", id)) + set_model_patch_replace(self.work_model, patch_kwargs, ("output", id)) patch_kwargs["number"] += 1 - set_model_patch_replace(work_model, patch_kwargs, ("middle", 0)) + set_model_patch_replace(self.work_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: - set_model_patch_replace(work_model, patch_kwargs, ("input", id, index)) + set_model_patch_replace(self.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, patch_kwargs, ("output", id, index)) + set_model_patch_replace(self.work_model, patch_kwargs, ("output", id, index)) patch_kwargs["number"] += 1 for index in range(10): - set_model_patch_replace(work_model, patch_kwargs, ("middle", 0, index)) + set_model_patch_replace(self.work_model, patch_kwargs, ("middle", 0, index)) patch_kwargs["number"] += 1 - return (work_model, ) +class IPAdapterTilesMaskedImport(IPAdapterApplyImport): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ipadapter": ("IPADAPTER", ), + "clip_vision": ("CLIP_VISION",), + "image": ("IMAGE",), + "model": ("MODEL", ), + "weight": ("FLOAT", { "default": 0.7, "min": -1, "max": 3, "step": 0.05 }), + "noise": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01 }), + "weight_type": (["original", "linear", "channel penalty"], ), + "start_at": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001 }), + "end_at": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001 }), + "short_side_tiles": ("INT", { "default": 1, "min": 0, "max": 12, "step": 1 }), + "tile_weight": ("FLOAT", { "default": 0.6, "min": 0.0, "max": 1.0, "step": 0.05 }), + #"tile_blur": ("INT", { "default": 0, "min": 0, "max": 112, "step": 1 }), + }, + } -def prep_image(image, interpolation="LANCZOS", crop_position="center", sharpening=0.0): + RETURN_TYPES = ("MODEL",) + +class IPAdapterApplyFaceIDImport(IPAdapterApplyImport): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ipadapter": ("IPADAPTER", ), + "clip_vision": ("CLIP_VISION",), + "insightface": ("INSIGHTFACE",), + "image": ("IMAGE",), + "model": ("MODEL", ), + "weight": ("FLOAT", { "default": 1.0, "min": -1, "max": 3, "step": 0.05 }), + "noise": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01 }), + "weight_type": (["original", "linear", "channel penalty"], ), + "start_at": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001 }), + "end_at": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001 }), + "faceid_v2": ("BOOLEAN", { "default": False }), + "weight_v2": ("FLOAT", { "default": 1.0, "min": -1, "max": 3, "step": 0.05 }), + "unfold_batch": ("BOOLEAN", { "default": False }), + }, + "optional": { + "attn_mask": ("MASK",), + } + } + +def prepImage(image, interpolation="LANCZOS", crop_position="center", size=(224,224), sharpening=0.0, padding=0): _, oh, ow, _ = image.shape output = image.permute([0,3,1,2]) @@ -597,62 +933,68 @@ def prep_image(image, interpolation="LANCZOS", crop_position="center", sharpenin imgs = [] for i in range(output.shape[0]): img = TT.ToPILImage()(output[i]) - img = img.resize((224,224), resample=Image.Resampling[interpolation]) + img = img.resize(size, resample=Image.Resampling[interpolation]) imgs.append(TT.ToTensor()(img)) output = torch.stack(imgs, dim=0) + imgs = None # zelous GC if sharpening > 0: output = contrast_adaptive_sharpening(output, sharpening) + if padding > 0: + output = F.pad(output, (padding, padding, padding, padding), value=255, mode="constant") + output = output.permute([0,2,3,1]) - return (output,) + return output -class ResamplerImport(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) +class PrepImageForInsightFaceImport: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE",), + "crop_position": (["center", "top", "bottom", "left", "right"],), + "sharpening": ("FLOAT", {"default": 0.0, "min": 0, "max": 1, "step": 0.05}), + "pad_around": ("BOOLEAN", { "default": True }), + }, + } - 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), - ] - ) - ) + RETURN_TYPES = ("IMAGE",) + FUNCTION = "prep_image" - 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) + CATEGORY = "ipadapter" + def prep_image(self, image, crop_position, sharpening=0.0, pad_around=True): + if pad_around: + padding = 30 + size = (580, 580) + else: + padding = 0 + size = (640, 640) + output = prepImage(image, "LANCZOS", crop_position, size, sharpening, padding) + + return (output, ) + +class PrepImageForClipVisionImport: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "image": ("IMAGE",), + "interpolation": (["LANCZOS", "BICUBIC", "HAMMING", "BILINEAR", "BOX", "NEAREST"],), + "crop_position": (["top", "bottom", "left", "right", "center", "pad"],), + "sharpening": ("FLOAT", {"default": 0.0, "min": 0, "max": 1, "step": 0.05}), + }, + } + + RETURN_TYPES = ("IMAGE",) + FUNCTION = "prep_image" + + CATEGORY = "ipadapter" + + def prep_image(self, image, interpolation="LANCZOS", crop_position="center", sharpening=0.0): + size = (224, 224) + output = prepImage(image, interpolation, crop_position, size, sharpening, 0) + return (output, ) class IPAdapterEncoderImport: @classmethod @@ -680,13 +1022,9 @@ class IPAdapterEncoderImport: def preprocess(self, clip_vision, image_1, ipadapter_plus, noise, weight_1, image_2=None, image_3=None, image_4=None, weight_2=1.0, weight_3=1.0, weight_4=1.0): weight_1 *= (0.1 + (weight_1 - 0.1)) - weight_1 = 1.19e-05 if weight_1 <= 1.19e-05 else weight_1 weight_2 *= (0.1 + (weight_2 - 0.1)) - weight_2 = 1.19e-05 if weight_2 <= 1.19e-05 else weight_2 weight_3 *= (0.1 + (weight_3 - 0.1)) - weight_3 = 1.19e-05 if weight_3 <= 1.19e-05 else weight_3 weight_4 *= (0.1 + (weight_4 - 0.1)) - weight_5 = 1.19e-05 if weight_4 <= 1.19e-05 else weight_4 image = image_1 weight = [weight_1]*image_1.shape[0] @@ -731,7 +1069,67 @@ class IPAdapterEncoderImport: return( output, ) +class IPAdapterApplyEncodedImport(IPAdapterApplyImport): + @classmethod + def INPUT_TYPES(s): + return { + "required": { + "ipadapter": ("IPADAPTER", ), + "embeds": ("EMBEDS",), + "model": ("MODEL", ), + "weight": ("FLOAT", { "default": 1.0, "min": -1, "max": 3, "step": 0.05 }), + "weight_type": (["original", "linear", "channel penalty"], ), + "start_at": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.001 }), + "end_at": ("FLOAT", { "default": 1.0, "min": 0.0, "max": 1.0, "step": 0.001 }), + "unfold_batch": ("BOOLEAN", { "default": False }), + }, + "optional": { + "attn_mask": ("MASK",), + } + } +class IPAdapterSaveEmbedsImport: + def __init__(self): + self.output_dir = folder_paths.get_output_directory() + + @classmethod + def INPUT_TYPES(s): + return {"required": { + "embeds": ("EMBEDS",), + "filename_prefix": ("STRING", {"default": "embeds/IPAdapter"}) + }, + } + + RETURN_TYPES = () + FUNCTION = "save" + OUTPUT_NODE = True + CATEGORY = "ipadapter" + + def save(self, embeds, filename_prefix): + full_output_folder, filename, counter, subfolder, filename_prefix = folder_paths.get_save_image_path(filename_prefix, self.output_dir) + file = f"{filename}_{counter:05}_.ipadpt" + file = os.path.join(full_output_folder, file) + + torch.save(embeds, file) + return (None, ) + + +class IPAdapterLoadEmbedsImport: + @classmethod + def INPUT_TYPES(s): + input_dir = folder_paths.get_input_directory() + files = [os.path.relpath(os.path.join(root, file), input_dir) for root, dirs, files in os.walk(input_dir) for file in files if file.endswith('.ipadpt')] + return {"required": {"embeds": [sorted(files), ]}, } + + RETURN_TYPES = ("EMBEDS", ) + FUNCTION = "load" + CATEGORY = "ipadapter" + + def load(self, embeds): + path = folder_paths.get_annotated_filepath(embeds) + output = torch.load(path).cpu() + + return (output, ) class IPAdapterBatchEmbedsImport: @@ -747,5 +1145,208 @@ class IPAdapterBatchEmbedsImport: CATEGORY = "ipadapter" def batch(self, embed1, embed2): - output = torch.cat((embed1, embed2), dim=1) - return (output, ) \ No newline at end of file + return (torch.cat((embed1, embed2), dim=1), ) + + + +import math + +import torch +import torch.nn as nn +from einops import rearrange +from einops.layers.torch import Rearrange + + +# FFN +def FeedForwardImport(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 PerceiverAttentionImport(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 ResamplerImport(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, + max_seq_len: int = 257, # CLIP tokens + CLS token + apply_pos_emb: bool = False, + num_latents_mean_pooled: int = 0, # number of latents derived from mean pooled representation of the sequence + ): + super().__init__() + self.pos_emb = nn.Embedding(max_seq_len, embedding_dim) if apply_pos_emb else None + + 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.to_latents_from_mean_pooled_seq = ( + nn.Sequential( + nn.LayerNorm(dim), + nn.Linear(dim, dim * num_latents_mean_pooled), + Rearrange("b (n d) -> b n d", n=num_latents_mean_pooled), + ) + if num_latents_mean_pooled > 0 + else None + ) + + self.layers = nn.ModuleList([]) + for _ in range(depth): + self.layers.append( + nn.ModuleList( + [ + PerceiverAttentionImport(dim=dim, dim_head=dim_head, heads=heads), + FeedForwardImport(dim=dim, mult=ff_mult), + ] + ) + ) + + def forward(self, x): + if self.pos_emb is not None: + n, device = x.shape[1], x.device + pos_emb = self.pos_emb(torch.arange(n, device=device)) + x = x + pos_emb + + latents = self.latents.repeat(x.size(0), 1, 1) + + x = self.proj_in(x) + + if self.to_latents_from_mean_pooled_seq: + meanpooled_seq = masked_mean(x, dim=1, mask=torch.ones(x.shape[:2], device=x.device, dtype=torch.bool)) + meanpooled_latents = self.to_latents_from_mean_pooled_seq(meanpooled_seq) + latents = torch.cat((meanpooled_latents, latents), dim=-2) + + 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) + + +def masked_mean(t, *, dim, mask=None): + if mask is None: + return t.mean(dim=dim) + + denom = mask.sum(dim=dim, keepdim=True) + mask = rearrange(mask, "b n -> b n 1") + masked_t = t.masked_fill(~mask, 0.0) + + return masked_t.sum(dim=dim) / denom.clamp(min=1e-5) + + +def prep_image(image, interpolation="LANCZOS", crop_position="center", sharpening=0.0): + _, oh, ow, _ = image.shape + output = image.permute([0,3,1,2]) + + if "pad" in crop_position: + target_length = max(oh, ow) + pad_l = (target_length - ow) // 2 + pad_r = (target_length - ow) - pad_l + pad_t = (target_length - oh) // 2 + pad_b = (target_length - oh) - pad_t + output = F.pad(output, (pad_l, pad_r, pad_t, pad_b), value=0, mode="constant") + else: + crop_size = min(oh, ow) + x = (ow-crop_size) // 2 + y = (oh-crop_size) // 2 + if "top" in crop_position: + y = 0 + elif "bottom" in crop_position: + y = oh-crop_size + elif "left" in crop_position: + x = 0 + elif "right" in crop_position: + x = ow-crop_size + + x2 = x+crop_size + y2 = y+crop_size + + # crop + output = output[:, :, y:y2, x:x2] + + # resize (apparently PIL resize is better than tourchvision interpolate) + imgs = [] + for i in range(output.shape[0]): + img = TT.ToPILImage()(output[i]) + img = img.resize((224,224), resample=Image.Resampling[interpolation]) + imgs.append(TT.ToTensor()(img)) + output = torch.stack(imgs, dim=0) + + if sharpening > 0: + output = contrast_adaptive_sharpening(output, sharpening) + + output = output.permute([0,2,3,1]) + + return (output,) \ No newline at end of file