hcpdiff 0.9.1__py3-none-any.whl → 2.1__py3-none-any.whl

hcpdiff/models/cfg_context.py

@@ -1,6 +1,7 @@
 import torch
 from einops import repeat
 import math
+from typing import Union, Callable
 
 class CFGContext:
     def pre(self, noisy_latents, timesteps):
@@ -10,9 +11,11 @@ class CFGContext:
         return model_pred
 
 class DreamArtistPTContext(CFGContext):
-    def __init__(self, cfg_scale, num_train_timesteps):
-        self.cfg_scale=cfg_scale
-        self.num_train_timesteps=num_train_timesteps
+    def __init__(self, cfg_low: float, cfg_high: float=None, cfg_func: Union[str, Callable]=None, num_train_timesteps=1000):
+        self.cfg_low = cfg_low
+        self.cfg_high = cfg_high or cfg_low
+        self.cfg_func = cfg_func
+        self.num_train_timesteps = num_train_timesteps
 
     def pre(self, noisy_latents, timesteps):
         self.t_raw = timesteps
@@ -22,18 +25,18 @@ class DreamArtistPTContext(CFGContext):
 
     def post(self, model_pred):
         e_t_uncond, e_t = model_pred.chunk(2)
-        if self.cfg_scale[0] != self.cfg_scale[1]:
-            rate = self.t_raw / (self.num_train_timesteps - 1)
-            if self.cfg_scale[2] == 'cos':
-                rate = torch.cos((rate - 1) * math.pi / 2)
-            elif self.cfg_scale[2] == 'cos2':
-                rate = 1 - torch.cos(rate * math.pi / 2)
-            elif self.cfg_scale[2] == 'ln':
+        if self.cfg_low != self.cfg_high:
+            rate = self.t_raw/(self.num_train_timesteps-1)
+            if self.cfg_func == 'cos':
+                rate = torch.cos((rate-1)*math.pi/2)
+            elif self.cfg_func == 'cos2':
+                rate = 1-torch.cos(rate*math.pi/2)
+            elif self.cfg_func == 'ln':
                 pass
             else:
-                rate = eval(self.cfg_scale[2])
-            rate = rate.view(-1,1,1,1)
+                rate = self.cfg_func(rate)
+            rate = rate.view(-1, 1, 1, 1)
         else:
             rate = 1
-        model_pred = e_t_uncond + ((self.cfg_scale[1] - self.cfg_scale[0]) * rate + self.cfg_scale[0]) * (e_t - e_t_uncond)
-        return model_pred
+        model_pred = e_t_uncond+((self.cfg_high-self.cfg_low)*rate+self.cfg_low)*(e_t-e_t_uncond)
+        return model_pred

hcpdiff/models/compose/compose_hook.py

@@ -38,42 +38,42 @@ class ComposeEmbPTHook(nn.Module):
             hook.remove()
 
     @classmethod
-    def hook(cls, ex_words_emb: Dict[str, ParameterGroup], tokenizer, text_encoder, log=False, **kwargs):
+    def hook(cls, ex_words_emb: Dict[str, ParameterGroup], tokenizer, text_encoder, **kwargs):
         if isinstance(text_encoder, ComposeTextEncoder):
             hook_list = []
 
             emb_len = 0
-            for i, (name, tokenizer_i) in enumerate(tokenizer.tokenizer_list):
+            for i, name in enumerate(tokenizer.tokenizer_names):
                 text_encoder_i = getattr(text_encoder, name)
-                if log:
-                    logger.info(f'compose hook: {name}')
+                tokenizer_i = getattr(tokenizer, name)
                 embedding_dim = text_encoder_i.get_input_embeddings().embedding_dim
                 ex_words_emb_i = {k:v[i] for k, v in ex_words_emb.items()}
                 emb_len += embedding_dim
-                hook_list.append((name, EmbeddingPTHook.hook(ex_words_emb_i, tokenizer_i, text_encoder_i, log=log, **kwargs)))
+                hook_list.append((name, EmbeddingPTHook.hook(ex_words_emb_i, tokenizer_i, text_encoder_i, **kwargs)))
 
             return cls(hook_list)
         else:
-            return EmbeddingPTHook.hook(ex_words_emb, tokenizer, text_encoder, log, **kwargs)
+            return EmbeddingPTHook.hook(ex_words_emb, tokenizer, text_encoder, **kwargs)
 
     @classmethod
-    def hook_from_dir(cls, emb_dir, tokenizer, text_encoder, log=True, device='cuda:0', **kwargs) -> Union[
+    def hook_from_dir(cls, emb_dir, tokenizer, text_encoder, device='cuda:0', **kwargs) -> Union[
         Tuple['ComposeEmbPTHook', Dict], Tuple[EmbeddingPTHook, Dict]]:
         if isinstance(text_encoder, ComposeTextEncoder):
             # multi text encoder
-            #ex_words_emb = {file[:-3]:load_emb(os.path.join(emb_dir, file)).to(device) for file in os.listdir(emb_dir) if file.endswith('.pt')}
+            # ex_words_emb = {file[:-3]:load_emb(os.path.join(emb_dir, file)).to(device) for file in os.listdir(emb_dir) if file.endswith('.pt')}
 
             # slice of nn.Parameter cannot return grad. Split the tensor
             ex_words_emb = {}
-            emb_dims = [x.embedding_dim for x in text_encoder.get_input_embeddings()]
-            for file in os.listdir(emb_dir):
-                if file.endswith('.pt'):
-                    emb = load_emb(os.path.join(emb_dir, file)).to(device)
-                    emb = ParameterGroup([nn.Parameter(item, requires_grad=False) for item in emb.split(emb_dims, dim=1)])
-                    ex_words_emb[file[:-3]] = emb
-            return cls.hook(ex_words_emb, tokenizer, text_encoder, log, **kwargs), ex_words_emb
+            if emb_dir is not None and os.path.exists(emb_dir):
+                emb_dims = [x.embedding_dim for x in text_encoder.get_input_embeddings()]
+                for file in os.listdir(emb_dir):
+                    if file.endswith('.pt'):
+                        emb = load_emb(os.path.join(emb_dir, file)).to(device)
+                        emb = ParameterGroup([nn.Parameter(item, requires_grad=False) for item in emb.split(emb_dims, dim=1)])
+                        ex_words_emb[file[:-3]] = emb
+            return cls.hook(ex_words_emb, tokenizer, text_encoder, **kwargs), ex_words_emb
         else:
-            return EmbeddingPTHook.hook_from_dir(emb_dir, tokenizer, text_encoder, log, device, **kwargs)
+            return EmbeddingPTHook.hook_from_dir(emb_dir, tokenizer, text_encoder, **kwargs)
 
 class ComposeTEEXHook:
     def __init__(self, tehook_list: List[Tuple[str, TEEXHook]], cat_dim=-1):
@@ -98,10 +98,28 @@ class ComposeTEEXHook:
         for name, tehook in self.tehook_list:
             tehook.clip_skip = value
 
+    @property
+    def clip_final_norm(self):
+        return self.tehook_list[0][1].clip_final_norm
+
+    @clip_final_norm.setter
+    def clip_final_norm(self, value: bool):
+        for name, tehook in self.tehook_list:
+            tehook.clip_final_norm = value
+
+    @property
+    def use_attention_mask(self):
+        return self.tehook_list[0][1].use_attention_mask
+
+    @use_attention_mask.setter
+    def use_attention_mask(self, value: bool):
+        for name, tehook in self.tehook_list:
+            tehook.use_attention_mask = value
+
     def encode_prompt_to_emb(self, prompt):
         emb_list = [tehook.encode_prompt_to_emb(prompt) for name, tehook in self.tehook_list]
-        encoder_hidden_states, pooled_output = list(zip(*emb_list))
-        return torch.cat(encoder_hidden_states, dim=self.cat_dim), pooled_output
+        encoder_hidden_states, pooled_output, attention_mask = list(zip(*emb_list))
+        return torch.cat(encoder_hidden_states, dim=self.cat_dim), pooled_output, attention_mask[0]
 
     def enable_xformers(self):
         for name, tehook in self.tehook_list:
@@ -112,16 +130,19 @@ class ComposeTEEXHook:
         return TEEXHook.mult_attn(prompt_embeds, attn_mult)
 
     @classmethod
-    def hook(cls, text_enc: nn.Module, tokenizer, N_repeats=3, clip_skip=0, clip_final_norm=True, device='cuda', use_attention_mask=False) -> Union['ComposeTEEXHook', TEEXHook]:
+    def hook(cls, text_enc: nn.Module, tokenizer, N_repeats=3, clip_skip=0, clip_final_norm=True, use_attention_mask=False) -> Union[
+        'ComposeTEEXHook', TEEXHook]:
         if isinstance(text_enc, ComposeTextEncoder):
             # multi text encoder
-            tehook_list = [(name, TEEXHook.hook(getattr(text_enc, name), tokenizer_i, N_repeats, clip_skip, clip_final_norm, device=device, use_attention_mask=use_attention_mask))
-                for name, tokenizer_i in tokenizer.tokenizer_list]
+            tehook_list = [(name, TEEXHook.hook(getattr(text_enc, name), getattr(tokenizer, name), N_repeats, clip_skip, clip_final_norm,
+                                                use_attention_mask=use_attention_mask))
+                for name in tokenizer.tokenizer_names]
             return cls(tehook_list)
         else:
             # single text encoder
-            return TEEXHook.hook(text_enc, tokenizer, N_repeats, clip_skip, device=device, use_attention_mask=use_attention_mask)
+            return TEEXHook.hook(text_enc, tokenizer, N_repeats, clip_skip, clip_final_norm, use_attention_mask=use_attention_mask)
 
     @classmethod
     def hook_pipe(cls, pipe, N_repeats=3, clip_skip=0, clip_final_norm=True, use_attention_mask=False):
-        return cls.hook(pipe.text_encoder, pipe.tokenizer, N_repeats=N_repeats, device='cuda', clip_skip=clip_skip, clip_final_norm=clip_final_norm, use_attention_mask=use_attention_mask)
+        return cls.hook(pipe.text_encoder, pipe.tokenizer, N_repeats=N_repeats, clip_skip=clip_skip, clip_final_norm=clip_final_norm,
+                        use_attention_mask=use_attention_mask)

hcpdiff/models/compose/compose_tokenizer.py

@@ -18,14 +18,19 @@ from transformers.tokenization_utils_base import BatchEncoding
 class ComposeTokenizer(PreTrainedTokenizer):
     def __init__(self, tokenizer_list: List[Tuple[str, CLIPTokenizer]], cat_dim=-1):
         self.cat_dim = cat_dim
-        self.tokenizer_list = tokenizer_list
+
+        self.tokenizer_names = []
+        for name, tokenizer in tokenizer_list:
+            setattr(self, name, tokenizer)
+            self.tokenizer_names.append(name)
+
         super().__init__()
 
-        self.model_max_length = self.first_tokenizer.model_max_length
+        self.model_max_length = torch.tensor([tokenizer.model_max_length for name, tokenizer in tokenizer_list])
 
     @property
     def first_tokenizer(self):
-        return self.tokenizer_list[0][1]
+        return getattr(self, self.tokenizer_names[0])
 
     @property
     def vocab_size(self):
@@ -40,18 +45,26 @@ class ComposeTokenizer(PreTrainedTokenizer):
         return self.first_tokenizer.bos_token_id
 
     def get_vocab(self):
-        return dict(self.first_tokenizer.encoder, **self.first_tokenizer.added_tokens_encoder)
+        return self.first_tokenizer.get_vocab()
 
     def tokenize(self, text, **kwargs) -> List[str]:
         return self.first_tokenizer.tokenize(text, **kwargs)
 
     def add_tokens( self, new_tokens, special_tokens: bool = False) -> List[int]:
-        return [tokenizer.add_tokens(new_tokens, special_tokens) for name, tokenizer in self.tokenizer_list]
+        return [getattr(self, name).add_tokens(new_tokens, special_tokens) for name in self.tokenizer_names]
+    
+    def save_vocabulary(self, save_directory: str, filename_prefix = None) -> Tuple[str]:
+        return self.first_tokenizer.save_vocabulary(save_directory, filename_prefix)
+
+    def __call__(self, text, *args, max_length=None, **kwargs):
+        if isinstance(max_length, torch.Tensor):
+            token_list: List[BatchEncoding] = [getattr(self, name)(text, *args, max_length=max_length_i, **kwargs)
+                for name, max_length_i in zip(self.tokenizer_names, max_length)]
+        else:
+            token_list: List[BatchEncoding] = [getattr(self, name)(text, *args, max_length=max_length, **kwargs) for name in self.tokenizer_names]
 
-    def __call__(self, text, *args, **kwargs):
-        token_list: List[BatchEncoding] = [tokenizer(text, *args, **kwargs) for name, tokenizer in self.tokenizer_list]
         input_ids = torch.cat([token.input_ids for token in token_list], dim=-1)  # [N_tokenizer, N_token]
-        attention_mask = [token.attention_mask for token in token_list]
+        attention_mask = torch.cat([token.attention_mask for token in token_list], dim=-1)
         return BatchEncoding({'input_ids':input_ids, 'attention_mask':attention_mask})
 
     @classmethod

hcpdiff/models/compose/sdxl_composer.py

@@ -27,13 +27,13 @@ class CLIPTextModelWithProjection_Align(CLIPTextModelWithProjection):
 class SDXLTextEncoder(ComposeTextEncoder):
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: str, *args, subfolder=None, revision:str=None, **kwargs):
-        clip_B = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder', **kwargs)
+        clip_L = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder', **kwargs)
         clip_bigG = CLIPTextModelWithProjection_Align.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder_2', **kwargs)
-        return cls([('clip_B', clip_B), ('clip_bigG', clip_bigG)])
+        return cls([('clip_L', clip_L), ('clip_bigG', clip_bigG)])
 
 class SDXLTokenizer(ComposeTokenizer):
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: str, *args, subfolder=None, revision:str=None, **kwargs):
-        clip_B = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer', **kwargs)
+        clip_L = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer', **kwargs)
         clip_bigG = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer_2', **kwargs)
-        return cls([('clip_B', clip_B), ('clip_bigG', clip_bigG)])
+        return cls([('clip_L', clip_L), ('clip_bigG', clip_bigG)])

hcpdiff/models/controlnet.py

@@ -5,7 +5,7 @@ import torch
 from torch import nn
 from copy import deepcopy
 
-from .plugin import MultiPluginBlock, BasePluginBlock
+from rainbowneko.models.plugin import MultiPluginBlock, BasePluginBlock
 from hcpdiff.utils.net_utils import remove_all_hooks, remove_layers
 
 class ControlNetPlugin(MultiPluginBlock):
@@ -55,25 +55,25 @@ class ControlNetPlugin(MultiPluginBlock):
         self.cond_head = nn.Sequential(*cond_head)
 
     def reset_parameters(self) -> None:
-        def weight_init(m):
-            if isinstance(m, nn.Conv2d):
-                nn.init.constant_(m.weight, 0)
-        self.controlnet_down_blocks.apply(weight_init)
-        self.controlnet_mid_block.apply(weight_init)
-        self.cond_head[-1].apply(weight_init)
-
-    def from_layer_hook(self, host, fea_in:Tuple[torch.Tensor], fea_out:Tuple[torch.Tensor], idx: int):
+        def zero_weight_init(m):
+            for p in m.parameters():
+                p.detach().zero_()
+        self.controlnet_down_blocks.apply(zero_weight_init)
+        self.controlnet_mid_block.apply(zero_weight_init)
+        self.cond_head[-1].apply(zero_weight_init)
+
+    def from_layer_hook(self, host, idx: int, args: Tuple[Any, ...], kwargs: Dict[str, Any], fea_out: Any=None):
         if idx==0:
-            self.data_input = fea_in
+            self.data_input = (args, kwargs)
         elif idx==1:
-            self.feat_to = self(*self.data_input)
+            self.feat_to = self(*self.data_input[0], **self.data_input[1])
 
-    def to_layer_hook(self, host, fea_in:Tuple[torch.Tensor], fea_out:Tuple[torch.Tensor], idx: int):
+    def to_layer_hook(self, host, idx: int, args: Tuple[Any, ...], kwargs: Dict[str, Any], fea_out: Any=None):
         if idx == 5:
-            sp = fea_in[0].shape[1]//2
-            new_feat = fea_in[0].clone()
-            new_feat[:, sp:, ...] = fea_in[0][:, sp:, ...] + self.feat_to[0]
-            return (new_feat, fea_in[1])
+            sp = args[0].shape[1]//2
+            new_feat = args[0].clone()
+            new_feat[:, sp:, ...] = args[0][:, sp:, ...] + self.feat_to[0]
+            return (new_feat, args[1])
         elif idx == 3:
             return (fea_out[0], tuple(fea_out[1][i] + self.feat_to[(idx) * 3 + i+1] for i in range(2)))
         elif idx == 4:

hcpdiff/models/lora_base_patch.py

@@ -13,7 +13,7 @@ from torch import nn
 from torch.nn import functional as F
 
 from hcpdiff.utils.utils import make_mask, low_rank_approximate, isinstance_list
-from .plugin import PatchPluginBlock, PluginGroup, PatchPluginContainer
+from rainbowneko.models.plugin import PatchPluginBlock, PluginGroup, PatchPluginContainer
 
 from typing import Union, Tuple, Dict, Type
 
@@ -38,9 +38,9 @@ class LoraBlock(PatchPluginBlock):
     container_cls = LoraPatchContainer
     wrapable_classes = (nn.Linear, nn.Conv2d)
 
-    def __init__(self, lora_id:int, host:Union[nn.Linear, nn.Conv2d], rank, dropout=0.1, alpha=1.0, bias=False,
+    def __init__(self, name:int, host:Union[nn.Linear, nn.Conv2d], rank, dropout=0.1, alpha=1.0, bias=False,
                  alpha_auto_scale=True, parent_block=None, host_name=None, **kwargs):
-        super().__init__(f'lora_block_{lora_id}', host, parent_block=parent_block, host_name=host_name)
+        super().__init__(name, host, parent_block=parent_block, host_name=host_name)
 
         self.bias=bias
 
@@ -56,8 +56,14 @@ class LoraBlock(PatchPluginBlock):
         self.dropout = nn.Dropout(dropout)
 
         self.rank = self.layer.rank
+        self.alpha_auto_scale = alpha_auto_scale
         self.register_buffer('alpha', torch.tensor(alpha/self.rank if alpha_auto_scale else alpha))
 
+    def set_hyper_params(self, alpha=None, **kwargs):
+        if alpha is not None:
+            self.register_buffer('alpha', torch.tensor(alpha/self.rank if self.alpha_auto_scale else alpha))
+        super().set_hyper_params(**kwargs)
+
     def get_weight(self):
         return self.layer.get_weight() * self.alpha
 
@@ -91,7 +97,7 @@ class LoraBlock(PatchPluginBlock):
             host.weight.data * base_alpha + alpha * re_w.to(host.weight.device, dtype=host.weight.dtype)
         )
 
-        if self.layer.lora_up.bias is not None:
+        if re_b is not None:
             if host.bias is None:
                 host.bias = nn.Parameter(re_b.to(host.weight.device, dtype=host.weight.dtype))
             else:
@@ -145,32 +151,15 @@ class LoraBlock(PatchPluginBlock):
             pass
 
     @classmethod
-    def wrap_layer(cls, lora_id:int, layer: Union[nn.Linear, nn.Conv2d], rank=1, dropout=0.0, alpha=1.0, svd_init=False,
+    def wrap_layer(cls, name:str, host: Union[nn.Linear, nn.Conv2d], rank=1, dropout=0.0, alpha=1.0, svd_init=False,
                    bias=False, mask=None, **kwargs):# -> LoraBlock:
-        lora_block = cls(lora_id, layer, rank, dropout, alpha, bias=bias, **kwargs)
+        lora_block = cls(name, host, rank, dropout, alpha, bias=bias, **kwargs)
         lora_block.init_weights(svd_init)
         return lora_block
 
     @classmethod
-    def wrap_model(cls, lora_id:int, model: nn.Module, **kwargs):# -> Dict[str, LoraBlock]:
-        return super(LoraBlock, cls).wrap_model(lora_id, model, exclude_classes=(LoraBlock,), **kwargs)
-
-    @staticmethod
-    def extract_lora_state(model:nn.Module):
-        return {k:v for k,v in model.state_dict().items() if 'lora_block_' in k}
-
-    @staticmethod
-    def extract_state_without_lora(model:nn.Module):
-        return {k:v for k,v in model.state_dict().items() if 'lora_block_' not in k}
-
-    @staticmethod
-    def extract_param_without_lora(model:nn.Module):
-        return {k:v for k,v in model.named_parameters() if 'lora_block_' not in k}
-
-    @staticmethod
-    def extract_trainable_state_without_lora(model:nn.Module):
-        trainable_keys = {k for k,v in model.named_parameters() if ('lora_block_' not in k) and v.requires_grad}
-        return {k: v for k, v in model.state_dict().items() if k in trainable_keys}
+    def wrap_model(cls, name:str, host: nn.Module, **kwargs):# -> Dict[str, LoraBlock]:
+        return super().wrap_model(name, host, exclude_classes=(LoraBlock,), **kwargs)
 
 class LoraGroup(PluginGroup):
     def set_mask(self, batch_mask):

hcpdiff/models/lora_layers.py

@@ -15,7 +15,7 @@ from einops import repeat, rearrange, einsum
 from torch import nn
 
 from .lora_base import LoraBlock
-from .layers import GroupLinear
+from rainbowneko.models.layers import GroupLinear
 import warnings
 
 class LoraLayer(LoraBlock):
@@ -59,8 +59,8 @@ class LoraLayerGroup(LoraBlock):
         def __init__(self, host, rank, bias, dropout, block):
             super().__init__(host, rank, bias, dropout, block)
             self.register_buffer('rank_groups', torch.tensor(block.rank_groups_raw, dtype=torch.int))
-            self.lora_down = GroupLinear(host.in_features*self.rank_groups, self.rank, groups=self.rank_groups, bias=False)
-            self.lora_up = GroupLinear(self.rank, host.out_features*self.rank_groups, groups=self.rank_groups, bias=bias)
+            self.lora_down = GroupLinear(host.in_features, self.rank//self.rank_groups, group=self.rank_groups, bias=False)
+            self.lora_up = GroupLinear(self.rank//self.rank_groups, host.out_features, group=self.rank_groups, bias=bias)
 
         def feed_svd(self, U, V, weight):
             self.lora_up.weight.data = rearrange(U, 'o (g ri) -> g ri o', g=self.rank_groups).to(device=weight.device, dtype=weight.dtype)
@@ -137,9 +137,3 @@ class LohaLayer(LoraBlock):
             w = torch.prod(einsum(self.W_up.data, self.W_down.data, 'g o r ..., g r i ... -> g o i ...'), dim=0)
             b = None
             return w, b
-
-lora_layer_map={
-    'lora': LoraLayer,
-    'loha_group': LoraLayerGroup,
-    'loha': LohaLayer,
-}

hcpdiff/models/lora_layers_patch.py

@@ -8,19 +8,18 @@ lora_layers.py
     :Licence:     Apache-2.0
 """
 
+import math
+
 import torch
-from einops import einsum, rearrange
+from einops import einsum
 from torch import nn
 from torch.nn import functional as F
 
 from .lora_base_patch import LoraBlock, PatchPluginContainer
-from .layers import GroupLinear
-import math
-from typing import Union, List
 
 class LoraLayer(LoraBlock):
-    def __init__(self, lora_id: int, host, rank=1, dropout=0.1, alpha=1.0, bias=False, alpha_auto_scale=True, **kwargs):
-        super().__init__(lora_id, host, rank, dropout, alpha=alpha, bias=bias, alpha_auto_scale=alpha_auto_scale, **kwargs)
+    def __init__(self, name: str, host, rank=1, dropout=0.0, alpha=1.0, bias=False, alpha_auto_scale=True, **kwargs):
+        super().__init__(name, host, rank, dropout, alpha=alpha, bias=bias, alpha_auto_scale=alpha_auto_scale, **kwargs)
 
     class LinearLayer(LoraBlock.LinearLayer):
         def __init__(self, host:nn.Linear, rank, bias, block):
@@ -99,6 +98,11 @@ class LoraLayer(LoraBlock):
             b = self.bias.data if self.bias else None
             return w, b
 
+def none_add(a, b):
+    if a is None:
+        return b
+    return a+b
+
 class DAPPPatchContainer(PatchPluginContainer):
     def forward(self, x, *args, **kwargs):
         weight_p = None
@@ -107,25 +111,11 @@ class DAPPPatchContainer(PatchPluginContainer):
         bias_n = None
         for name in self.plugin_names:
             if self[name].branch=='p':
-                if weight_p is None:
-                    weight_p = self[name].get_weight()
-                else:
-                    weight_p = weight_p + self[name].get_weight()
-
-                if bias_p is None:
-                    bias_p = self[name].get_bias()
-                else:
-                    bias_p = bias_p+self[name].get_bias()
+                weight_p = none_add(weight_p, self[name].get_weight())
+                bias_p = none_add(bias_p, self[name].get_bias())
             elif self[name].branch=='n':
-                if weight_n is None:
-                    weight_n = self[name].get_weight()
-                else:
-                    weight_n = weight_n + self[name].get_weight()
-
-                if bias_n is None:
-                    bias_n = self[name].get_bias()
-                else:
-                    bias_n = bias_n+self[name].get_bias()
+                weight_n = none_add(weight_n, self[name].get_weight())
+                bias_n = none_add(bias_n, self[name].get_bias())
 
         B = x.shape[0]//2
         x_p = self[name].post_forward(x[B:], self._host.weight, weight_p, self._host.bias, bias_p)

hcpdiff/models/text_emb_ex.py

@@ -7,16 +7,17 @@ text_emb_ex.py
     :Created:     10/03/2023
     :Licence:     Apache-2.0
 """
-from typing import Tuple
+from typing import Tuple, Dict, Any
 
 import torch
 from torch import nn
 import os
-from loguru import logger
+from rainbowneko import _share
 from einops import rearrange, repeat
+import torch.nn.functional as F
 
 from ..utils.net_utils import load_emb
-from .plugin import SinglePluginBlock
+from rainbowneko.models.plugin import SinglePluginBlock
 
 class EmbeddingPTHook(SinglePluginBlock):
     def __init__(self, token_embedding:nn.Embedding, N_word=75, N_repeats=3):
@@ -37,6 +38,84 @@ class EmbeddingPTHook(SinglePluginBlock):
         self.input_ids = rearrange(input_ids[0], '(b r) w -> b (r w)', r=self.N_repeats)  # 兼容Attention mask
         return self.input_ids.clip(0, self.num_embeddings-1)
 
+    def forward(self, inputs_embeds:torch.Tensor, *args: Tuple[Any, ...], **kwargs: Dict[str, Any]):
+        '''
+        :param input_ids: [B, N_ids]
+        :param inputs_embeds: [B, N_repeat*(N_word+2), N_emb]
+        :return: [B, N_repeat, N_word+2, N_emb]
+        '''
+        rep_idxs_B = self.input_ids >= self.num_embeddings
+        BOS = repeat(inputs_embeds[:,0,:], 'b e -> b r 1 e', r=self.N_repeats)
+        EOS = repeat(inputs_embeds[:,-1,:], 'b e -> b r 1 e', r=self.N_repeats)
+
+        replaced_embeds = []
+        for i, (item, rep_idxs, ids_raw) in enumerate(zip(inputs_embeds, rep_idxs_B, self.input_ids)):
+            # insert pt to embeddings
+            rep_idxs=torch.where(rep_idxs)[0]
+            item_new=[]
+            rep_idx_last=0
+            for rep_idx in rep_idxs:
+                rep_idx=rep_idx.item()
+                item_new.append(item[rep_idx_last:rep_idx, :])
+                item_new.append(self.emb[ids_raw[rep_idx].item()].to(dtype=item.dtype))
+                rep_idx_last=rep_idx+1
+            item_new.append(item[rep_idx_last:, :])
+
+            # split to N_repeat sentence
+            replaced_item = torch.cat(item_new, dim=0)[1:self.N_word*self.N_repeats+1, :]
+            replaced_item = rearrange(replaced_item, '(r w) e -> r w e', r=self.N_repeats, w=self.N_word)
+            replaced_item = torch.cat([BOS[i], replaced_item, EOS[i]], dim=1) # [N_repeat, N_word+2, N_emb]
+
+            replaced_embeds.append(replaced_item)
+        return torch.cat(replaced_embeds, dim=0) # [B*N_repeat, N_word+2, N_emb]
+
+    def remove(self):
+        super(EmbeddingPTHook, self).remove()
+        self.handle_pre.remove()
+
+    @classmethod
+    def hook(cls, ex_words_emb, tokenizer, text_encoder, **kwargs):
+        word_list = list(ex_words_emb.keys())
+        tokenizer.add_tokens(word_list)
+        token_ids = tokenizer(' '.join(word_list)).input_ids[1:-1]
+
+        embedding_hook = cls(text_encoder.get_input_embeddings(), N_word=tokenizer.model_max_length-2, **kwargs)
+        #text_encoder.text_model.embeddings.token_embedding = embedding_hook
+        for tid, word in zip(token_ids, word_list):
+            embedding_hook.add_emb(ex_words_emb[word], tid)
+            _share.loggers.info(f'hook: {word}, len: {ex_words_emb[word].shape[0]}, id: {tid}')
+        return embedding_hook
+
+    @classmethod
+    def hook_from_dir(cls, emb_dir, tokenizer, text_encoder, device='cuda:0', **kwargs):
+        ex_words_emb = {file[:-3]: nn.Parameter(load_emb(os.path.join(emb_dir, file)).to(device), requires_grad=False)
+                        for file in os.listdir(emb_dir) if file.endswith('.pt')}
+        return cls.hook(ex_words_emb, tokenizer, text_encoder, **kwargs), ex_words_emb
+
+class EmbeddingPTInterpHook(SinglePluginBlock):
+    def __init__(self, token_embedding:nn.Embedding, N_word=75, N_repeats=3):
+        super().__init__('emb_ex', token_embedding)
+        self.handle_pre = token_embedding.register_forward_pre_hook(self.pre_hook)
+
+        new_len = int(token_embedding.num_embeddings*N_repeats)
+        original_weights = token_embedding.weight.data.unsqueeze(1)
+        token_embedding.weight.data = F.interpolate(original_weights, size=new_len, mode='linear', align_corners=False).squeeze(1)
+        token_embedding.num_embeddings = new_len
+
+        self.N_word=N_word
+        self.N_repeats=N_repeats
+        self.num_embeddings=token_embedding.num_embeddings
+        self.embedding_dim=token_embedding.embedding_dim
+        self.emb={}
+        self.emb_train=nn.ParameterList()
+
+    def add_emb(self, emb:nn.Parameter, token_id:int):
+        self.emb[token_id]=emb
+
+    def pre_hook(self, host, input_ids: Tuple[torch.Tensor]):
+        self.input_ids = rearrange(input_ids[0], '(b r) w -> b (r w)', r=self.N_repeats)  # 兼容Attention mask
+        return self.input_ids.clip(0, self.num_embeddings-1)
+
     def forward(self, fea_in:Tuple[torch.Tensor], inputs_embeds:torch.Tensor):
         '''
         :param input_ids: [B, N_ids]
@@ -83,12 +162,11 @@ class EmbeddingPTHook(SinglePluginBlock):
         for tid, word in zip(token_ids, word_list):
             embedding_hook.add_emb(ex_words_emb[word], tid)
             if log:
-                logger.info(f'hook: {word}, len: {ex_words_emb[word].shape[0]}, id: {tid}')
+                _share.logger.info(f'hook: {word}, len: {ex_words_emb[word].shape[0]}, id: {tid}')
         return embedding_hook
 
     @classmethod
     def hook_from_dir(cls, emb_dir, tokenizer, text_encoder, log=True, device='cuda:0', **kwargs):
         ex_words_emb = {file[:-3]: nn.Parameter(load_emb(os.path.join(emb_dir, file)).to(device), requires_grad=False)
                         for file in os.listdir(emb_dir) if file.endswith('.pt')}
-        return cls.hook(ex_words_emb, tokenizer, text_encoder, log, **kwargs), ex_words_emb
-
+        return cls.hook(ex_words_emb, tokenizer, text_encoder, log, **kwargs), ex_words_emb