import torch import ldm_patched.modules.clip_vision import safetensors.torch as sf import ldm_patched.modules.model_management as model_management import ldm_patched.ldm.modules.attention as attention from extras.resampler import Resampler from ldm_patched.modules.model_patcher import ModelPatcher from modules.core import numpy_to_pytorch from modules.ops import use_patched_ops from ldm_patched.modules.ops import manual_cast SD_V12_CHANNELS = [320] * 4 + [640] * 4 + [1280] * 4 + [1280] * 6 + [640] * 6 + [320] * 6 + [1280] * 2 SD_XL_CHANNELS = [640] * 8 + [1280] * 40 + [1280] * 60 + [640] * 12 + [1280] * 20 def sdp(q, k, v, extra_options): return attention.optimized_attention(q, k, v, heads=extra_options["n_heads"], mask=None) class ImageProjModel(torch.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 = torch.nn.Linear(clip_embeddings_dim, self.clip_extra_context_tokens * cross_attention_dim) self.norm = torch.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(torch.nn.Module): def __init__(self, cross_attention_dim): super().__init__() channels = SD_XL_CHANNELS if cross_attention_dim == 2048 else SD_V12_CHANNELS self.to_kvs = torch.nn.ModuleList( [torch.nn.Linear(cross_attention_dim, channel, bias=False) for channel in channels]) def load_state_dict_ordered(self, sd): state_dict = [] for i in range(4096): for k in ['k', 'v']: key = f'{i}.to_{k}_ip.weight' if key in sd: state_dict.append(sd[key]) for i, v in enumerate(state_dict): self.to_kvs[i].weight = torch.nn.Parameter(v, requires_grad=False) class IPAdapterModel(torch.nn.Module): def __init__(self, state_dict, plus, cross_attention_dim=768, clip_embeddings_dim=1024, clip_extra_context_tokens=4, sdxl_plus=False): super().__init__() self.plus = plus if self.plus: self.image_proj_model = Resampler( dim=1280 if sdxl_plus else cross_attention_dim, depth=4, dim_head=64, heads=20 if sdxl_plus else 12, num_queries=clip_extra_context_tokens, embedding_dim=clip_embeddings_dim, output_dim=cross_attention_dim, ff_mult=4 ) else: self.image_proj_model = ImageProjModel( cross_attention_dim=cross_attention_dim, clip_embeddings_dim=clip_embeddings_dim, clip_extra_context_tokens=clip_extra_context_tokens ) self.image_proj_model.load_state_dict(state_dict["image_proj"]) self.ip_layers = To_KV(cross_attention_dim) self.ip_layers.load_state_dict_ordered(state_dict["ip_adapter"]) clip_vision: ldm_patched.modules.clip_vision.ClipVisionModel = None ip_negative: torch.Tensor = None ip_adapters: dict = {} def load_ip_adapter(clip_vision_path, ip_negative_path, ip_adapter_path): global clip_vision, ip_negative, ip_adapters if clip_vision is None and isinstance(clip_vision_path, str): clip_vision = ldm_patched.modules.clip_vision.load(clip_vision_path) if ip_negative is None and isinstance(ip_negative_path, str): ip_negative = sf.load_file(ip_negative_path)['data'] if not isinstance(ip_adapter_path, str) or ip_adapter_path in ip_adapters: return load_device = model_management.get_torch_device() offload_device = torch.device('cpu') use_fp16 = model_management.should_use_fp16(device=load_device) ip_state_dict = torch.load(ip_adapter_path, map_location="cpu") plus = "latents" in ip_state_dict["image_proj"] cross_attention_dim = ip_state_dict["ip_adapter"]["1.to_k_ip.weight"].shape[1] sdxl = cross_attention_dim == 2048 sdxl_plus = sdxl and plus if plus: clip_extra_context_tokens = ip_state_dict["image_proj"]["latents"].shape[1] clip_embeddings_dim = ip_state_dict["image_proj"]["latents"].shape[2] else: clip_extra_context_tokens = ip_state_dict["image_proj"]["proj.weight"].shape[0] // cross_attention_dim clip_embeddings_dim = None with use_patched_ops(manual_cast): ip_adapter = IPAdapterModel( ip_state_dict, plus=plus, cross_attention_dim=cross_attention_dim, clip_embeddings_dim=clip_embeddings_dim, clip_extra_context_tokens=clip_extra_context_tokens, sdxl_plus=sdxl_plus ) ip_adapter.sdxl = sdxl ip_adapter.load_device = load_device ip_adapter.offload_device = offload_device ip_adapter.dtype = torch.float16 if use_fp16 else torch.float32 ip_adapter.to(offload_device, dtype=ip_adapter.dtype) image_proj_model = ModelPatcher(model=ip_adapter.image_proj_model, load_device=load_device, offload_device=offload_device) ip_layers = ModelPatcher(model=ip_adapter.ip_layers, load_device=load_device, offload_device=offload_device) ip_adapters[ip_adapter_path] = dict( ip_adapter=ip_adapter, image_proj_model=image_proj_model, ip_layers=ip_layers, ip_unconds=None ) return @torch.no_grad() @torch.inference_mode() def clip_preprocess(image): mean = torch.tensor([0.48145466, 0.4578275, 0.40821073], device=image.device, dtype=image.dtype).view([1, 3, 1, 1]) std = torch.tensor([0.26862954, 0.26130258, 0.27577711], device=image.device, dtype=image.dtype).view([1, 3, 1, 1]) image = image.movedim(-1, 1) # https://github.com/tencent-ailab/IP-Adapter/blob/d580c50a291566bbf9fc7ac0f760506607297e6d/README.md?plain=1#L75 B, C, H, W = image.shape assert H == 224 and W == 224 return (image - mean) / std @torch.no_grad() @torch.inference_mode() def preprocess(img, ip_adapter_path): global ip_adapters entry = ip_adapters[ip_adapter_path] ldm_patched.modules.model_management.load_model_gpu(clip_vision.patcher) pixel_values = clip_preprocess(numpy_to_pytorch(img).to(clip_vision.load_device)) outputs = clip_vision.model(pixel_values=pixel_values, output_hidden_states=True) ip_adapter = entry['ip_adapter'] ip_layers = entry['ip_layers'] image_proj_model = entry['image_proj_model'] ip_unconds = entry['ip_unconds'] if ip_adapter.plus: cond = outputs.hidden_states[-2] else: cond = outputs.image_embeds cond = cond.to(device=ip_adapter.load_device, dtype=ip_adapter.dtype) ldm_patched.modules.model_management.load_model_gpu(image_proj_model) cond = image_proj_model.model(cond).to(device=ip_adapter.load_device, dtype=ip_adapter.dtype) ldm_patched.modules.model_management.load_model_gpu(ip_layers) if ip_unconds is None: uncond = ip_negative.to(device=ip_adapter.load_device, dtype=ip_adapter.dtype) ip_unconds = [m(uncond).cpu() for m in ip_layers.model.to_kvs] entry['ip_unconds'] = ip_unconds ip_conds = [m(cond).cpu() for m in ip_layers.model.to_kvs] return ip_conds, ip_unconds @torch.no_grad() @torch.inference_mode() def patch_model(model, tasks): new_model = model.clone() def make_attn_patcher(ip_index): def patcher(n, context_attn2, value_attn2, extra_options): org_dtype = n.dtype current_step = float(model.model.diffusion_model.current_step.detach().cpu().numpy()[0]) cond_or_uncond = extra_options['cond_or_uncond'] q = n k = [context_attn2] v = [value_attn2] b, _, _ = q.shape for (cs, ucs), cn_stop, cn_weight in tasks: if current_step < cn_stop: ip_k_c = cs[ip_index * 2].to(q) ip_v_c = cs[ip_index * 2 + 1].to(q) ip_k_uc = ucs[ip_index * 2].to(q) ip_v_uc = ucs[ip_index * 2 + 1].to(q) ip_k = torch.cat([(ip_k_c, ip_k_uc)[i] for i in cond_or_uncond], dim=0) ip_v = torch.cat([(ip_v_c, ip_v_uc)[i] for i in cond_or_uncond], dim=0) # Midjourney's attention formulation of image prompt (non-official reimplementation) # Written by Lvmin Zhang at Stanford University, 2023 Dec # For non-commercial use only - if you use this in commercial project then # probably it has some intellectual property issues. # Contact lvminzhang@acm.org if you are not sure. # Below is the sensitive part with potential intellectual property issues. ip_v_mean = torch.mean(ip_v, dim=1, keepdim=True) ip_v_offset = ip_v - ip_v_mean B, F, C = ip_k.shape channel_penalty = float(C) / 1280.0 weight = cn_weight * channel_penalty ip_k = ip_k * weight ip_v = ip_v_offset + ip_v_mean * weight k.append(ip_k) v.append(ip_v) k = torch.cat(k, dim=1) v = torch.cat(v, dim=1) out = sdp(q, k, v, extra_options) return out.to(dtype=org_dtype) return patcher def set_model_patch_replace(model, number, 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"]: to["patches_replace"]["attn2"][key] = make_attn_patcher(number) number = 0 for id in [4, 5, 7, 8]: block_indices = range(2) if id in [4, 5] else range(10) for index in block_indices: set_model_patch_replace(new_model, number, ("input", id, index)) number += 1 for id in range(6): block_indices = range(2) if id in [3, 4, 5] else range(10) for index in block_indices: set_model_patch_replace(new_model, number, ("output", id, index)) number += 1 for index in range(10): set_model_patch_replace(new_model, number, ("middle", 0, index)) number += 1 return new_model