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

hcpdiff/loss/ssim.py

@@ -0,0 +1,37 @@
+from pytorch_msssim import SSIM, MS_SSIM
+from torch.nn.modules.loss import _Loss
+import torch
+
+class SSIMLoss(_Loss):
+    target_type = 'x0'
+
+    def __init__(self, size_average=None, reduce=None, reduction: str = 'mean'):
+        super().__init__(size_average=size_average, reduce=reduce, reduction=reduction)
+        self.ssim = SSIM(data_range=1., size_average=False, channel=4)
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        '''
+
+        :param input: [B,C,H,W]
+        :param target: [B,C,H,W]
+        :return: [B,1,1,1]
+        '''
+        input = (input+1)/2
+        target = (target+1)/2
+        return 1-self.ssim(input, target).view(-1,1,1,1)
+
+class MS_SSIMLoss(_Loss):
+    def __init__(self, size_average=None, reduce=None, reduction: str = 'mean'):
+        super().__init__(size_average=size_average, reduce=reduce, reduction=reduction)
+        self.ssim = MS_SSIM(data_range=1., size_average=False, channel=4)
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor) -> torch.Tensor:
+        '''
+
+        :param input: [B,C,H,W]
+        :param target: [B,C,H,W]
+        :return: [B,1,1,1]
+        '''
+        input = (input+1)/2
+        target = (target+1)/2
+        return 1-self.ssim(input, target).view(-1,1,1,1)

hcpdiff/loss/vlb.py

@@ -0,0 +1,79 @@
+import torch
+from torch import nn
+import numpy as np
+
+class VLBLoss(nn.Module):
+    need_sigma = True
+    need_timesteps = True
+    need_sampler = True
+    var_pred = True
+
+    def __init__(self, loss, weight: float = 1.):
+        super().__init__()
+        self.loss = loss
+        self.weight = weight
+
+    def normal_kl(self, mean1, logvar1, mean2, logvar2):
+        """
+        Compute the KL divergence between two gaussians.
+        """
+
+        return 0.5*(-1.0+logvar2-logvar1+(logvar1-logvar2).exp()+((mean1-mean2)**2)*(-logvar2).exp())
+
+    def forward(self, input: torch.Tensor, target: torch.Tensor, sigma, timesteps: torch.Tensor, x_t: torch.Tensor, sampler):
+        eps_pred, var_pred = input.chunk(2, dim=1)
+        x0_pred = sampler.eps_to_x0(eps_pred, x_t, sigma)
+
+        true_mean = sampler.sigma_scheduler.get_post_mean(timesteps, target, x_t)
+        true_logvar = sampler.sigma_scheduler.get_post_log_var(timesteps)
+
+        pred_mean = sampler.sigma_scheduler.get_post_mean(timesteps, x0_pred, x_t)
+        pred_logvar = sampler.sigma_scheduler.get_post_log_var(timesteps, x_t_var=var_pred)
+
+        kl = self.normal_kl(true_mean, true_logvar, pred_mean, pred_logvar)
+        kl = kl.mean(dim=(1,2,3))/np.log(2.0)
+
+        decoder_nll = -self.discretized_gaussian_log_likelihood(target, means=pred_mean, log_scales=0.5*pred_logvar)
+        assert decoder_nll.shape == target.shape
+        decoder_nll = decoder_nll.mean(dim=(1,2,3))/np.log(2.0)
+
+        # At the first timestep return the decoder NLL,
+        # otherwise return KL(q(x_{t-1}|x_t,x_0) || p(x_{t-1}|x_t))
+        output = torch.where((timesteps == 0), decoder_nll, kl)
+
+        return self.weight*output
+
+    def approx_standard_normal_cdf(self, x):
+        """
+        A fast approximation of the cumulative distribution function of the
+        standard normal.
+        """
+        return 0.5*(1.0+torch.tanh(np.sqrt(2.0/np.pi)*(x+0.044715*torch.pow(x, 3))))
+
+    def discretized_gaussian_log_likelihood(self, x, *, means, log_scales):
+        """
+        Compute the log-likelihood of a Gaussian distribution discretizing to a
+        given image.
+        :param x: the target images. It is assumed that this was uint8 values,
+                  rescaled to the range [-1, 1].
+        :param means: the Gaussian mean Tensor.
+        :param log_scales: the Gaussian log stddev Tensor.
+        :return: a tensor like x of log probabilities (in nats).
+        """
+        assert x.shape == means.shape == log_scales.shape
+        centered_x = x-means
+        inv_stdv = torch.exp(-log_scales)
+        plus_in = inv_stdv*(centered_x+1.0/255.0)
+        cdf_plus = self.approx_standard_normal_cdf(plus_in)
+        min_in = inv_stdv*(centered_x-1.0/255.0)
+        cdf_min = self.approx_standard_normal_cdf(min_in)
+        log_cdf_plus = torch.log(cdf_plus.clamp(min=1e-12))
+        log_one_minus_cdf_min = torch.log((1.0-cdf_min).clamp(min=1e-12))
+        cdf_delta = cdf_plus-cdf_min
+        log_probs = torch.where(
+            x<-0.999,
+            log_cdf_plus,
+            torch.where(x>0.999, log_one_minus_cdf_min, torch.log(cdf_delta.clamp(min=1e-12))),
+        )
+        assert log_probs.shape == x.shape
+        return log_probs

hcpdiff/loss/weighting.py

@@ -0,0 +1,66 @@
+from torch import nn
+
+from .base import DiffusionLossContainer
+
+class LossWeight(nn.Module):
+    def __init__(self, loss: DiffusionLossContainer):
+        super().__init__()
+        self.loss = loss
+
+    def get_weight(self, pred, inputs):
+        '''
+
+        :param input: [B,C,H,W]
+        :param target: [B,C,H,W]
+        :return: [B,1,1,1] or [B,C,H,W]
+        '''
+        raise NotImplementedError
+
+    def forward(self, pred, inputs):
+        '''
+        weight: [B,1,1,1] or [B,C,H,W]
+        loss: [B,*,*,*]
+        '''
+        return self.get_weight(pred, inputs)*self.loss(pred, inputs)
+
+class SNRWeight(LossWeight):
+    def get_weight(self, pred, inputs):
+        if self.loss.target_type == 'eps':
+            return 1
+        elif self.loss.target_type == "x0":
+            sigma = pred['sigma']
+            return (1./sigma**2).view(-1, 1, 1, 1)
+        else:
+            raise ValueError(f"{self.__class__.__name__} is not support for target_type {self.loss.target_type}")
+
+class MinSNRWeight(LossWeight):
+    def __init__(self, loss: DiffusionLossContainer, gamma: float = 1.):
+        super().__init__(loss)
+        self.gamma = gamma
+
+    def get_weight(self, pred, inputs):
+        sigma = pred['sigma']
+        if self.loss.target_type == 'eps':
+            w_snr = (self.gamma*sigma**2).clip(max=1).float()
+        elif self.loss.target_type == "x0":
+            w_snr = (1/(sigma**2)).clip(max=self.gamma).float()
+        else:
+            raise ValueError(f"{self.__class__.__name__} is not support for target_type {self.loss.target_type}")
+
+        return w_snr.view(-1, 1, 1, 1)
+
+class EDMWeight(LossWeight):
+    def __init__(self, loss: DiffusionLossContainer, gamma: float = 1.):
+        super().__init__(loss)
+        self.gamma = gamma
+
+    def get_weight(self, pred, inputs):
+        sigma = pred['sigma']
+        if self.loss.target_type == 'eps':
+            w_snr = ((sigma**2+self.gamma**2)/(self.gamma**2)).float()
+        elif self.loss.target_type == "x0":
+            w_snr = ((sigma**2+self.gamma**2)/((sigma*self.gamma)**2)).float()
+        else:
+            raise ValueError(f"{self.__class__.__name__} is not support for target_type {self.loss.target_type}")
+
+        return w_snr.view(-1, 1, 1, 1)

hcpdiff/models/__init__.py

@@ -1,4 +1,3 @@
-from .plugin import PluginBlock, PluginGroup, SinglePluginBlock, MultiPluginBlock, PatchPluginBlock
 # from .lora_base import LoraBlock, LoraGroup
 # from .lora_layers import lora_layer_map
 from .lora_base_patch import LoraBlock, LoraGroup
@@ -7,4 +6,5 @@ from .text_emb_ex import EmbeddingPTHook
 from .textencoder_ex import TEEXHook
 from .tokenizer_ex import TokenizerHook
 from .cfg_context import CFGContext, DreamArtistPTContext
-from .wrapper import TEUnetWrapper, SDXLTEUnetWrapper
+from .wrapper import SD15Wrapper, SDXLWrapper, PixArtWrapper, TEHookCFG
+from .controlnet import ControlNetPlugin

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,
-}