i (#559)

fooocus_extras/ip_adapter.py

@@ -0,0 +1,304 @@
+import torch
+import comfy.clip_vision
+import safetensors.torch as sf
+import comfy.model_management as model_management
+import contextlib
+
+from fooocus_extras.resampler import Resampler
+from comfy.model_patcher import ModelPatcher
+
+
+if model_management.xformers_enabled():
+    import xformers
+    import xformers.ops
+
+
+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):
+    if model_management.xformers_enabled():
+        b, _, _ = q.shape
+        q, k, v = map(
+            lambda t: t.unsqueeze(3)
+            .reshape(b, t.shape[1], extra_options["n_heads"], extra_options["dim_head"])
+            .permute(0, 2, 1, 3)
+            .reshape(b * extra_options["n_heads"], t.shape[1], extra_options["dim_head"])
+            .contiguous(),
+            (q, k, v),
+        )
+        out = xformers.ops.memory_efficient_attention(q, k, v, attn_bias=None, op=None)
+        out = (
+            out.unsqueeze(0)
+            .reshape(b, extra_options["n_heads"], out.shape[1], extra_options["dim_head"])
+            .permute(0, 2, 1, 3)
+            .reshape(b, out.shape[1], extra_options["n_heads"] * extra_options["dim_head"])
+        )
+    else:
+        b, _, _ = q.shape
+        q, k, v = map(
+            lambda t: t.view(b, -1, extra_options["n_heads"], extra_options["dim_head"]).transpose(1, 2),
+            (q, k, v),
+        )
+        out = torch.nn.functional.scaled_dot_product_attention(q, k, v, attn_mask=None, dropout_p=0.0, is_causal=False)
+        out = out.transpose(1, 2).reshape(b, -1, extra_options["n_heads"] * extra_options["dim_head"])
+    return out
+
+
+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: comfy.clip_vision.ClipVisionModel = None
+ip_negative: torch.Tensor = None
+image_proj_model: ModelPatcher = None
+ip_layers: ModelPatcher = None
+ip_adapter: IPAdapterModel = None
+
+
+def load_ip_adapter(clip_vision_path, ip_negative_path, ip_adapter_path):
+    global clip_vision, image_proj_model, ip_layers, ip_negative, ip_adapter
+
+    if clip_vision_path is None:
+        return
+    if ip_negative_path is None:
+        return
+    if ip_adapter_path is None:
+        return
+    if clip_vision is not None and image_proj_model is not None and ip_layers is not None and ip_negative is not None:
+        return
+
+    ip_negative = sf.load_file(ip_negative_path)['data']
+    clip_vision = comfy.clip_vision.load(clip_vision_path)
+
+    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
+
+    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)
+
+    return
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def preprocess(img):
+    inputs = clip_vision.processor(images=img, return_tensors="pt")
+    comfy.model_management.load_models_gpu([clip_vision.patcher, image_proj_model])
+    pixel_values = inputs['pixel_values'].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=pixel_values, output_hidden_states=True)
+
+    if ip_adapter.plus:
+        cond = outputs.hidden_states[-2].to(ip_adapter.dtype)
+    else:
+        cond = outputs.image_embeds.to(ip_adapter.dtype)
+
+    outputs = image_proj_model.model(cond)
+
+    return outputs
+
+
+@torch.no_grad()
+@torch.inference_mode()
+def patch_model(model, ip_tasks):
+    new_model = model.clone()
+
+    tasks = []
+    for cn_img, cn_stop, cn_weight in ip_tasks:
+        tasks.append((cn_img, cn_stop, cn_weight, {}))
+
+    def make_attn_patcher(ip_index):
+        ip_model_k = ip_layers.model.to_kvs[ip_index * 2]
+        ip_model_v = ip_layers.model.to_kvs[ip_index * 2 + 1]
+
+        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']
+
+            with torch.autocast("cuda", dtype=ip_adapter.dtype):
+                q = n
+                k = [context_attn2]
+                v = [value_attn2]
+                b, _, _ = q.shape
+                batch_prompt = b // len(cond_or_uncond)
+
+                for cn_img, cn_stop, cn_weight, cache in tasks:
+                    if current_step < cn_stop:
+                        if ip_index in cache:
+                            ip_k, ip_v = cache[ip_index]
+                        else:
+                            ip_model_k.to(device=ip_adapter.load_device, dtype=ip_adapter.dtype)
+                            ip_model_v.to(device=ip_adapter.load_device, dtype=ip_adapter.dtype)
+                            cond = cn_img.to(device=ip_adapter.load_device, dtype=ip_adapter.dtype).repeat(batch_prompt, 1, 1)
+                            uncond = ip_negative.to(device=ip_adapter.load_device, dtype=ip_adapter.dtype).repeat(batch_prompt, 1, 1)
+                            uncond_cond = torch.cat([(cond, uncond)[i] for i in cond_or_uncond], dim=0)
+                            ip_k = ip_model_k(uncond_cond)
+                            ip_v = ip_model_v(uncond_cond)
+
+                            # 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
+
+                            # The sensitive part ends here.
+
+                            cache[ip_index] = ip_k, ip_v
+                            ip_model_k.to(device=ip_adapter.offload_device, dtype=ip_adapter.dtype)
+                            ip_model_v.to(device=ip_adapter.offload_device, dtype=ip_adapter.dtype)
+
+                        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
+    if not ip_adapter.sdxl:
+        for id in [1, 2, 4, 5, 7, 8]:  # id of input_blocks that have cross attention
+            set_model_patch_replace(new_model, number, ("input", id))
+            number += 1
+        for id in [3, 4, 5, 6, 7, 8, 9, 10, 11]:  # id of output_blocks that have cross attention
+            set_model_patch_replace(new_model, number, ("output", id))
+            number += 1
+        set_model_patch_replace(new_model, number, ("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(new_model, number, ("input", id, index))
+                number += 1
+        for id in range(6):  # id of output_blocks that have cross attention
+            block_indices = range(2) if id in [3, 4, 5] else range(10)  # transformer_depth
+            for index in block_indices:
+                set_model_patch_replace(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

fooocus_extras/preprocessors.py

@@ -0,0 +1,42 @@
+import cv2
+import numpy as np
+
+
+def canny_k(x, k=0.5):
+    import cv2
+    H, W, C = x.shape
+    Hs, Ws = int(H * k), int(W * k)
+    small = cv2.resize(x, (Ws, Hs), interpolation=cv2.INTER_AREA)
+    return cv2.Canny(small, 100, 200).astype(np.float32) / 255.0
+
+
+def canny_pyramid(x):
+    # For some reasons, SAI's Control-lora Canny seems to be trained on canny maps with non-standard resolutions.
+    # Then we use pyramid to use all resolutions to avoid missing any structure in specific resolutions.
+
+    ks = [0.2, 0.3, 0.4, 0.5, 0.6, 0.7, 0.8, 0.9, 1.0]
+    cs = [canny_k(x, k) for k in ks]
+    cur = None
+
+    for c in cs:
+        if cur is None:
+            cur = c
+        else:
+            H, W = c.shape
+            cur = cv2.resize(cur, (W, H), interpolation=cv2.INTER_LINEAR)
+            cur = cur * 0.75 + c * 0.25
+
+    cur *= 400.0
+
+    return cur.clip(0, 255).astype(np.uint8)
+
+
+def cpds(x):
+    import cv2
+    # cv2.decolor is not "decolor", it is Cewu Lu's method
+    # See http://www.cse.cuhk.edu.hk/leojia/projects/color2gray/index.html
+    # See https://docs.opencv.org/3.0-beta/modules/photo/doc/decolor.html
+
+    y = np.ascontiguousarray(x[:, :, ::-1].copy())
+    y = cv2.decolor(y)[0]
+    return y

fooocus_extras/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)