diff --git a/IPAdapterPlus.py b/IPAdapterPlus.py index 272a361..6bf1a7f 100644 --- a/IPAdapterPlus.py +++ b/IPAdapterPlus.py @@ -14,7 +14,7 @@ from PIL import Image import torch.nn.functional as F import torchvision.transforms as TT -from .resampler import Resampler +from .resampler import PerceiverAttention, FeedForward, Resampler # set the models directory backward compatible GLOBAL_MODELS_DIR = os.path.join(folder_paths.models_dir, "ipadapter") @@ -27,8 +27,43 @@ folder_paths.folder_names_and_paths["ipadapter"] = (current_paths, folder_paths. INSIGHTFACE_DIR = os.path.join(folder_paths.models_dir, "insightface") +class FacePerceiverResampler(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( + [ + PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads), + FeedForward(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 MLPProjModel(torch.nn.Module): - """SD model with image prompt""" def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024): super().__init__() @@ -44,7 +79,6 @@ class MLPProjModel(torch.nn.Module): return clip_extra_context_tokens class MLPProjModelFaceId(torch.nn.Module): - """SD model with image prompt""" def __init__(self, cross_attention_dim=768, id_embeddings_dim=512, num_tokens=4): super().__init__() @@ -64,6 +98,37 @@ class MLPProjModelFaceId(torch.nn.Module): clip_extra_context_tokens = self.norm(clip_extra_context_tokens) return clip_extra_context_tokens +class ProjModelFaceIdPlus(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 = FacePerceiverResampler( + 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): + x = self.proj(id_embeds) + x = x.reshape(-1, self.num_tokens, self.cross_attention_dim) + x = self.norm(x) + x = self.perceiver_resampler(x, clip_embeds) + return x + class ImageProjModel(nn.Module): def __init__(self, cross_attention_dim=1024, clip_embeddings_dim=1024, clip_extra_context_tokens=4): super().__init__() @@ -176,20 +241,11 @@ def contrast_adaptive_sharpening(image, amount): return (output) -def tensorToCV(image): - import numpy as np +def tensorToNP(image): + out = torch.clamp(255. * image.detach().cpu(), 0, 255).to(torch.uint8) + out = out[..., [2, 1, 0]] + out = out.numpy() - # TODO: there must be a better way - out = [] - image = image.detach().cpu() - for i in range(image.shape[0]): - img = image[i] - img = img[..., [2, 1, 0]] # Convert from RGB to BGR - img = torch.clamp(img * 255, 0, 255).to(torch.uint8) # Scale pixel values up to [0, 255] if they are in [0, 1] - img = img.numpy() - out.append(img) - - out = np.stack(img, axis=0) return out class IPAdapter(nn.Module): @@ -202,6 +258,7 @@ class IPAdapter(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() @@ -241,11 +298,19 @@ class IPAdapter(nn.Module): return image_proj_model def init_proj_faceid(self): - image_proj_model = MLPProjModelFaceId( - cross_attention_dim=self.cross_attention_dim, - id_embeddings_dim=512, - num_tokens=self.clip_extra_context_tokens, - ) + if self.is_plus: + image_proj_model = ProjModelFaceIdPlus( + cross_attention_dim=self.cross_attention_dim, + id_embeddings_dim=512, + clip_embeddings_dim=1280, + num_tokens=4, + ) + else: + image_proj_model = MLPProjModelFaceId( + 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() @@ -254,6 +319,11 @@ class IPAdapter(nn.Module): 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): + embeds = self.image_proj_model(face_embed, clip_embed) + return embeds + class CrossAttentionPatch: # forward for patching def __init__(self, weight, ipadapter, number, cond, uncond, weight_type, mask=None, sigma_start=0.0, sigma_end=1.0, unfold_batch=False): @@ -447,7 +517,7 @@ class InsightFaceLoader: }, } - RETURN_TYPES = ("CLIP_VISION",) + RETURN_TYPES = ("INSIGHTFACE",) FUNCTION = "load_insight_face" CATEGORY = "ipadapter" @@ -481,6 +551,7 @@ class IPAdapterApply: }, "optional": { "attn_mask": ("MASK",), + "insightface": ("INSIGHTFACE",), } } @@ -488,13 +559,16 @@ class IPAdapterApply: 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): + 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): 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_faceid = "0.to_q_lora.down.weight" in ipadapter["ip_adapter"] # TODO: better way to detect faceid? - self.is_plus = (self.is_full or "latents" in ipadapter["image_proj"]) and not self.is_faceid + self.is_full = "proj.3.weight" in ipadapter["image_proj"] + self.is_faceid = "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"]) + + 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 @@ -507,13 +581,37 @@ class IPAdapterApply: clip_embed_zeroed = embeds[1].cpu() else: if self.is_faceid: - clip_embed = clip_vision.get(tensorToCV(image)) # TODO: support multiple images (is it needed?) - if not clip_embed: - print("\033[33mWARNING!!! InsightFace wasn't able to detect the face. Try to use the PrepImageForInsightFace node.\033[0m") + insightface.det_model.input_size = (640,640) # reset the detection size + face_img = tensorToNP(image) + face_embed = [] - #clip_embed = torch.from_numpy(clip_embed[0].normed_embedding).unsqueeze(0) - clip_embed = torch.stack([torch.from_numpy(elem.normed_embedding).unsqueeze(0) for elem in clip_embed], dim=0) - clip_embed_zeroed = torch.zeros_like(clip_embed) + 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)) + if 640 not in size: + print(f"\033[33mINFO: InsightFace detection resolution lowered to {size}.\033[0m") + break + else: + print("\033[33mWARNING!!! InsightFace didn't detect any face.\033[0m") + + face_embed = torch.stack(face_embed, 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") @@ -550,7 +648,12 @@ class IPAdapterApply: 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, dtype=self.dtype), clip_embed_zeroed.to(self.device, dtype=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)) + 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)) + 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) @@ -667,12 +770,11 @@ class PrepImageForInsightFace: def prep_image(self, image, crop_position, sharpening=0.0, pad_around=True): if pad_around: - padding = 60 + padding = 30 size = (580, 580) else: padding = 0 size = (640, 640) - output = prepImage(image, "LANCZOS", crop_position, size, sharpening, padding) return (output, ) diff --git a/README.md b/README.md index 179f71b..26eba4a 100644 --- a/README.md +++ b/README.md @@ -5,6 +5,8 @@ IPAdapter implementation that follows the ComfyUI way of doing things. The code ## Important updates +**2023/12/28**: Added support for FaceID Plus models. **Important:** this update breaks the previous implementation of FaceID. Check the updated workflows in the exmaple directory! + **2023/12/22**: Added support for FaceID models. [Read the documentation](#faceid) for details. **2023/12/05**: Added `batch embeds` node. This lets you encode images in batches and merge them together into an `IPAdapter Apply Encoded` node. Useful mostly for animations because the clip vision encoder takes a lot of VRAM. My suggestion is to split the animation in batches of about 120 frames. @@ -59,7 +61,7 @@ IPAdapter also needs the image encoders. You need the [CLIP-**ViT-H**-14-laion2B The following table shows the combination of Checkpoint and Image encoder to use for each IPAdapter Model. Any Tensor size error you may get it is likely caused by a wrong combination. -| SD v. | IPadapter | Img encoder | Nodes | +| SD v. | IPadapter | Img encoder | Notes | |---|---|---|---| | v1.5 | [ip-adapter_sd15](https://huggingface.co/h94/IP-Adapter/resolve/main/models/ip-adapter_sd15.safetensors) | ViT-H | Basic model, average strength | | v1.5 | [ip-adapter_sd15_light](https://huggingface.co/h94/IP-Adapter/blob/main/models/ip-adapter_sd15_light.safetensors) | ViT-H | Light model, very light impact | @@ -72,12 +74,18 @@ The following table shows the combination of Checkpoint and Image encoder to use | SDXL | [ip-adapter-plus_sdxl_vit-h](https://huggingface.co/h94/IP-Adapter/resolve/main/sdxl_models/ip-adapter-plus_sdxl_vit-h.safetensors) | ViT-H | SDXL plus model, stronger | | SDXL | [ip-adapter-plus-face_sdxl_vit-h](https://huggingface.co/h94/IP-Adapter/resolve/main/sdxl_models/ip-adapter-plus-face_sdxl_vit-h.safetensors) | ViT-H | SDXL face model | -**FaceID** requires `insightface` and `onnxruntime`, you need to install them in your ComfyUI environment with `pip`, it's also a good idea to try to upgrade them with `pip install --upgrade ...`. +**FaceID** requires `insightface`, you need to install them in your ComfyUI environment. Check [this issue](https://github.com/cubiq/ComfyUI_IPAdapter_plus/issues/162) for help. When the dependencies are satisfied you need: -- The [main SD1.5 model](https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid_sd15.bin) to be placed into the ipadapter models directory. -- The [Lora](https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid_sd15_lora.safetensors) to be planced into `ComfyUI/models/loras/` directory. +| SD v. | IPadapter | Img encoder | Lora | +|---|---|---|---| +| v1.5 | [FaceID](https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid_sd15.bin) | (not used¹) | [FaceID Lora](https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid_sd15_lora.safetensors) | +| v1.5 | [FaceID Plus](https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid-plus_sd15.bin) | ViT-H | [FaceID Plus Lora](https://huggingface.co/h94/IP-Adapter-FaceID/resolve/main/ip-adapter-faceid-plus_sd15_lora.safetensors) | + +¹ The base FaceID model doesn't make use of a CLIP vision encoder. + +The loras need to be placed into `ComfyUI/models/loras/` directory. **There is no SDXL model at the moment.** diff --git a/examples/FaceID_Plus.json b/examples/FaceID_Plus.json new file mode 100644 index 0000000..12f9049 --- /dev/null +++ b/examples/FaceID_Plus.json @@ -0,0 +1,703 @@ +{ + "last_node_id": 32, + "last_link_id": 80, + "nodes": [ + { + "id": 6, + "type": "VAEDecode", + "pos": [ + 1670, + 620 + ], + "size": { + "0": 140, + "1": 50 + }, + "flags": {}, + "order": 12, + "mode": 0, + "inputs": [ + { + "name": "samples", + "type": "LATENT", + "link": 7 + }, + { + "name": "vae", + "type": "VAE", + "link": 8 + } + ], + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 72 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAEDecode" + } + }, + { + "id": 7, + "type": "VAELoader", + "pos": [ + 1380, + 920 + ], + "size": { + "0": 240, + "1": 60 + }, + "flags": {}, + "order": 0, + "mode": 0, + "outputs": [ + { + "name": "VAE", + "type": "VAE", + "links": [ + 8 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "VAELoader" + }, + "widgets_values": [ + "vae-ft-mse-840000-ema-pruned.safetensors" + ] + }, + { + "id": 2, + "type": "CheckpointLoaderSimple", + "pos": [ + 50, + 700 + ], + "size": { + "0": 290, + "1": 100 + }, + "flags": {}, + "order": 1, + "mode": 0, + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 15 + ], + "slot_index": 0 + }, + { + "name": "CLIP", + "type": "CLIP", + "links": [ + 5, + 6 + ], + "slot_index": 1 + }, + { + "name": "VAE", + "type": "VAE", + "links": [], + "slot_index": 2 + } + ], + "properties": { + "Node name for S&R": "CheckpointLoaderSimple" + }, + "widgets_values": [ + "sd15/realisticVisionV51_v51VAE.safetensors" + ] + }, + { + "id": 4, + "type": "CLIPTextEncode", + "pos": [ + 710, + 570 + ], + "size": { + "0": 400, + "1": 160 + }, + "flags": {}, + "order": 8, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 5 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 2 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "closeup photo of a woman wearing a white spring dress in a garden\n\nhigh quality, diffuse light, highly detailed, 4k" + ], + "color": "#232", + "bgcolor": "#353" + }, + { + "id": 5, + "type": "CLIPTextEncode", + "pos": [ + 710, + 777 + ], + "size": { + "0": 400, + "1": 160 + }, + "flags": {}, + "order": 9, + "mode": 0, + "inputs": [ + { + "name": "clip", + "type": "CLIP", + "link": 6 + } + ], + "outputs": [ + { + "name": "CONDITIONING", + "type": "CONDITIONING", + "links": [ + 3 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPTextEncode" + }, + "widgets_values": [ + "blurry, malformed, distorted, naked" + ], + "color": "#322", + "bgcolor": "#533" + }, + { + "id": 3, + "type": "EmptyLatentImage", + "pos": [ + 892, + 994 + ], + "size": { + "0": 210, + "1": 110 + }, + "flags": {}, + "order": 2, + "mode": 0, + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 4 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "EmptyLatentImage" + }, + "widgets_values": [ + 512, + 512, + 1 + ] + }, + { + "id": 10, + "type": "LoadImage", + "pos": [ + 180, + -50 + ], + "size": { + "0": 390.733154296875, + "1": 482.0174560546875 + }, + "flags": {}, + "order": 3, + "mode": 0, + "outputs": [ + { + "name": "IMAGE", + "type": "IMAGE", + "links": [ + 77 + ], + "shape": 3, + "slot_index": 0 + }, + { + "name": "MASK", + "type": "MASK", + "links": null, + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "LoadImage" + }, + "widgets_values": [ + "cwf_portrait.jpg", + "image" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 1, + "type": "KSampler", + "pos": [ + 1380, + 620 + ], + "size": { + "0": 240, + "1": 262 + }, + "flags": {}, + "order": 11, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 78 + }, + { + "name": "positive", + "type": "CONDITIONING", + "link": 2 + }, + { + "name": "negative", + "type": "CONDITIONING", + "link": 3 + }, + { + "name": "latent_image", + "type": "LATENT", + "link": 4 + } + ], + "outputs": [ + { + "name": "LATENT", + "type": "LATENT", + "links": [ + 7 + ], + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "KSampler" + }, + "widgets_values": [ + 2, + "fixed", + 25, + 6, + "dpmpp_2m", + "karras", + 1 + ] + }, + { + "id": 22, + "type": "PreviewImage", + "pos": [ + 1861, + 542 + ], + "size": { + "0": 527.009765625, + "1": 548.077880859375 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "images", + "type": "IMAGE", + "link": 72 + } + ], + "properties": { + "Node name for S&R": "PreviewImage" + } + }, + { + "id": 11, + "type": "IPAdapterApply", + "pos": [ + 1060, + 240 + ], + "size": { + "0": 210, + "1": 278 + }, + "flags": {}, + "order": 10, + "mode": 0, + "inputs": [ + { + "name": "ipadapter", + "type": "IPADAPTER", + "link": 10 + }, + { + "name": "clip_vision", + "type": "CLIP_VISION", + "link": 79, + "slot_index": 1 + }, + { + "name": "image", + "type": "IMAGE", + "link": 77 + }, + { + "name": "model", + "type": "MODEL", + "link": 16 + }, + { + "name": "attn_mask", + "type": "MASK", + "link": null + }, + { + "name": "insightface", + "type": "INSIGHTFACE", + "link": 80, + "slot_index": 5 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 78 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "IPAdapterApply" + }, + "widgets_values": [ + 0.9, + 0, + "original", + 0, + 1, + false + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 32, + "type": "InsightFaceLoader", + "pos": [ + 650, + 290 + ], + "size": { + "0": 290, + "1": 60 + }, + "flags": {}, + "order": 4, + "mode": 0, + "outputs": [ + { + "name": "INSIGHTFACE", + "type": "INSIGHTFACE", + "links": [ + 80 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "InsightFaceLoader" + }, + "widgets_values": [ + "CPU" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 14, + "type": "LoraLoaderModelOnly", + "pos": [ + 770, + 410 + ], + "size": { + "0": 210, + "1": 82 + }, + "flags": {}, + "order": 7, + "mode": 0, + "inputs": [ + { + "name": "model", + "type": "MODEL", + "link": 15 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 16 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "LoraLoaderModelOnly" + }, + "widgets_values": [ + "ip-adapter-faceid-plus_sd15_lora.safetensors", + 0.6 + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 31, + "type": "CLIPVisionLoader", + "pos": [ + 650, + 190 + ], + "size": { + "0": 290, + "1": 60 + }, + "flags": {}, + "order": 5, + "mode": 0, + "outputs": [ + { + "name": "CLIP_VISION", + "type": "CLIP_VISION", + "links": [ + 79 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "CLIPVisionLoader" + }, + "widgets_values": [ + "IPAdapter_image_encoder_sd15.safetensors" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 9, + "type": "IPAdapterModelLoader", + "pos": [ + 650, + 90 + ], + "size": { + "0": 290, + "1": 60 + }, + "flags": {}, + "order": 6, + "mode": 0, + "outputs": [ + { + "name": "IPADAPTER", + "type": "IPADAPTER", + "links": [ + 10 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "IPAdapterModelLoader" + }, + "widgets_values": [ + "ip-adapter-faceid-plus_sd15.bin" + ], + "color": "#223", + "bgcolor": "#335" + } + ], + "links": [ + [ + 2, + 4, + 0, + 1, + 1, + "CONDITIONING" + ], + [ + 3, + 5, + 0, + 1, + 2, + "CONDITIONING" + ], + [ + 4, + 3, + 0, + 1, + 3, + "LATENT" + ], + [ + 5, + 2, + 1, + 4, + 0, + "CLIP" + ], + [ + 6, + 2, + 1, + 5, + 0, + "CLIP" + ], + [ + 7, + 1, + 0, + 6, + 0, + "LATENT" + ], + [ + 8, + 7, + 0, + 6, + 1, + "VAE" + ], + [ + 10, + 9, + 0, + 11, + 0, + "IPADAPTER" + ], + [ + 15, + 2, + 0, + 14, + 0, + "MODEL" + ], + [ + 16, + 14, + 0, + 11, + 3, + "MODEL" + ], + [ + 72, + 6, + 0, + 22, + 0, + "IMAGE" + ], + [ + 77, + 10, + 0, + 11, + 2, + "IMAGE" + ], + [ + 78, + 11, + 0, + 1, + 0, + "MODEL" + ], + [ + 79, + 31, + 0, + 11, + 1, + "CLIP_VISION" + ], + [ + 80, + 32, + 0, + 11, + 5, + "INSIGHTFACE" + ] + ], + "groups": [], + "config": {}, + "extra": {}, + "version": 0.4 +} \ No newline at end of file diff --git a/examples/FaceID_basic.json b/examples/FaceID_basic.json index c6b06ad..6a420ab 100644 --- a/examples/FaceID_basic.json +++ b/examples/FaceID_basic.json @@ -1,6 +1,6 @@ { - "last_node_id": 30, - "last_link_id": 78, + "last_node_id": 32, + "last_link_id": 80, "nodes": [ { "id": 6, @@ -14,7 +14,7 @@ "1": 50 }, "flags": {}, - "order": 11, + "order": 12, "mode": 0, "inputs": [ { @@ -74,74 +74,6 @@ "vae-ft-mse-840000-ema-pruned.safetensors" ] }, - { - "id": 12, - "type": "InsightFaceLoader", - "pos": [ - 710, - 240 - ], - "size": { - "0": 210, - "1": 58 - }, - "flags": {}, - "order": 1, - "mode": 0, - "outputs": [ - { - "name": "CLIP_VISION", - "type": "CLIP_VISION", - "links": [ - 11 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "InsightFaceLoader" - }, - "widgets_values": [ - "CPU" - ], - "color": "#223", - "bgcolor": "#335" - }, - { - "id": 9, - "type": "IPAdapterModelLoader", - "pos": [ - 630, - 130 - ], - "size": { - "0": 300, - "1": 60 - }, - "flags": {}, - "order": 2, - "mode": 0, - "outputs": [ - { - "name": "IPADAPTER", - "type": "IPADAPTER", - "links": [ - 10 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "IPAdapterModelLoader" - }, - "widgets_values": [ - "ip-adapter-faceid_sd15.bin" - ], - "color": "#223", - "bgcolor": "#335" - }, { "id": 2, "type": "CheckpointLoaderSimple", @@ -154,7 +86,7 @@ "1": 100 }, "flags": {}, - "order": 3, + "order": 1, "mode": 0, "outputs": [ { @@ -200,7 +132,7 @@ "1": 160 }, "flags": {}, - "order": 7, + "order": 8, "mode": 0, "inputs": [ { @@ -240,7 +172,7 @@ "1": 160 }, "flags": {}, - "order": 8, + "order": 9, "mode": 0, "inputs": [ { @@ -280,7 +212,7 @@ "1": 110 }, "flags": {}, - "order": 4, + "order": 2, "mode": 0, "outputs": [ { @@ -308,12 +240,12 @@ 180, -50 ], - "size": [ - 390.73314006347687, - 482.01744979858427 - ], + "size": { + "0": 390.733154296875, + "1": 482.0174560546875 + }, "flags": {}, - "order": 5, + "order": 3, "mode": 0, "outputs": [ { @@ -349,12 +281,12 @@ 770, 410 ], - "size": [ - 210, - 82 - ], + "size": { + "0": 210, + "1": 82 + }, "flags": {}, - "order": 6, + "order": 7, "mode": 0, "inputs": [ { @@ -393,10 +325,10 @@ ], "size": { "0": 240, - "1": 260 + "1": 262 }, "flags": {}, - "order": 10, + "order": 11, "mode": 0, "inputs": [ { @@ -450,12 +382,12 @@ 1861, 542 ], - "size": [ - 527.0097535195305, - 548.0778934960933 - ], + "size": { + "0": 527.009765625, + "1": 548.077880859375 + }, "flags": {}, - "order": 12, + "order": 13, "mode": 0, "inputs": [ { @@ -477,10 +409,10 @@ ], "size": { "0": 210, - "1": 258 + "1": 278 }, "flags": {}, - "order": 9, + "order": 10, "mode": 0, "inputs": [ { @@ -491,7 +423,8 @@ { "name": "clip_vision", "type": "CLIP_VISION", - "link": 11 + "link": 79, + "slot_index": 1 }, { "name": "image", @@ -507,6 +440,12 @@ "name": "attn_mask", "type": "MASK", "link": null + }, + { + "name": "insightface", + "type": "INSIGHTFACE", + "link": 80, + "slot_index": 5 } ], "outputs": [ @@ -533,6 +472,106 @@ ], "color": "#223", "bgcolor": "#335" + }, + { + "id": 9, + "type": "IPAdapterModelLoader", + "pos": [ + 650, + 90 + ], + "size": [ + 290, + 60 + ], + "flags": {}, + "order": 4, + "mode": 0, + "outputs": [ + { + "name": "IPADAPTER", + "type": "IPADAPTER", + "links": [ + 10 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "IPAdapterModelLoader" + }, + "widgets_values": [ + "ip-adapter-faceid_sd15.bin" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 31, + "type": "CLIPVisionLoader", + "pos": [ + 650, + 190 + ], + "size": [ + 290, + 60 + ], + "flags": {}, + "order": 5, + "mode": 0, + "outputs": [ + { + "name": "CLIP_VISION", + "type": "CLIP_VISION", + "links": [ + 79 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "CLIPVisionLoader" + }, + "widgets_values": [ + "IPAdapter_image_encoder_sd15.safetensors" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 32, + "type": "InsightFaceLoader", + "pos": [ + 650, + 290 + ], + "size": [ + 290, + 60 + ], + "flags": {}, + "order": 6, + "mode": 0, + "outputs": [ + { + "name": "INSIGHTFACE", + "type": "INSIGHTFACE", + "links": [ + 80 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "InsightFaceLoader" + }, + "widgets_values": [ + "CPU" + ], + "color": "#223", + "bgcolor": "#335" } ], "links": [ @@ -600,14 +639,6 @@ 0, "IPADAPTER" ], - [ - 11, - 12, - 0, - 11, - 1, - "CLIP_VISION" - ], [ 15, 2, @@ -647,6 +678,22 @@ 1, 0, "MODEL" + ], + [ + 79, + 31, + 0, + 11, + 1, + "CLIP_VISION" + ], + [ + 80, + 32, + 0, + 11, + 5, + "INSIGHTFACE" ] ], "groups": [], diff --git a/examples/FaceID_enhanced.json b/examples/FaceID_enhanced.json index 2f602a7..09599f5 100644 --- a/examples/FaceID_enhanced.json +++ b/examples/FaceID_enhanced.json @@ -1,6 +1,6 @@ { - "last_node_id": 30, - "last_link_id": 77, + "last_node_id": 32, + "last_link_id": 81, "nodes": [ { "id": 6, @@ -74,74 +74,6 @@ "vae-ft-mse-840000-ema-pruned.safetensors" ] }, - { - "id": 12, - "type": "InsightFaceLoader", - "pos": [ - 380, - 240 - ], - "size": { - "0": 210, - "1": 58 - }, - "flags": {}, - "order": 1, - "mode": 0, - "outputs": [ - { - "name": "CLIP_VISION", - "type": "CLIP_VISION", - "links": [ - 11 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "InsightFaceLoader" - }, - "widgets_values": [ - "CPU" - ], - "color": "#223", - "bgcolor": "#335" - }, - { - "id": 9, - "type": "IPAdapterModelLoader", - "pos": [ - 300, - 130 - ], - "size": { - "0": 300, - "1": 60 - }, - "flags": {}, - "order": 2, - "mode": 0, - "outputs": [ - { - "name": "IPADAPTER", - "type": "IPADAPTER", - "links": [ - 10 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "IPAdapterModelLoader" - }, - "widgets_values": [ - "ip-adapter-faceid_sd15.bin" - ], - "color": "#223", - "bgcolor": "#335" - }, { "id": 19, "type": "PrepImageForClipVision", @@ -149,10 +81,10 @@ 700, 80 ], - "size": [ - 243.60000610351562, - 106 - ], + "size": { + "0": 243.60000610351562, + "1": 106 + }, "flags": {}, "order": 11, "mode": 0, @@ -185,74 +117,6 @@ "color": "#223", "bgcolor": "#335" }, - { - "id": 17, - "type": "CLIPVisionLoader", - "pos": [ - 730, - -40 - ], - "size": [ - 210, - 58 - ], - "flags": {}, - "order": 3, - "mode": 0, - "outputs": [ - { - "name": "CLIP_VISION", - "type": "CLIP_VISION", - "links": [ - 70 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "CLIPVisionLoader" - }, - "widgets_values": [ - "IPAdapter_image_encoder_sd15.safetensors" - ], - "color": "#223", - "bgcolor": "#335" - }, - { - "id": 18, - "type": "IPAdapterModelLoader", - "pos": [ - 640, - -160 - ], - "size": { - "0": 300, - "1": 60 - }, - "flags": {}, - "order": 4, - "mode": 0, - "outputs": [ - { - "name": "IPADAPTER", - "type": "IPADAPTER", - "links": [ - 28 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "IPAdapterModelLoader" - }, - "widgets_values": [ - "ip-adapter-plus-face_sd15.safetensors" - ], - "color": "#223", - "bgcolor": "#335" - }, { "id": 22, "type": "PreviewImage", @@ -290,7 +154,7 @@ "1": 100 }, "flags": {}, - "order": 5, + "order": 1, "mode": 0, "outputs": [ { @@ -333,7 +197,7 @@ ], "size": { "0": 240, - "1": 260 + "1": 262 }, "flags": {}, "order": 14, @@ -392,7 +256,7 @@ ], "size": { "0": 210, - "1": 258 + "1": 278 }, "flags": {}, "order": 12, @@ -406,7 +270,7 @@ { "name": "clip_vision", "type": "CLIP_VISION", - "link": 11 + "link": 81 }, { "name": "image", @@ -422,6 +286,11 @@ "name": "attn_mask", "type": "MASK", "link": null + }, + { + "name": "insightface", + "type": "INSIGHTFACE", + "link": 80 } ], "outputs": [ @@ -541,7 +410,7 @@ "1": 110 }, "flags": {}, - "order": 6, + "order": 2, "mode": 0, "outputs": [ { @@ -562,73 +431,6 @@ 1 ] }, - { - "id": 16, - "type": "IPAdapterApply", - "pos": [ - 1070, - 180 - ], - "size": { - "0": 210, - "1": 258 - }, - "flags": {}, - "order": 13, - "mode": 0, - "inputs": [ - { - "name": "ipadapter", - "type": "IPADAPTER", - "link": 28 - }, - { - "name": "clip_vision", - "type": "CLIP_VISION", - "link": 70, - "slot_index": 1 - }, - { - "name": "image", - "type": "IMAGE", - "link": 30 - }, - { - "name": "model", - "type": "MODEL", - "link": 35 - }, - { - "name": "attn_mask", - "type": "MASK", - "link": null - } - ], - "outputs": [ - { - "name": "MODEL", - "type": "MODEL", - "links": [ - 36 - ], - "shape": 3, - "slot_index": 0 - } - ], - "properties": { - "Node name for S&R": "IPAdapterApply" - }, - "widgets_values": [ - 0.35000000000000003, - 0.33, - "original", - 0, - 1, - false - ], - "color": "#223", - "bgcolor": "#335" - }, { "id": 10, "type": "LoadImage", @@ -636,12 +438,12 @@ -150, -50 ], - "size": [ - 390.73314006347687, - 482.01744979858427 - ], + "size": { + "0": 390.733154296875, + "1": 482.0174560546875 + }, "flags": {}, - "order": 7, + "order": 3, "mode": 0, "outputs": [ { @@ -678,10 +480,10 @@ 440, 410 ], - "size": [ - 210, - 82 - ], + "size": { + "0": 210, + "1": 82 + }, "flags": {}, "order": 8, "mode": 0, @@ -712,6 +514,215 @@ ], "color": "#223", "bgcolor": "#335" + }, + { + "id": 16, + "type": "IPAdapterApply", + "pos": [ + 1070, + 180 + ], + "size": { + "0": 210, + "1": 278 + }, + "flags": {}, + "order": 13, + "mode": 0, + "inputs": [ + { + "name": "ipadapter", + "type": "IPADAPTER", + "link": 28 + }, + { + "name": "clip_vision", + "type": "CLIP_VISION", + "link": 78, + "slot_index": 1 + }, + { + "name": "image", + "type": "IMAGE", + "link": 30 + }, + { + "name": "model", + "type": "MODEL", + "link": 35 + }, + { + "name": "attn_mask", + "type": "MASK", + "link": null + }, + { + "name": "insightface", + "type": "INSIGHTFACE", + "link": null, + "slot_index": 5 + } + ], + "outputs": [ + { + "name": "MODEL", + "type": "MODEL", + "links": [ + 36 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "IPAdapterApply" + }, + "widgets_values": [ + 0.35000000000000003, + 0.33, + "original", + 0, + 1, + false + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 32, + "type": "InsightFaceLoader", + "pos": [ + 310, + 260 + ], + "size": [ + 290, + 60 + ], + "flags": {}, + "order": 4, + "mode": 0, + "outputs": [ + { + "name": "INSIGHTFACE", + "type": "INSIGHTFACE", + "links": [ + 80 + ], + "shape": 3 + } + ], + "properties": { + "Node name for S&R": "InsightFaceLoader" + }, + "widgets_values": [ + "CPU" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 9, + "type": "IPAdapterModelLoader", + "pos": [ + 310, + 160 + ], + "size": [ + 290, + 60 + ], + "flags": {}, + "order": 5, + "mode": 0, + "outputs": [ + { + "name": "IPADAPTER", + "type": "IPADAPTER", + "links": [ + 10 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "IPAdapterModelLoader" + }, + "widgets_values": [ + "ip-adapter-faceid_sd15.bin" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 31, + "type": "CLIPVisionLoader", + "pos": [ + 310, + 60 + ], + "size": [ + 290, + 60 + ], + "flags": {}, + "order": 6, + "mode": 0, + "outputs": [ + { + "name": "CLIP_VISION", + "type": "CLIP_VISION", + "links": [ + 78, + 81 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "CLIPVisionLoader" + }, + "widgets_values": [ + "IPAdapter_image_encoder_sd15.safetensors" + ], + "color": "#223", + "bgcolor": "#335" + }, + { + "id": 18, + "type": "IPAdapterModelLoader", + "pos": [ + 650, + -30 + ], + "size": { + "0": 290, + "1": 60 + }, + "flags": {}, + "order": 7, + "mode": 0, + "outputs": [ + { + "name": "IPADAPTER", + "type": "IPADAPTER", + "links": [ + 28 + ], + "shape": 3, + "slot_index": 0 + } + ], + "properties": { + "Node name for S&R": "IPAdapterModelLoader" + }, + "widgets_values": [ + "ip-adapter-plus-face_sd15.safetensors" + ], + "color": "#223", + "bgcolor": "#335" } ], "links": [ @@ -779,14 +790,6 @@ 0, "IPADAPTER" ], - [ - 11, - 12, - 0, - 11, - 1, - "CLIP_VISION" - ], [ 15, 2, @@ -843,14 +846,6 @@ 0, "MODEL" ], - [ - 70, - 17, - 0, - 16, - 1, - "CLIP_VISION" - ], [ 72, 6, @@ -866,6 +861,30 @@ 11, 2, "IMAGE" + ], + [ + 78, + 31, + 0, + 16, + 1, + "CLIP_VISION" + ], + [ + 80, + 32, + 0, + 11, + 5, + "INSIGHTFACE" + ], + [ + 81, + 31, + 0, + 11, + 1, + "CLIP_VISION" ] ], "groups": [], diff --git a/examples/face_id_wf.jpg b/examples/face_id_wf.jpg index c82922f..076a040 100644 Binary files a/examples/face_id_wf.jpg and b/examples/face_id_wf.jpg differ diff --git a/resampler.py b/resampler.py index 4521c8c..2426667 100644 --- a/resampler.py +++ b/resampler.py @@ -1,8 +1,12 @@ # modified from https://github.com/mlfoundations/open_flamingo/blob/main/open_flamingo/src/helpers.py +# and https://github.com/lucidrains/imagen-pytorch/blob/main/imagen_pytorch/imagen_pytorch.py + import math import torch import torch.nn as nn +from einops import rearrange +from einops.layers.torch import Rearrange # FFN @@ -14,11 +18,11 @@ def FeedForward(dim, mult=4): 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) + # (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) @@ -42,7 +46,6 @@ class PerceiverAttention(nn.Module): 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: @@ -53,23 +56,23 @@ class PerceiverAttention(nn.Module): """ 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 = (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) @@ -86,16 +89,30 @@ class Resampler(nn.Module): 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( @@ -108,14 +125,34 @@ class Resampler(nn.Module): ) 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) \ No newline at end of file + 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)