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

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

hcpdiff/models/textencoder_ex.py

@@ -8,29 +8,53 @@ textencoder_ex.py
     :Licence:     Apache-2.0
 """
 
-from typing import Tuple, Optional, List
+from typing import Tuple, Optional
 
 import torch
 from einops import repeat, rearrange
 from einops.layers.torch import Rearrange
+from loguru import logger
 from torch import nn
+from transformers import CLIPTextModelWithProjection, T5EncoderModel
 from transformers.models.clip.modeling_clip import CLIPAttention
 
 class TEEXHook:
-    def __init__(self, text_enc: nn.Module, tokenizer, N_repeats=3, clip_skip=0, clip_final_norm=True, device='cuda', use_attention_mask=False):
+    def __init__(self, text_enc: nn.Module, tokenizer, N_repeats=1, clip_skip=0, clip_final_norm=True, use_attention_mask=False):
         self.text_enc = text_enc
         self.tokenizer = tokenizer
 
         self.N_repeats = N_repeats
         self.clip_skip = clip_skip
         self.clip_final_norm = clip_final_norm
-        self.device = device
-        self.attn_mult = None
         self.use_attention_mask = use_attention_mask
 
         text_enc.register_forward_hook(self.forward_hook)
         text_enc.register_forward_pre_hook(self.forward_hook_input)
 
+    def find_final_norm(self, text_enc: nn.Module):
+        for module in text_enc.modules():
+            if 'final_layer_norm' in module._modules:
+                logger.info(f'find final_layer_norm in {type(module)}')
+                return module.final_layer_norm
+
+        logger.info(f'final_layer_norm not found in {type(text_enc)}')
+        return None
+
+    @property
+    def clip_final_norm(self):
+        return self.final_layer_norm is not None
+
+    @clip_final_norm.setter
+    def clip_final_norm(self, value: bool):
+        if value:
+            self.final_layer_norm = self.find_final_norm(self.text_enc)
+        else:
+            self.final_layer_norm = None
+
+    @property
+    def device(self):
+        return self.text_enc.device
+
     def encode_prompt_to_emb(self, prompt):
         text_inputs = self.tokenizer(
             prompt,
@@ -50,12 +74,23 @@ class TEEXHook:
         if position_ids is not None:
             position_ids = position_ids.to(self.device)
 
-        prompt_embeds, pooled_output = self.text_enc(
-            text_input_ids.to(self.device),
-            attention_mask=attention_mask,
-            position_ids=position_ids,
-            output_hidden_states=True,
-        )
+        # align with sd-webui
+        if isinstance(self.text_enc, CLIPTextModelWithProjection):
+            self.text_enc.text_projection.weight.data = self.text_enc.text_projection.weight.data.t()
+
+        if isinstance(self.text_enc, T5EncoderModel):
+            prompt_embeds, pooled_output = self.text_enc(
+                text_input_ids.to(self.device),
+                attention_mask=attention_mask,
+                output_hidden_states=True,
+            )
+        else:
+            prompt_embeds, pooled_output = self.text_enc(
+                text_input_ids.to(self.device),
+                attention_mask=attention_mask,
+                position_ids=position_ids,
+                output_hidden_states=True,
+            )
         return prompt_embeds, pooled_output, attention_mask
 
     def forward_hook_input(self, host, feat_in):
@@ -64,13 +99,12 @@ class TEEXHook:
 
     def forward_hook(self, host, feat_in: Tuple[torch.Tensor], feat_out):
         encoder_hidden_states = feat_out['hidden_states'][-self.clip_skip-1]
-        if self.clip_final_norm:
-            encoder_hidden_states = self.text_enc.text_model.final_layer_norm(encoder_hidden_states)
+        if self.clip_final_norm and self.final_layer_norm is not None:
+            encoder_hidden_states = self.final_layer_norm(encoder_hidden_states)
         if self.text_enc.training and self.clip_skip>0:
             encoder_hidden_states = encoder_hidden_states+0*feat_out['last_hidden_state'].mean()  # avoid unused parameters, make gradient checkpointing happy
-
         encoder_hidden_states = rearrange(encoder_hidden_states, '(b r) ... -> b r ...', r=self.N_repeats)  # [B, N_repeat, N_word+2, N_emb]
-        pooled_output = feat_out.pooler_output
+        pooled_output = feat_out.get('pooler_output', feat_out.get('text_embeds', None))
         # TODO: may have better fusion method
         if pooled_output is not None:
             pooled_output = rearrange(pooled_output, '(b r) ... -> b r ...', r=self.N_repeats).mean(dim=1)
@@ -81,7 +115,7 @@ class TEEXHook:
         return encoder_hidden_states, pooled_output
 
     def pool_hidden_states(self, encoder_hidden_states, input_ids):
-        pooled_output = encoder_hidden_states[:, :, -1, :].mean(dim=1) # [B, N_emb]
+        pooled_output = encoder_hidden_states[:, :, -1, :].mean(dim=1)  # [B, N_emb]
         return pooled_output
 
     @staticmethod
@@ -147,9 +181,11 @@ class TEEXHook:
         layer.forward = forward
 
     @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):
-        return cls(text_enc, tokenizer, N_repeats=N_repeats, clip_skip=clip_skip, clip_final_norm=clip_final_norm, device=device, use_attention_mask=use_attention_mask)
+    def hook(cls, text_enc: nn.Module, tokenizer, N_repeats=3, clip_skip=0, clip_final_norm=True, use_attention_mask=False):
+        return cls(text_enc, tokenizer, N_repeats=N_repeats, clip_skip=clip_skip, clip_final_norm=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(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(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/wrapper/__init__.py

@@ -0,0 +1,3 @@
+from .sd import SD15Wrapper, SDXLWrapper
+from .pixart import PixArtWrapper
+from .utils import TEHookCFG, SD15_TEHookCFG, SDXL_TEHookCFG

hcpdiff/models/wrapper/pixart.py

@@ -0,0 +1,19 @@
+from .sd import SD15Wrapper
+from hcpdiff.utils import pad_attn_bias
+
+class PixArtWrapper(SD15Wrapper):
+    def forward_denoiser(self, x_t, prompt_ids, encoder_hidden_states, timesteps, attn_mask=None, position_ids=None, resolution=None, aspect_ratio=None,
+                     plugin_input={}, **kwargs):
+        if attn_mask is not None:
+            attn_mask[:, :self.min_attnmask] = 1
+            encoder_hidden_states, attn_mask = pad_attn_bias(encoder_hidden_states, attn_mask)
+
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask,
+                         encoder_hidden_states=encoder_hidden_states, **plugin_input)
+        if hasattr(self.denoiser, 'input_feeder'):
+            for feeder in self.denoiser.input_feeder:
+                feeder(input_all)
+        added_cond_kwargs = {"resolution":resolution, "aspect_ratio":aspect_ratio}
+        model_pred = self.denoiser(x_t, encoder_hidden_states, timesteps, encoder_attention_mask=attn_mask,
+                               added_cond_kwargs=added_cond_kwargs).sample  # Predict the noise residual
+        return model_pred

hcpdiff/models/wrapper/sd.py

@@ -0,0 +1,218 @@
+from contextlib import nullcontext
+from functools import partial
+from typing import Dict, Union
+
+import torch
+from diffusers import AutoencoderKL, UNet2DConditionModel
+from rainbowneko.models.wrapper import BaseWrapper
+from torch import Tensor
+from torch import nn
+
+from hcpdiff.diffusion.sampler import BaseSampler
+from hcpdiff.models import TEEXHook
+from hcpdiff.models.compose import ComposeTEEXHook
+from hcpdiff.utils import pad_attn_bias
+from .utils import TEHookCFG, SD15_TEHookCFG, SDXL_TEHookCFG
+from ..cfg_context import CFGContext
+
+class SD15Wrapper(BaseWrapper):
+    def __init__(self, denoiser: UNet2DConditionModel, TE, vae: AutoencoderKL, noise_sampler: BaseSampler, tokenizer, min_attnmask=0,
+                 pred_type='eps', TE_hook_cfg:TEHookCFG=SD15_TEHookCFG, cfg_context=CFGContext(), key_map_in=None, key_map_out=None):
+        super().__init__()
+        self.key_mapper_in = self.build_mapper(key_map_in, None, (
+            'prompt -> prompt_ids', 'image -> image', 'attn_mask -> attn_mask', 'position_ids -> position_ids', 'neg_prompt -> neg_prompt_ids',
+            'neg_attn_mask -> neg_attn_mask', 'neg_position_ids -> neg_position_ids', 'plugin_input -> plugin_input'))
+        self.key_mapper_out = self.build_mapper(key_map_out, None, None)
+
+        self.denoiser = denoiser
+        self.TE = TE
+        self.vae = vae
+        self.noise_sampler = noise_sampler
+        self.tokenizer = tokenizer
+        self.min_attnmask = min_attnmask
+
+        self.pred_type = pred_type
+
+        self.TE_hook_cfg = TEHookCFG.create(TE_hook_cfg)
+        self.cfg_context = cfg_context
+        self.tokenizer.N_repeats = self.TE_hook_cfg.tokenizer_repeats
+
+    def post_init(self):
+        self.make_TE_hook(self.TE_hook_cfg)
+
+        self.vae_trainable = False
+        if self.vae is not None:
+            for p in self.vae.parameters():
+                if p.requires_grad:
+                    self.vae_trainable = True
+                    break
+
+        self.TE_trainable = False
+        for p in self.TE.parameters():
+            if p.requires_grad:
+                self.TE_trainable = True
+                break
+
+    def make_TE_hook(self, TE_hook_cfg):
+        # Hook and extend text_encoder
+        self.text_enc_hook = TEEXHook.hook(self.TE, self.tokenizer, N_repeats=TE_hook_cfg.tokenizer_repeats,
+                                           clip_skip=TE_hook_cfg.clip_skip, clip_final_norm=TE_hook_cfg.clip_final_norm)
+
+    def get_latents(self, image: Tensor):
+        if image.shape[1] == 3:
+            with torch.no_grad() if self.vae_trainable else nullcontext():
+                latents = self.vae.encode(image.to(dtype=self.vae.dtype)).latent_dist.sample()
+                latents = latents*self.vae.config.scaling_factor
+        else:
+            latents = image  # Cached latents
+        return latents
+
+    def forward_TE(self, prompt_ids, timesteps, attn_mask=None, position_ids=None, plugin_input={}, **kwargs):
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask, **plugin_input)
+        if hasattr(self.TE, 'input_feeder'):
+            for feeder in self.TE.input_feeder:
+                feeder(input_all)
+        # Get the text embedding for conditioning
+        encoder_hidden_states = self.TE(prompt_ids, position_ids=position_ids, attention_mask=attn_mask, output_hidden_states=True)[0]
+        return encoder_hidden_states
+
+    def forward_denoiser(self, x_t, prompt_ids, encoder_hidden_states, timesteps, attn_mask=None, position_ids=None, plugin_input={}, **kwargs):
+        if attn_mask is not None:
+            attn_mask[:, :self.min_attnmask] = 1
+            encoder_hidden_states, attn_mask = pad_attn_bias(encoder_hidden_states, attn_mask)
+
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask,
+                         encoder_hidden_states=encoder_hidden_states, **plugin_input)
+        if hasattr(self.denoiser, 'input_feeder'):
+            for feeder in self.denoiser.input_feeder:
+                feeder(input_all)
+        model_pred = self.denoiser(x_t, timesteps, encoder_hidden_states, encoder_attention_mask=attn_mask).sample  # Predict the noise residual
+        return model_pred
+
+    def model_forward(self, prompt_ids, image, attn_mask=None, position_ids=None, neg_prompt_ids=None, neg_attn_mask=None, neg_position_ids=None,
+                      plugin_input={}, **kwargs):
+        # input prepare
+        x_0 = self.get_latents(image)
+        x_t, noise, sigma, timesteps = self.noise_sampler.add_noise_rand_t(x_0)
+        x_t_in = x_t*self.noise_sampler.c_in(sigma).to(dtype=x_t.dtype)
+
+        if neg_prompt_ids:
+            prompt_ids = torch.cat([neg_prompt_ids, prompt_ids], dim=0)
+            if neg_attn_mask:
+                attn_mask = torch.cat([neg_attn_mask, attn_mask], dim=0)
+            if neg_position_ids:
+                position_ids = torch.cat([neg_position_ids, position_ids], dim=0)
+
+        # model forward
+        x_t_in, timesteps = self.cfg_context.pre(x_t_in, timesteps)
+        encoder_hidden_states = self.forward_TE(prompt_ids, timesteps, attn_mask=attn_mask, position_ids=position_ids,
+                                                plugin_input=plugin_input, **kwargs)
+        model_pred = self.forward_denoiser(x_t_in, prompt_ids, encoder_hidden_states, timesteps, attn_mask=attn_mask, position_ids=position_ids,
+                                           plugin_input=plugin_input, **kwargs)
+        model_pred = self.cfg_context.post(model_pred)
+
+        return dict(model_pred=model_pred, noise=noise, sigma=sigma, timesteps=timesteps, x_0=x_0, x_t=x_t, pred_type=self.pred_type,
+                    noise_sampler=self.noise_sampler)
+
+    def forward(self, ds_name=None, **kwargs):
+        model_args, model_kwargs = self.get_map_data(self.key_mapper_in, kwargs, ds_name)
+        out = self.model_forward(*model_args, **model_kwargs)
+        return self.get_map_data(self.key_mapper_out, out, ds_name=ds_name)[1]
+
+    def enable_gradient_checkpointing(self):
+        def grad_ckpt_enable(m):
+            if getattr(m, 'gradient_checkpointing', False):
+                m.training = True
+
+        self.denoiser.enable_gradient_checkpointing()
+        if self.TE_trainable:
+            self.TE.gradient_checkpointing_enable()
+        self.apply(grad_ckpt_enable)
+
+    def enable_xformers(self):
+        self.denoiser.enable_xformers_memory_efficient_attention()
+
+    @property
+    def trainable_parameters(self):
+        return [p for p in self.parameters() if p.requires_grad]
+
+    @property
+    def trainable_models(self) -> Dict[str, nn.Module]:
+        return {'self':self}
+
+    def set_dtype(self, dtype, vae_dtype):
+        self.dtype = dtype
+        self.vae_dtype = vae_dtype
+        # Move vae and text_encoder to device and cast to weight_dtype
+        if self.vae is not None:
+            self.vae = self.vae.to(dtype=vae_dtype)
+        if not self.TE_trainable:
+            self.TE = self.TE.to(dtype=dtype)
+
+    @classmethod
+    def from_pretrained(cls, models: Union[partial, Dict[str, nn.Module]], **kwargs):
+        models = models() if isinstance(models, partial) else models
+        return cls(models['denoiser'], models['TE'], models['vae'], models['noise_sampler'], models['tokenizer'], **kwargs)
+
+class SDXLWrapper(SD15Wrapper):
+    def __init__(self, denoiser: UNet2DConditionModel, TE, vae: AutoencoderKL, noise_sampler: BaseSampler, tokenizer, min_attnmask=0,
+                 pred_type='eps', TE_hook_cfg:TEHookCFG=SDXL_TEHookCFG, cfg_context=CFGContext(), key_map_in=None, key_map_out=None):
+        super().__init__(denoiser, TE, vae, noise_sampler, tokenizer, min_attnmask, pred_type, TE_hook_cfg, cfg_context, key_map_in, key_map_out)
+        self.key_mapper_in = self.build_mapper(key_map_in, None, (
+            'prompt -> prompt_ids', 'image -> image', 'attn_mask -> attn_mask', 'position_ids -> position_ids', 'neg_prompt -> neg_prompt_ids',
+            'neg_attn_mask -> neg_attn_mask', 'neg_position_ids -> neg_position_ids', 'plugin_input -> plugin_input', 'coord -> crop_info'))
+
+    def make_TE_hook(self, TE_hook_cfg):
+        # Hook and extend text_encoder
+        self.text_enc_hook = ComposeTEEXHook.hook(self.TE, self.tokenizer, N_repeats=TE_hook_cfg.tokenizer_repeats,
+                                                  clip_skip=TE_hook_cfg.clip_skip, clip_final_norm=TE_hook_cfg.clip_final_norm)
+
+    def forward_TE(self, prompt_ids, timesteps, attn_mask=None, position_ids=None, plugin_input={}, **kwargs):
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask, **plugin_input)
+        if hasattr(self.TE, 'input_feeder'):
+            for feeder in self.TE.input_feeder:
+                feeder(input_all)
+        # Get the text embedding for conditioning
+        encoder_hidden_states, pooled_output = self.TE(prompt_ids, position_ids=position_ids, attention_mask=attn_mask, output_hidden_states=True)
+        return encoder_hidden_states, pooled_output
+
+    def forward_denoiser(self, x_t, prompt_ids, encoder_hidden_states, timesteps, added_cond_kwargs, attn_mask=None, position_ids=None,
+                         plugin_input={}, **kwargs):
+        if attn_mask is not None:
+            attn_mask[:, :self.min_attnmask] = 1
+            encoder_hidden_states, attn_mask = pad_attn_bias(encoder_hidden_states, attn_mask)
+
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, position_ids=position_ids, attn_mask=attn_mask,
+                         encoder_hidden_states=encoder_hidden_states, added_cond_kwargs=added_cond_kwargs, **plugin_input)
+        if hasattr(self.denoiser, 'input_feeder'):
+            for feeder in self.denoiser.input_feeder:
+                feeder(input_all)
+        model_pred = self.denoiser(x_t, timesteps, encoder_hidden_states, encoder_attention_mask=attn_mask,
+                                   added_cond_kwargs=added_cond_kwargs).sample  # Predict the noise residual
+        return model_pred
+
+    def model_forward(self, prompt_ids, image, attn_mask=None, position_ids=None, neg_prompt_ids=None, neg_attn_mask=None, neg_position_ids=None,
+                      crop_info=None, plugin_input={}):
+        # input prepare
+        x_0 = self.get_latents(image)
+        x_t, noise, sigma, timesteps = self.noise_sampler.add_noise_rand_t(x_0)
+        x_t_in = x_t*self.noise_sampler.c_in(sigma).to(dtype=x_t.dtype)
+
+        if neg_prompt_ids:
+            prompt_ids = torch.cat([neg_prompt_ids, prompt_ids], dim=0)
+            if neg_attn_mask:
+                attn_mask = torch.cat([neg_attn_mask, attn_mask], dim=0)
+            if neg_position_ids:
+                position_ids = torch.cat([neg_position_ids, position_ids], dim=0)
+
+        # model forward
+        x_t_in, timesteps = self.cfg_context.pre(x_t_in, timesteps)
+        encoder_hidden_states, pooled_output = self.forward_TE(prompt_ids, timesteps, attn_mask=attn_mask, position_ids=position_ids,
+                                                               plugin_input=plugin_input)
+        added_cond_kwargs = {"text_embeds":pooled_output[-1], "time_ids":crop_info}
+        model_pred = self.forward_denoiser(x_t_in, prompt_ids, encoder_hidden_states, timesteps, added_cond_kwargs=added_cond_kwargs,
+                                           attn_mask=attn_mask, position_ids=position_ids, plugin_input=plugin_input)
+        model_pred = self.cfg_context.post(model_pred)
+
+        return dict(model_pred=model_pred, noise=noise, sigma=sigma, timesteps=timesteps, x_0=x_0, x_t=x_t, pred_type=self.pred_type,
+                    noise_sampler=self.noise_sampler)

hcpdiff/models/wrapper/utils.py

@@ -0,0 +1,20 @@
+from dataclasses import dataclass
+from rainbowneko.utils import is_dict
+
+class TEHookCFG:
+    def __init__(self, tokenizer_repeats: int = 1, clip_skip: int = 0, clip_final_norm: bool = True):
+        self.tokenizer_repeats = tokenizer_repeats
+        self.clip_skip = clip_skip
+        self.clip_final_norm = clip_final_norm
+
+    @classmethod
+    def create(cls, cfg):
+        if is_dict(cfg):
+            return cls(**cfg)
+        elif isinstance(cfg, cls):
+            return cfg
+        else:
+            raise ValueError(f'Invalid TEHookCFG type: {type(cfg)}')
+        
+SD15_TEHookCFG = TEHookCFG()
+SDXL_TEHookCFG = TEHookCFG(clip_skip=1, clip_final_norm=False)

hcpdiff/parser/__init__.py

@@ -0,0 +1 @@
+from .embpt import CfgEmbPTParser

hcpdiff/parser/embpt.py

@@ -0,0 +1,32 @@
+from typing import Dict, Tuple, List
+from rainbowneko.utils import Path_Like
+from hcpdiff.models import EmbeddingPTHook
+from torch import Tensor
+
+class CfgEmbPTParser:
+    def __init__(self, emb_dir: Path_Like, cfg_pt: Dict[str, Dict], lr: float = 1e-5, weight_decay: float = 0):
+        self.emb_dir = emb_dir
+        self.cfg_pt = cfg_pt
+        self.lr = lr
+        self.weight_decay = weight_decay
+
+    def get_params_group(self, model) -> Tuple[List, Dict[str, Tensor]]:
+        self.embedding_hook, self.ex_words_emb = EmbeddingPTHook.hook_from_dir(
+            self.emb_dir, model.tokenizer, model.TE, N_repeats=model.tokenizer.N_repeats)
+        self.embedding_hook.requires_grad_(False)
+
+        train_params_emb = []
+        train_pts = {}
+        for pt_name, info in self.cfg_pt.items():
+            word_emb = self.ex_words_emb[pt_name]
+            train_pts[pt_name] = word_emb
+            word_emb.requires_grad = True
+            self.embedding_hook.emb_train.append(word_emb)
+            param_group = {'params':word_emb}
+            if 'lr' in info:
+                param_group['lr'] = info.lr
+            if 'weight_decay' in info:
+                param_group['weight_decay'] = info.weight_decay
+            train_params_emb.append(param_group)
+
+        return train_params_emb, train_pts

hcpdiff/tools/convert_caption_txt2json.py

@@ -2,7 +2,7 @@ import argparse
 import json
 import os
 
-from hcpdiff.utils.img_size_tool import types_support
+from rainbowneko.utils import types_support
 
 parser = argparse.ArgumentParser(description='Stable Diffusion Training')
 parser.add_argument('--data_root', type=str, default='')