Spaces:
Runtime error
Runtime error
File size: 7,977 Bytes
96d7ad8 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 |
import copy
from typing import Any, Callable, Dict, Iterable, Union
import PIL
import cv2
import torch
import argparse
import datetime
import logging
import inspect
import math
import os
import shutil
from typing import Dict, List, Optional, Tuple
from pprint import pprint
from collections import OrderedDict
from dataclasses import dataclass
import gc
import time
import numpy as np
from omegaconf import OmegaConf
from omegaconf import SCMode
import torch
from torch import nn
import torch.nn.functional as F
import torch.utils.checkpoint
from einops import rearrange, repeat
import pandas as pd
import h5py
from diffusers.models.modeling_utils import load_state_dict
from diffusers.utils import (
logging,
)
from diffusers.utils.import_utils import is_xformers_available
from ip_adapter.resampler import Resampler
from ip_adapter.ip_adapter import ImageProjModel
from ip_adapter.ip_adapter_faceid import ProjPlusModel, MLPProjModel
from mmcm.vision.feature_extractor.clip_vision_extractor import (
ImageClipVisionFeatureExtractor,
ImageClipVisionFeatureExtractorV2,
)
from mmcm.vision.feature_extractor.insight_face_extractor import (
InsightFaceExtractorNormEmb,
)
from .unet_loader import update_unet_with_sd
from .unet_3d_condition import UNet3DConditionModel
from .ip_adapter_loader import ip_adapter_keys_list
logger = logging.get_logger(__name__) # pylint: disable=invalid-name
# refer https://github.com/tencent-ailab/IP-Adapter/issues/168#issuecomment-1846771651
unet_keys_list = [
"down_blocks.0.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"down_blocks.0.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"down_blocks.0.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"down_blocks.0.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"down_blocks.1.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"down_blocks.1.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"down_blocks.1.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"down_blocks.1.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"down_blocks.2.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"down_blocks.2.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"down_blocks.2.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"down_blocks.2.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"up_blocks.1.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"up_blocks.1.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"up_blocks.1.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"up_blocks.1.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"up_blocks.1.attentions.2.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"up_blocks.1.attentions.2.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"up_blocks.2.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"up_blocks.2.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"up_blocks.2.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"up_blocks.2.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"up_blocks.2.attentions.2.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"up_blocks.2.attentions.2.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"up_blocks.3.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"up_blocks.3.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"up_blocks.3.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"up_blocks.3.attentions.1.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"up_blocks.3.attentions.2.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"up_blocks.3.attentions.2.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
"mid_block.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_k_ip.weight",
"mid_block.attentions.0.transformer_blocks.0.attn2.processor.ip_adapter_face_to_v_ip.weight",
]
UNET2IPAadapter_Keys_MAPIING = {
k: v for k, v in zip(unet_keys_list, ip_adapter_keys_list)
}
def load_ip_adapter_face_extractor_and_proj_by_name(
model_name: str,
ip_ckpt: Tuple[str, nn.Module],
ip_image_encoder: Tuple[str, nn.Module] = None,
cross_attention_dim: int = 768,
clip_embeddings_dim: int = 1024,
clip_extra_context_tokens: int = 4,
ip_scale: float = 0.0,
dtype: torch.dtype = torch.float16,
device: str = "cuda",
unet: nn.Module = None,
) -> nn.Module:
if model_name == "IPAdapterFaceID":
if ip_image_encoder is not None:
ip_adapter_face_emb_extractor = InsightFaceExtractorNormEmb(
pretrained_model_name_or_path=ip_image_encoder,
dtype=dtype,
device=device,
)
else:
ip_adapter_face_emb_extractor = None
ip_adapter_image_proj = MLPProjModel(
cross_attention_dim=cross_attention_dim,
id_embeddings_dim=clip_embeddings_dim,
num_tokens=clip_extra_context_tokens,
).to(device, dtype=dtype)
else:
raise ValueError(
f"unsupport model_name={model_name}, only support IPAdapter, IPAdapterPlus, IPAdapterFaceID"
)
ip_adapter_state_dict = torch.load(
ip_ckpt,
map_location="cpu",
)
ip_adapter_image_proj.load_state_dict(ip_adapter_state_dict["image_proj"])
if unet is not None and "ip_adapter" in ip_adapter_state_dict:
update_unet_ip_adapter_cross_attn_param(
unet,
ip_adapter_state_dict["ip_adapter"],
)
logger.info(
f"update unet.spatial_cross_attn_ip_adapter parameter with {ip_ckpt}"
)
return (
ip_adapter_face_emb_extractor,
ip_adapter_image_proj,
)
def update_unet_ip_adapter_cross_attn_param(
unet: UNet3DConditionModel, ip_adapter_state_dict: Dict
) -> None:
"""use independent ip_adapter attn 中的 to_k, to_v in unet
ip_adapter： like ['1.to_k_ip.weight', '1.to_v_ip.weight', '3.to_k_ip.weight']
Args:
unet (UNet3DConditionModel): _description_
ip_adapter_state_dict (Dict): _description_
"""
unet_spatial_cross_atnns = unet.spatial_cross_attns[0]
unet_spatial_cross_atnns_dct = {k: v for k, v in unet_spatial_cross_atnns}
for i, (unet_key_more, ip_adapter_key) in enumerate(
UNET2IPAadapter_Keys_MAPIING.items()
):
ip_adapter_value = ip_adapter_state_dict[ip_adapter_key]
unet_key_more_spit = unet_key_more.split(".")
unet_key = ".".join(unet_key_more_spit[:-3])
suffix = ".".join(unet_key_more_spit[-3:])
logger.debug(
f"{i}: unet_key_more = {unet_key_more}, {unet_key}=unet_key, suffix={suffix}",
)
if ".ip_adapter_face_to_k" in suffix:
with torch.no_grad():
unet_spatial_cross_atnns_dct[
unet_key
].ip_adapter_face_to_k_ip.weight.copy_(ip_adapter_value.data)
else:
with torch.no_grad():
unet_spatial_cross_atnns_dct[
unet_key
].ip_adapter_face_to_v_ip.weight.copy_(ip_adapter_value.data)
|