diff --git a/config.yaml b/config.yaml deleted file mode 100644 index b99d4a7..0000000 --- a/config.yaml +++ /dev/null @@ -1,3 +0,0 @@ -#Plz don't delete this file, just edit it when neccessary. -ckpts_path: "./ckpts" -ops_backend: "cupy" #Either "taichi" or "cupy" \ No newline at end of file diff --git a/control/IPAdapterPlus.py b/control/IPAdapterPlus.py new file mode 100644 index 0000000..a2300bc --- /dev/null +++ b/control/IPAdapterPlus.py @@ -0,0 +1,751 @@ +import torch +import contextlib +import os +import math + +import comfy.utils +import comfy.model_management +from comfy.clip_vision import clip_preprocess +from comfy.ldm.modules.attention import optimized_attention +import folder_paths + +from torch import nn +from PIL import Image +import torch.nn.functional as F +import torchvision.transforms as TT + +from .resampler import Resampler + +# 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) +else: + folder_paths.folder_names_and_paths["ipadapter"][1].update(folder_paths.supported_pt_extensions) + +class MLPProjModel(torch.nn.Module): + """SD model with image prompt""" + def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024): + super().__init__() + + self.proj = torch.nn.Sequential( + torch.nn.Linear(clip_embeddings_dim, clip_embeddings_dim), + torch.nn.GELU(), + torch.nn.Linear(clip_embeddings_dim, cross_attention_dim), + torch.nn.LayerNorm(cross_attention_dim) + ) + + def forward(self, image_embeds): + clip_extra_context_tokens = self.proj(image_embeds) + return clip_extra_context_tokens + +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, state_dict): + super().__init__() + + self.to_kvs = nn.ModuleDict() + for key, value in state_dict.items(): + 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 set_model_patch_replace(model, patch_kwargs, key): + to = model.model_options["transformer_options"] + if "patches_replace" not in to: + to["patches_replace"] = {} + 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 image_add_noise(image, noise): + image = image.permute([0,3,1,2]) + torch.manual_seed(0) # use a fixed random for reproducible results + transforms = TT.Compose([ + TT.CenterCrop(min(image.shape[2], image.shape[3])), + TT.Resize((224, 224), interpolation=TT.InterpolationMode.BICUBIC, antialias=True), + TT.ElasticTransform(alpha=75.0, sigma=noise*3.5), # shuffle the image + TT.RandomVerticalFlip(p=1.0), # flip the image to change the geometry even more + TT.RandomHorizontalFlip(p=1.0), + ]) + image = transforms(image.cpu()) + image = image.permute([0,2,3,1]) + image = image + ((0.25*(1-noise)+0.05) * torch.randn_like(image) ) # add further random noise + return image + +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, output_hidden_states=True) + + # we only need the penultimate hidden states + outputs = outputs['hidden_states'][-2].cpu() if 'hidden_states' in outputs else None + + return outputs + +def min_(tensor_list): + # return the element-wise min of the tensor list. + x = torch.stack(tensor_list) + mn = x.min(axis=0)[0] + return torch.clamp(mn, min=0) + +def max_(tensor_list): + # return the element-wise max of the tensor list. + x = torch.stack(tensor_list) + mx = x.max(axis=0)[0] + return torch.clamp(mx, max=1) + +# From https://github.com/Jamy-L/Pytorch-Contrast-Adaptive-Sharpening/ +def contrast_adaptive_sharpening(image, amount): + img = F.pad(image, pad=(1, 1, 1, 1)).cpu() + + a = img[..., :-2, :-2] + b = img[..., :-2, 1:-1] + c = img[..., :-2, 2:] + d = img[..., 1:-1, :-2] + e = img[..., 1:-1, 1:-1] + f = img[..., 1:-1, 2:] + g = img[..., 2:, :-2] + h = img[..., 2:, 1:-1] + i = img[..., 2:, 2:] + + # Computing contrast + cross = (b, d, e, f, h) + mn = min_(cross) + mx = max_(cross) + + diag = (a, c, g, i) + mn2 = min_(diag) + mx2 = max_(diag) + mx = mx + mx2 + mn = mn + mn2 + + # Computing local weight + inv_mx = torch.reciprocal(mx) + amp = inv_mx * torch.minimum(mn, (2 - mx)) + + # scaling + amp = torch.sqrt(amp) + w = - amp * (amount * (1/5 - 1/8) + 1/8) + 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) + + return (output) + +class IPAdapter(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): + super().__init__() + + self.clip_embeddings_dim = clip_embeddings_dim + self.cross_attention_dim = cross_attention_dim + self.output_cross_attention_dim = output_cross_attention_dim + self.clip_extra_context_tokens = clip_extra_context_tokens + self.is_sdxl = is_sdxl + self.is_full = is_full + + 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_KV(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 + + def init_proj_plus(self): + if self.is_full: + image_proj_model = MLPProjModel( + cross_attention_dim=self.cross_attention_dim, + clip_embeddings_dim=self.clip_embeddings_dim + ) + else: + image_proj_model = Resampler( + dim=self.cross_attention_dim, + depth=4, + dim_head=64, + heads=20 if self.is_sdxl else 12, + num_queries=self.clip_extra_context_tokens, + embedding_dim=self.clip_embeddings_dim, + output_dim=self.output_cross_attention_dim, + ff_mult=4 + ) + 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 + +class CrossAttentionPatch: + # 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): + 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] + self.sigma_start = [sigma_start] + self.sigma_end = [sigma_end] + self.unfold_batch = [unfold_batch] + + 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): + 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) + self.unfold_batch.append(unfold_batch) + + 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 + + # 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 + + 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 + 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]) + + out_ip = out_ip * mask_downsample + + out = out + out_ip + + return out.to(dtype=org_dtype) + +class IPAdapterModelLoader: + @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,) + +class IPAdapterApply: + @classmethod + def INPUT_TYPES(s): + return { + "required": { + + "ipadapter": ("IPADAPTER", ), + "clip_vision": ("CLIP_VISION",), + "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 }), + "unfold_batch": ("BOOLEAN", { "default": False }), + }, + "optional": { + "attn_mask": ("MASK",), + } + } + + RETURN_TYPES = ("MODEL",) + 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): + for attr in list(self.__dict__): + delattr(self, attr) + self.dtype = model.model.diffusion_model.dtype + 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"] + + 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 + + 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") + + 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 + else: + clip_embed_zeroed = torch.zeros_like(clip_embed) + + clip_embeddings_dim = clip_embed.shape[-1] + + self.ipadapter = IPAdapter( + ipadapter, + cross_attention_dim=cross_attention_dim, + output_cross_attention_dim=output_cross_attention_dim, + clip_embeddings_dim=clip_embeddings_dim, + clip_extra_context_tokens=clip_extra_context_tokens, + is_sdxl=self.is_sdxl, + is_plus=self.is_plus, + is_full=self.is_full, + ) + + 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)) + 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() + + 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) + + 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, + } + + 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)) + 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)) + patch_kwargs["number"] += 1 + set_model_patch_replace(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)) + 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)) + patch_kwargs["number"] += 1 + for index in range(10): + set_model_patch_replace(work_model, patch_kwargs, ("middle", 0, index)) + patch_kwargs["number"] += 1 + + return (work_model, ) + + +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,) + +class IPAdapterEncoder: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "clip_vision": ("CLIP_VISION",), + "image_1": ("IMAGE",), + "ipadapter_plus": ("BOOLEAN", { "default": False }), + "noise": ("FLOAT", { "default": 0.0, "min": 0.0, "max": 1.0, "step": 0.01 }), + "weight_1": ("FLOAT", { "default": 1.0, "min": 0, "max": 1.0, "step": 0.01 }), + }, + "optional": { + "image_2": ("IMAGE",), + "image_3": ("IMAGE",), + "image_4": ("IMAGE",), + "weight_2": ("FLOAT", { "default": 1.0, "min": 0, "max": 1.0, "step": 0.01 }), + "weight_3": ("FLOAT", { "default": 1.0, "min": 0, "max": 1.0, "step": 0.01 }), + "weight_4": ("FLOAT", { "default": 1.0, "min": 0, "max": 1.0, "step": 0.01 }), + } + } + + RETURN_TYPES = ("EMBEDS",) + FUNCTION = "preprocess" + CATEGORY = "ipadapter" + + 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] + + if image_2 is not None: + if image_1.shape[1:] != image_2.shape[1:]: + image_2 = comfy.utils.common_upscale(image_2.movedim(-1,1), image.shape[2], image.shape[1], "bilinear", "center").movedim(1,-1) + image = torch.cat((image, image_2), dim=0) + weight += [weight_2]*image_2.shape[0] + if image_3 is not None: + if image.shape[1:] != image_3.shape[1:]: + image_3 = comfy.utils.common_upscale(image_3.movedim(-1,1), image.shape[2], image.shape[1], "bilinear", "center").movedim(1,-1) + image = torch.cat((image, image_3), dim=0) + weight += [weight_3]*image_3.shape[0] + if image_4 is not None: + if image.shape[1:] != image_4.shape[1:]: + image_4 = comfy.utils.common_upscale(image_4.movedim(-1,1), image.shape[2], image.shape[1], "bilinear", "center").movedim(1,-1) + image = torch.cat((image, image_4), dim=0) + weight += [weight_4]*image_4.shape[0] + + clip_embed = clip_vision.encode_image(image) + neg_image = image_add_noise(image, noise) if noise > 0 else None + + if ipadapter_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) + + if any(e != 1.0 for e in weight): + weight = torch.tensor(weight).unsqueeze(-1) if not ipadapter_plus else torch.tensor(weight).unsqueeze(-1).unsqueeze(-1) + clip_embed = clip_embed * weight + + output = torch.stack((clip_embed, clip_embed_zeroed)) + + return( output, ) + +class IPAdapterApplyEncoded(IPAdapterApply): + @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 IPAdapterSaveEmbeds: + 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 IPAdapterLoadEmbeds: + @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 IPAdapterBatchEmbeds: + @classmethod + def INPUT_TYPES(s): + return {"required": { + "embed1": ("EMBEDS",), + "embed2": ("EMBEDS",), + }} + + RETURN_TYPES = ("EMBEDS",) + FUNCTION = "batch" + CATEGORY = "ipadapter" + + def batch(self, embed1, embed2): + output = torch.cat((embed1, embed2), dim=1) + return (output, ) + +NODE_CLASS_MAPPINGS = { + "IPAdapterModelLoader": IPAdapterModelLoader, + "IPAdapterApply": IPAdapterApply, + "IPAdapterApplyEncoded": IPAdapterApplyEncoded, + "IPAdapterEncoder": IPAdapterEncoder, + "IPAdapterSaveEmbeds": IPAdapterSaveEmbeds, + "IPAdapterLoadEmbeds": IPAdapterLoadEmbeds, + "IPAdapterBatchEmbeds": IPAdapterBatchEmbeds, +} + +NODE_DISPLAY_NAME_MAPPINGS = { + "IPAdapterModelLoader": "Load IPAdapter Model", + "IPAdapterApply": "Apply IPAdapter", + "IPAdapterApplyEncoded": "Apply IPAdapter from Encoded", + "IPAdapterEncoder": "Encode IPAdapter Image", + "IPAdapterSaveEmbeds": "Save IPAdapter Embeds", + "IPAdapterLoadEmbeds": "Load IPAdapter Embeds", + "IPAdapterBatchEmbeds": "IPAdapter Batch Embeds", +} \ No newline at end of file diff --git a/control/film.py b/control/film.py deleted file mode 100644 index e1d2c0e..0000000 --- a/control/film.py +++ /dev/null @@ -1,201 +0,0 @@ -import torch -from comfy.model_management import get_torch_device, soft_empty_cache -import bisect -from typing import List -import numpy as np -import ast -import typing -import pathlib -import einops -import traceback -import os -from urllib.parse import urlparse -from torch.hub import download_url_to_file, get_dir -import yaml - - -BASE_MODEL_DOWNLOAD_URLS = [ - "https://github.com/styler00dollar/VSGAN-tensorrt-docker/releases/download/models/", - "https://github.com/Fannovel16/ComfyUI-Frame-Interpolation/releases/download/models/", - "https://github.com/dajes/frame-interpolation-pytorch/releases/download/v1.0.0/" -] - -config_path = os.path.join(os.path.dirname(__file__), "./config.yaml") -if os.path.exists(config_path): - config = yaml.load(open(config_path, "r"), Loader=yaml.FullLoader) -else: - raise Exception("config.yaml file is neccessary, plz recreate the config file by downloading it from https://github.com/Fannovel16/ComfyUI-Frame-Interpolation") -DEVICE = get_torch_device() - -def get_ckpt_container_path(model_type): - return os.path.abspath(os.path.join(os.path.dirname(__file__), config["ckpts_path"], model_type)) - - -def load_file_from_url(url, model_dir=None, progress=True, file_name=None): - - if model_dir is None: # use the pytorch hub_dir - hub_dir = get_dir() - model_dir = os.path.join(hub_dir, 'checkpoints') - - os.makedirs(model_dir, exist_ok=True) - - parts = urlparse(url) - file_name = os.path.basename(parts.path) - if file_name is not None: - file_name = file_name - cached_file = os.path.abspath(os.path.join(model_dir, file_name)) - if not os.path.exists(cached_file): - print(f'Downloading: "{url}" to {cached_file}\n') - download_url_to_file(url, cached_file, hash_prefix=None, progress=progress) - return cached_file - -class InterpolationStateList(): - - def __init__(self, frame_indices: typing.List[int], is_skip_list: bool): - self.frame_indices = frame_indices - self.is_skip_list = is_skip_list - - def is_frame_skipped(self, frame_index): - is_frame_in_list = frame_index in self.frame_indices - return self.is_skip_list and is_frame_in_list or not self.is_skip_list and not is_frame_in_list - - -def load_file_from_github_release(model_type, ckpt_name): - error_strs = [] - for i, base_model_download_url in enumerate(BASE_MODEL_DOWNLOAD_URLS): - try: - return load_file_from_url(base_model_download_url + ckpt_name, get_ckpt_container_path(model_type)) - except Exception: - traceback_str = traceback.format_exc() - if i < len(BASE_MODEL_DOWNLOAD_URLS) - 1: - print("Failed! Trying another endpoint.") - error_strs.append(f"Error when downloading from: {base_model_download_url + ckpt_name}\n\n{traceback_str}") - - error_str = '\n\n'.join(error_strs) - raise Exception(f"Tried all GitHub base urls to download {ckpt_name} but no suceess. Below is the error log:\n\n{error_str}") - - -def preprocess_frames(frames): - return einops.rearrange(frames, "n h w c -> n c h w") - -def postprocess_frames(frames): - return einops.rearrange(frames, "n c h w -> n h w c").cpu() - - -MODEL_TYPE = pathlib.Path(__file__).parent.name -DEVICE = get_torch_device() - -def inference(model, img_batch_1, img_batch_2, inter_frames): - results = [ - img_batch_1, - img_batch_2 - ] - - idxes = [0, inter_frames + 1] - remains = list(range(1, inter_frames + 1)) - - splits = torch.linspace(0, 1, inter_frames + 2) - - for _ in range(len(remains)): - starts = splits[idxes[:-1]] - ends = splits[idxes[1:]] - distances = ((splits[None, remains] - starts[:, None]) / (ends[:, None] - starts[:, None]) - .5).abs() - matrix = torch.argmin(distances).item() - start_i, step = np.unravel_index(matrix, distances.shape) - end_i = start_i + 1 - - x0 = results[start_i].to(DEVICE) - x1 = results[end_i].to(DEVICE) - dt = x0.new_full((1, 1), (splits[remains[step]] - splits[idxes[start_i]])) / (splits[idxes[end_i]] - splits[idxes[start_i]]) - - with torch.no_grad(): - prediction = model(x0, x1, dt) - insert_position = bisect.bisect_left(idxes, remains[step]) - idxes.insert(insert_position, remains[step]) - results.insert(insert_position, prediction.clamp(0, 1).float()) - del remains[step] - - return [tensor.flip(0) for tensor in results] - - -def film_interpolation( - frames: torch.Tensor, - frame_counts: List[int] = [0,5,32,40], - buffer: int = 10): - - frame_counts = sorted(frame_counts) - - max_gap = max(b-a for a, b in zip(frame_counts[:-1], frame_counts[1:])) - - model_path = load_file_from_github_release(MODEL_TYPE, "film_net_fp32.pt") - model = torch.jit.load(model_path, map_location='cpu') - model.eval() - model = model.to(DEVICE) - - frames = preprocess_frames(frames) - number_of_frames_processed_since_last_cleared_cuda_cache = 0 - clear_cache_after_n_frames = 10 # Example value - - # Generate buffer frames and attach them to the beginning of output_frames - first_frame = frames[0].unsqueeze(0) - buffer_frames = [first_frame] * buffer - output_frames = buffer_frames - - - for frame_itr in range(len(frames) - 1): - - - frame_0 = frames[frame_itr:frame_itr+1].to(DEVICE) - frame_1 = frames[frame_itr+1:frame_itr+2].to(DEVICE) - frame_output = [] - result = inference(model, frame_0, frame_1, max_gap - 1) - - # Find the current frame's position in the frame_counts list - current_frame = frame_counts[frame_itr] - next_frame = frame_counts[frame_itr + 1] - 1 - current_gap = next_frame - current_frame - # Determine the number of frames to drop based on the difference between the max gap and the current gap - frames_to_drop = max_gap - current_gap - - frame_output = result[:-1] - - # frames_to_drop = max_gap - 1 - len(frame_output) - if frames_to_drop > 0: - if frames_to_drop >= len(frame_output): - raise ValueError("Number of frames to drop is greater than or equal to total number of frames in the batch.") - - drop_interval = len(frame_output) / float(frames_to_drop) - result = [] - next_drop = drop_interval - - for i, frame in enumerate(frame_output): - if i >= next_drop: - next_drop += drop_interval - else: - result.append(frame) - - frame_output = result # Update frame_output with only the undropped frames - - # Detach and move to CPU all frames, whether or not any were dropped - frame_output = [frame.detach().cpu() for frame in frame_output] - - # Append the processed (and potentially dropped) frames to the main list - output_frames.extend(frame_output) - - number_of_frames_processed_since_last_cleared_cuda_cache += 1 - if number_of_frames_processed_since_last_cleared_cuda_cache >= clear_cache_after_n_frames: - print("Comfy-VFI: Clearing cache...") - soft_empty_cache() - number_of_frames_processed_since_last_cleared_cuda_cache = 0 - print("Comfy-VFI: Done cache clearing") - - output_frames.append(frames[-1:]) - out = torch.cat(output_frames, dim=0) - - # clear cache for courtesy - print("Comfy-VFI: Final clearing cache...") - soft_empty_cache() - print("Comfy-VFI: Done cache clearing") - - - return (postprocess_frames(out), ) \ No newline at end of file diff --git a/control/latent_keyframe_nodes.py b/control/latent_keyframe_nodes.py index baa0185..faf932d 100644 --- a/control/latent_keyframe_nodes.py +++ b/control/latent_keyframe_nodes.py @@ -164,63 +164,96 @@ class LatentKeyframeInterpolationNodeImport: interpolation: str, revert_direction_at_midpoint: bool=False, last_key_frame_position: int=0, + i=0, + number_of_items=0, + buffer=0, prev_latent_keyframe: LatentKeyframeGroupImport=None): - if (batch_index_from > batch_index_to_excl): - raise ValueError("batch_index_from must be less than or equal to batch_index_to.") - if (batch_index_from < 0 and batch_index_to_excl >= 0): - raise ValueError("batch_index_from and batch_index_to must be either both positive or both negative.") - if not prev_latent_keyframe: prev_latent_keyframe = LatentKeyframeGroupImport() + curr_latent_keyframe = LatentKeyframeGroupImport() - steps = batch_index_to_excl - batch_index_from - diff = strength_to - strength_from - - index = np.linspace(0, 1, steps // 2 + 1) if revert_direction_at_midpoint else np.linspace(0, 1, steps) - - if interpolation == "linear": - weights = np.linspace(strength_from, strength_to, len(index)) - elif interpolation == "ease-in": - weights = diff * np.power(index, 2) + strength_from - elif interpolation == "ease-out": - weights = diff * (1 - np.power(1 - index, 2)) + strength_from - elif interpolation == "ease-in-out": - weights = diff * ((1 - np.cos(index * np.pi)) / 2) + strength_from - - if revert_direction_at_midpoint: - weights = np.concatenate([weights, weights[::-1]]) - - # Generate frame numbers - frame_numbers = np.arange(batch_index_from, batch_index_from + len(weights)) - - # "Dropper" component: For keyframes with negative start, drop the weights - if batch_index_from < 0: - drop_count = abs(batch_index_from) - weights = weights[drop_count:] - frame_numbers = frame_numbers[drop_count:] - - # Dropper component: for keyframes a batch_index_to_excl is greater than last_key_frame_position, drop the weights - if batch_index_to_excl > last_key_frame_position: - drop_count = batch_index_to_excl - last_key_frame_position - weights = weights[:-drop_count] - frame_numbers = frame_numbers[:-drop_count] - - batch_index_from = frame_numbers[0] + weights, frame_numbers = calculate_weights(batch_index_from, batch_index_to_excl, strength_from, strength_to, interpolation, revert_direction_at_midpoint, last_key_frame_position,i,number_of_items, buffer) - for i in range(len(frame_numbers)): - keyframe = LatentKeyframeImport(batch_index_from + i, float(weights[i])) - logger.info(f"keyframe {batch_index_from + i}:{weights[i]}") + for i, frame_number in enumerate(frame_numbers): + keyframe = LatentKeyframeImport(frame_number, float(weights[i])) + logger.info(f"keyframe {frame_number}:{weights[i]}") curr_latent_keyframe.add(keyframe) for latent_keyframe in prev_latent_keyframe.keyframes: curr_latent_keyframe.add(latent_keyframe) - return (curr_latent_keyframe,) + return (weights, frame_numbers, curr_latent_keyframe,) + + +def calculate_weights(batch_index_from, batch_index_to, strength_from, strength_to, interpolation,revert_direction_at_midpoint, last_key_frame_position,i, number_of_items,buffer): + print("<-------- LOOK HERE s") + print("batch_index_from",batch_index_from) + print("batch_index_to",batch_index_to) + print("strength_from",strength_from) + print("strength_to",strength_to) + print("interpolation",interpolation) + print("revert_direction_at_midpoint",revert_direction_at_midpoint) + print("last_key_frame_position",last_key_frame_position) + print("i",i) + print("number_of_items",number_of_items) + # Initialize variables based on the position of the keyframe + range_start = batch_index_from + range_end = batch_index_to + # if it's the first value, set influence range from 1.0 to 0.0 + if buffer > 0: + if i == 0: + range_start = 0 + elif i == 1: + range_start = buffer + else: + if i == 1: + range_start = 0 + + if i == number_of_items - 1: + range_end = last_key_frame_position + + steps = range_end - range_start + diff = strength_to - strength_from + + # Calculate index for interpolation + index = np.linspace(0, 1, steps // 2 + 1) if revert_direction_at_midpoint else np.linspace(0, 1, steps) + + # Calculate weights based on interpolation type + if interpolation == "linear": + weights = np.linspace(strength_from, strength_to, len(index)) + elif interpolation == "ease-in": + weights = diff * np.power(index, 2) + strength_from + elif interpolation == "ease-out": + weights = diff * (1 - np.power(1 - index, 2)) + strength_from + elif interpolation == "ease-in-out": + weights = diff * ((1 - np.cos(index * np.pi)) / 2) + strength_from + + # If it's a middle keyframe, mirror the weights + if revert_direction_at_midpoint: + weights = np.concatenate([weights, weights[::-1]]) + + # Generate frame numbers + frame_numbers = np.arange(range_start, range_start + len(weights)) + + # "Dropper" component: For keyframes with negative start, drop the weights + if range_start < 0 and i > 0: + drop_count = abs(range_start) + weights = weights[drop_count:] + frame_numbers = frame_numbers[drop_count:] + + # Dropper component: for keyframes a range_End is greater than last_key_frame_position, drop the weights + if range_end > last_key_frame_position and i < number_of_items - 1: + drop_count = range_end - last_key_frame_position + weights = weights[:-drop_count] + frame_numbers = frame_numbers[:-drop_count] + + return weights, frame_numbers + class LatentKeyframeBatchedGroupNodeImport: diff --git a/control/nodes.py b/control/nodes.py index 0d68c80..5e1f31e 100644 --- a/control/nodes.py +++ b/control/nodes.py @@ -1,16 +1,28 @@ import numpy as np import torch import folder_paths +from PIL import Image from ast import literal_eval from .control import ControlNetAdvancedImport, T2IAdapterAdvancedImport, load_controlnet, ControlNetWeightsTypeImport, T2IAdapterWeightsTypeImport,\ LatentKeyframeGroupImport, TimestepKeyframeImport, TimestepKeyframeGroupImport, is_advanced_controlnet +import matplotlib.pyplot as plt +from .IPAdapterPlus import contrast_adaptive_sharpening, IPAdapterApply,prep_image from .weight_nodes import ScaledSoftControlNetWeightsImport, SoftControlNetWeightsImport, CustomControlNetWeightsImport, \ SoftT2IAdapterWeightsImport, CustomT2IAdapterWeightsImport -from .latent_keyframe_nodes import LatentKeyframeGroupNodeImport, LatentKeyframeInterpolationNodeImport, LatentKeyframeBatchedGroupNodeImport, LatentKeyframeNodeImport +from .latent_keyframe_nodes import LatentKeyframeGroupNodeImport, LatentKeyframeInterpolationNodeImport, LatentKeyframeBatchedGroupNodeImport, LatentKeyframeNodeImport,calculate_weights from .deprecated_nodes import LoadImagesFromDirectory from .logger import logger -from .film import film_interpolation +import torchvision.transforms as TT +import torch.nn.functional as F + +import comfy.utils +import comfy.model_management +from comfy.clip_vision import clip_preprocess +from comfy.ldm.modules.attention import optimized_attention +# import BytesIO +from io import BytesIO + class TimestepKeyframeNodeImport: @@ -196,35 +208,45 @@ class BatchCreativeInterpolationNode: "required": { "positive": ("CONDITIONING", ), "negative": ("CONDITIONING", ), - "control_net_name": (folder_paths.get_filename_list("controlnet"), ), "images": ("IMAGE", ), + "model": ("MODEL", ), + "ipadapter": ("IPADAPTER", ), + "clip_vision": ("CLIP_VISION",), + "control_net_name": (folder_paths.get_filename_list("controlnet"), ), "type_of_frame_distribution": (["linear", "dynamic"],), "linear_frame_distribution_value": ("INT", {"default": 16, "min": 4, "max": 64, "step": 1}), "dynamic_frame_distribution_values": ("STRING", {"multiline": True, "default": "0,10,26,40"}), "type_of_key_frame_influence": (["linear", "dynamic"],), - "linear_key_frame_influence_value": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.001}), + "linear_key_frame_influence_value": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 10.0, "step": 0.1}), "dynamic_key_frame_influence_values": ("STRING", {"multiline": True, "default": "1.0,1.0,1.0,0.5"}), "type_of_cn_strength_distribution": (["linear", "dynamic"],), - "linear_cn_strength_value": ("STRING", {"multiline": False, "default": "(0.0,1.0)"}), + "linear_cn_strength_value": ("STRING", {"multiline": False, "default": "(0.0,0.4)"}), "dynamic_cn_strength_values": ("STRING", {"multiline": True, "default": "(0.0,1.0),(0.0,1.0),(0.0,1.0),(0.0,1.0)"}), - "soft_scaled_cn_weights_multiplier": ("FLOAT", {"default": 0.85, "min": 0.0, "max": 10.0, "step": 0.01}), - "interpolation": (["ease-in", "ease-out", "ease-in-out"],), - "buffer": ("INT", {"default": 4, "min": 1, "max": 16, "step": 1}), - "intermediate_frame_mask_strength": ("FLOAT", {"default": 0.5, "min": 0.0, "max": 1.0, "step": 0.01}), + "soft_scaled_cn_weights_multiplier": ("FLOAT", {"default": 0.85, "min": 0.0, "max": 10.0, "step": 0.1}), + # "interpolation": (["ease-in-out", "ease-in", "ease-out"],), + "buffer": ("INT", {"default": 4, "min": 0, "max": 16, "step": 1}), + "relative_ipadapter_strength": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}), + "relative_ipadapter_influence": ("FLOAT", {"default": 1.0, "min": 0.0, "max": 5.0, "step": 0.1}), + "ipadapter_noise": ("FLOAT", {"default": 0.3, "min": 0.0, "max": 1.0, "step": 0.01}), }, "optional": { } } - RETURN_TYPES = ("CONDITIONING","CONDITIONING","IMAGE","MASK") - RETURN_NAMES = ("positive", "negative") + RETURN_TYPES = ("IMAGE","CONDITIONING","CONDITIONING","MODEL",) + RETURN_NAMES = ("GRAPH","POSITIVE", "NEGATIVE","MODEL") FUNCTION = "combined_function" CATEGORY = "Steerable-Motion/Interpolation" - def combined_function(self, positive, negative, control_net_name, images,type_of_frame_distribution,linear_frame_distribution_value,dynamic_frame_distribution_values,type_of_key_frame_influence,linear_key_frame_influence_value,dynamic_key_frame_influence_values,type_of_cn_strength_distribution,linear_cn_strength_value,dynamic_cn_strength_values,soft_scaled_cn_weights_multiplier,interpolation,buffer,intermediate_frame_mask_strength): + def combined_function(self, positive, negative, images,model,ipadapter,clip_vision,control_net_name, + type_of_frame_distribution,linear_frame_distribution_value,dynamic_frame_distribution_values, + type_of_key_frame_influence,linear_key_frame_influence_value,dynamic_key_frame_influence_values, + type_of_cn_strength_distribution,linear_cn_strength_value,dynamic_cn_strength_values, + soft_scaled_cn_weights_multiplier,buffer,relative_ipadapter_strength, + relative_ipadapter_influence,ipadapter_noise): - def calculate_dynamic_influence_ranges(keyframe_positions, key_frame_influence_values): + def calculate_dynamic_influence_ranges(keyframe_positions, key_frame_influence_values, allow_extension=True): if len(keyframe_positions) < 2 or len(keyframe_positions) != len(key_frame_influence_values): return [] @@ -240,8 +262,9 @@ class BatchCreativeInterpolationNode: end_influence = position + range_size # Adjust start and end influence to not exceed previous and next keyframes - start_influence = max(start_influence, keyframe_positions[i - 1] if i > 0 else 0) - end_influence = min(end_influence, keyframe_positions[i + 1] if i < len(keyframe_positions) - 1 else keyframe_positions[-1]) + if not allow_extension: + start_influence = max(start_influence, keyframe_positions[i - 1] if i > 0 else 0) + end_influence = min(end_influence, keyframe_positions[i + 1] if i < len(keyframe_positions) - 1 else keyframe_positions[-1]) influence_ranges.append((round(start_influence), round(end_influence))) @@ -298,143 +321,218 @@ class BatchCreativeInterpolationNode: # Return a list of tuples with the linear_key_frame_influence_value as a tuple repeated for each position return [linear_key_frame_influence_value for _ in keyframe_positions] - def create_mask_batch(keyframe_positions, buffer, masked_strength): + def create_mask_batch(last_key_frame_position, weights, frames): # Hardcoded dimensions width, height = 512, 512 - def generate_masks(number_of_masks, strength): - # Create masks with the specified strength - masks = [] - for _ in range(number_of_masks): - mask = torch.full((height, width), strength) - masks.append(mask) + # Calculate the reversed weights in a generalizable way (e.g., 0.6 becomes 0.4, 0.1 becomes 0.9) + reversed_weights = [1.0 - weight for weight in weights] - # Convert list of masks to a single tensor - masks_tensor = torch.stack(masks, dim=0) - return masks_tensor + # Map frames to their corresponding reversed weights for easy lookup + frame_to_weight = {frame: weights[i] for i, frame in enumerate(frames)} + # Create masks for each frame up to last_key_frame_position masks = [] + for frame_number in range(last_key_frame_position): + # Determine the strength of the mask + strength = frame_to_weight.get(frame_number, 0.0) - # Creating initial buffer masks if buffer > 0 - if buffer > 0: - buffer_masks = generate_masks(buffer, 1.0) - masks.append(buffer_masks) + # Create the mask with the determined strength + mask = torch.full((height, width), strength) + masks.append(mask) - # Iterating through the keyframe positions - last_position = -1 - for pos in keyframe_positions: - # Number of masks to create between last position and current keyframe - num_masks = pos - last_position + # Convert list of masks to a single tensor + masks_tensor = torch.stack(masks, dim=0) - if not (pos == 0 and buffer > 0): - # Create masks of strength 0.0 for the keyframe itself - keyframe_mask = generate_masks(1, 1.0) - masks.append(keyframe_mask) - - if num_masks > 1: - # Create masks of masked_strength for positions between keyframes - intermediate_masks = generate_masks(num_masks - 1, masked_strength) - masks.append(intermediate_masks) - - last_position = pos - - # Ensure a mask of 0.0 strength is created at the last keyframe position - if keyframe_positions[-1] == last_position: - last_keyframe_mask = generate_masks(1, 1.0) - masks.append(last_keyframe_mask) - - # Flatten the list of tensors into a single tensor - masks_tensor = torch.cat(masks, dim=0) return masks_tensor + + def adjust_influence_range(batch_index_from, batch_index_to_excl, last_key_frame_position, scale_factor, buffer): + # Calculate the midpoint of the current range + midpoint = (batch_index_from + batch_index_to_excl) // 2 + # Calculate the new range length + new_range_length = int((batch_index_to_excl - batch_index_from) * scale_factor) + # Adjusting both sides of the range + if batch_index_from == 0: + # Start is anchored at 0 + new_batch_index_from = 0 + new_batch_index_to_excl = batch_index_from + new_range_length + elif batch_index_to_excl == last_key_frame_position: + # End is anchored at last_key_frame_position + new_batch_index_from = batch_index_to_excl - new_range_length + new_batch_index_to_excl = last_key_frame_position + else: + # No anchoring, adjust both sides around the midpoint + new_batch_index_from = midpoint - new_range_length // 2 + new_batch_index_to_excl = midpoint + new_range_length // 2 + # Remove minimum and maximum constraints + + return new_batch_index_from, new_batch_index_to_excl + + def adjust_strength_values(strength_from, strength_to, multiplier): + mid_point = (strength_from + strength_to) / 2 + range_half = abs(strength_to - strength_from) / 2 + + # Adjust the range with the multiplier + new_range_half = min(range_half * multiplier, 0.5) + + # Calculate new strength values, ensuring they stay within [0.0, 1.0] + new_strength_from = max(mid_point - new_range_half, 0.0) + new_strength_to = min(mid_point + new_range_half, 1.0) + + # Preserve the order of the original strength values + if strength_from > strength_to: + new_strength_from, new_strength_to = new_strength_to, new_strength_from + + return (new_strength_from, new_strength_to) + def plot_weight_comparison(cn_frame_numbers, cn_weights, ipadapter_frame_numbers, ipadapter_weights, buffer): + plt.figure(figsize=(12, 8)) + + # Defining colors for each set of data + colors = ['b', 'g', 'r', 'c', 'm', 'y', 'k'] + + # Alternating the data sets with labels and colors + max_length = max(len(cn_frame_numbers), len(ipadapter_frame_numbers)) + label_counter = 1 if buffer < 0 else 0 # Start from 1 if buffer < 0, else start from 0 + for i in range(max_length): + # Label for cn_strength + if i < len(cn_frame_numbers): + if i == 0 and buffer > 0: + label = 'cn_strength_buffer' + else: + label = f'cn_strength_{label_counter}' + plt.plot(cn_frame_numbers[i], cn_weights[i], marker='o', color=colors[i % len(colors)], label=label) + + # Label for ipa_strength + if i < len(ipadapter_frame_numbers): + if i == 0 and buffer > 0: + label = 'ipa_strength_buffer' + else: + label = f'ipa_strength_{label_counter}' + plt.plot(ipadapter_frame_numbers[i], ipadapter_weights[i], marker='x', linestyle='--', color=colors[i % len(colors)], label=label) + + if label_counter == 0 or buffer < 0 or i > 0: + label_counter += 1 + + plt.legend() + max_weight = max([weight.max() for weight in cn_weights + ipadapter_weights]) * 1.5 + plt.ylim(0, max_weight) + + buffer_io = BytesIO() + plt.savefig(buffer_io, format='png', bbox_inches='tight') + plt.close() + + buffer_io.seek(0) + img = Image.open(buffer_io) + + img_tensor = TT.ToTensor()(img) + img_tensor = img_tensor.unsqueeze(0) + + img_tensor = img_tensor.permute([0, 2, 3, 1]) + + return (img_tensor,) + keyframe_positions = get_keyframe_positions(type_of_frame_distribution, dynamic_frame_distribution_values, images, linear_frame_distribution_value) cn_strength_values = extract_start_and_endpoint_values(type_of_cn_strength_distribution, dynamic_cn_strength_values, keyframe_positions, linear_cn_strength_value) key_frame_influence_values = extract_keyframe_values(type_of_key_frame_influence, dynamic_key_frame_influence_values, keyframe_positions, linear_key_frame_influence_value) influence_ranges = calculate_dynamic_influence_ranges(keyframe_positions,key_frame_influence_values) - if buffer > 0: - influence_ranges = add_starting_buffer(influence_ranges, buffer) - print(f"keyframe_positions: {keyframe_positions}") + + influence_ranges = add_starting_buffer(influence_ranges, buffer) cn_strength_values = [literal_eval(val) if isinstance(val, str) else val for val in cn_strength_values] - ipadapter_input, = film_interpolation(images, keyframe_positions, buffer) - - - generated_masks_tensor = create_mask_batch(keyframe_positions, buffer, intermediate_frame_mask_strength) + cn_frame_numbers = [] + cn_weights = [] + ipadapter_frame_numbers = [] + ipadapter_weights = [] - last_key_frame_position = (keyframe_positions[-1]) + buffer control_net = [] for i, (start, end) in enumerate(influence_ranges): batch_index_from, batch_index_to_excl = influence_ranges[i] + ipadapter_strength_multiplier = relative_ipadapter_strength + ipadapter_influence_multiplier = relative_ipadapter_influence - if i == 0 and buffer > 0: # First image with buffer - image = images[0] - strength_from = strength_to = cn_strength_values[0][1] if len(cn_strength_values) > 0 else (1.0, 1.0) - return_at_midpoint = False + if i == 0: + if buffer > 0: # First image with buffer + image = images[0] + strength_from = strength_to = cn_strength_values[0][1] if len(cn_strength_values) > 0 else (1.0, 1.0) + revert_direction_at_midpoint = False + ipadapter_strength_multiplier = 1.0 + ipadapter_influence_multiplier = 1.0 + interpolation = "ease-in-out" + else: + continue # Skip first image without buffer elif i == 1: # First image image = images[0] strength_to, strength_from = cn_strength_values[0] if len(cn_strength_values) > 0 else (0.0, 1.0) - return_at_midpoint = False + revert_direction_at_midpoint = False + interpolation = "ease-in" + + + elif i == len(images): # Last image image = images[i-1] strength_from, strength_to = cn_strength_values[i-1] if i-1 < len(cn_strength_values) else (0.0, 1.0) - return_at_midpoint = False + revert_direction_at_midpoint = False + interpolation = "ease-out" + else: # Middle images image = images[i-1] strength_from, strength_to = cn_strength_values[i-1] if i-1 < len(cn_strength_values) else (0.0, 1.0) - return_at_midpoint = True - + revert_direction_at_midpoint = True + interpolation = "ease-in-out" + + latent_keyframe_interpolation_node = LatentKeyframeInterpolationNodeImport() - latent_keyframe, = latent_keyframe_interpolation_node.load_keyframe( - batch_index_from, - strength_from, - batch_index_to_excl, - strength_to, - interpolation, - return_at_midpoint, - last_key_frame_position) - - scaled_soft_control_net_weights = ScaledSoftControlNetWeightsImport() - control_net_weights, _ = scaled_soft_control_net_weights.load_weights( - soft_scaled_cn_weights_multiplier, - False) + weights, frame_numbers, latent_keyframe, = latent_keyframe_interpolation_node.load_keyframe(batch_index_from,strength_from,batch_index_to_excl,strength_to,interpolation,revert_direction_at_midpoint,last_key_frame_position,i,len(influence_ranges),buffer) + cn_frame_numbers.append(frame_numbers) + cn_weights.append(weights) + + scaled_soft_control_net_weights = ScaledSoftControlNetWeightsImport() + control_net_weights, _ = scaled_soft_control_net_weights.load_weights(soft_scaled_cn_weights_multiplier,False) timestep_keyframe_node = TimestepKeyframeNodeImport() - timestep_keyframe, = timestep_keyframe_node.load_keyframe( - start_percent=0.0, - control_net_weights=control_net_weights, - t2i_adapter_weights=None, - latent_keyframe=latent_keyframe, - prev_timestep_keyframe=None - ) + timestep_keyframe, = timestep_keyframe_node.load_keyframe(start_percent=0.0,control_net_weights=control_net_weights,t2i_adapter_weights=None,latent_keyframe=latent_keyframe,prev_timestep_keyframe=None) control_net_loader = ControlNetLoaderAdvancedImport() - control_net, = control_net_loader.load_controlnet( - control_net_name, - timestep_keyframe) + control_net, = control_net_loader.load_controlnet(control_net_name, timestep_keyframe) apply_advanced_control_net = AdvancedControlNetApplyImport() - positive, negative = apply_advanced_control_net.apply_controlnet( - positive, - negative, - control_net, - image.unsqueeze(0), - 1.0, - 0.0, - 1.0) + positive, negative = apply_advanced_control_net.apply_controlnet(positive,negative,control_net,image.unsqueeze(0),1.0,0.0,1.0) + + prepped_image, = prep_image(image=image.unsqueeze(0), interpolation="LANCZOS", crop_position="pad", sharpening=0.0) + + ipadapter_application = IPAdapterApply() + + ipa_strength_from, ipa_strength_to = adjust_strength_values(strength_from, strength_to, ipadapter_strength_multiplier) + + ipa_batch_index_from, ipa_batch_index_to_excl = adjust_influence_range(batch_index_from, batch_index_to_excl, last_key_frame_position, ipadapter_influence_multiplier, buffer) + + ipa_weights, ipa_frame_numbers = calculate_weights(ipa_batch_index_from,ipa_batch_index_to_excl,ipa_strength_from, ipa_strength_to,interpolation, revert_direction_at_midpoint, last_key_frame_position,i, len(influence_ranges),buffer) + + + ipadapter_frame_numbers.append(ipa_frame_numbers) + ipadapter_weights.append(ipa_weights) + + + masks = create_mask_batch(last_key_frame_position, weights, frame_numbers) + + model, = ipadapter_application.apply_ipadapter(ipadapter=ipadapter, model=model, weight=1.0, clip_vision=clip_vision, image=prepped_image, weight_type="original", noise=ipadapter_noise, embeds=None, attn_mask=masks, start_at=0.0, end_at=1.0, unfold_batch=True) - return positive, negative, ipadapter_input,generated_masks_tensor + comparison_diagram, = plot_weight_comparison(cn_frame_numbers, cn_weights, ipadapter_frame_numbers, ipadapter_weights, buffer) + + return comparison_diagram, positive, negative, model # NODE MAPPING NODE_CLASS_MAPPINGS = { # Combined - "BatchCreativeInterpolation": BatchCreativeInterpolationNode, - "MaskGenerator": MaskGeneratorNode + "BatchCreativeInterpolation": BatchCreativeInterpolationNode + # "MaskGenerator": MaskGeneratorNode # "FILMVFIImport": FILMVFINode # Keyframes # "TimestepKeyframe": TimestepKeyframeNodeImport, @@ -459,8 +557,8 @@ NODE_CLASS_MAPPINGS = { NODE_DISPLAY_NAME_MAPPINGS = { # Combined - "BatchCreativeInterpolation": "Batch Creative Interpolation 🎞️🅟🅞🅜", - "MaskGenerator": "Mask Generator 🎞️🅟🅞🅜" + "BatchCreativeInterpolation": "Batch Creative Interpolation 🎞️🅟🅞🅜" + # "MaskGenerator": "Mask Generator 🎞️🅟🅞🅜" # Keyframes # "TimestepKeyframe": "Timestep Keyframe 🎞️🅟🅞🅜", # "LatentKeyframe": "Latent Keyframe 🛂🅐🅒🅝", diff --git a/control/resampler.py b/control/resampler.py new file mode 100644 index 0000000..4521c8c --- /dev/null +++ b/control/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