hcpdiff 2.3.1__py3-none-any.whl → 2.4__py3-none-any.whl

hcpdiff/models/compose/compose_tokenizer.py

@@ -14,19 +14,20 @@ from typing import Dict, Tuple, List
 import torch
 from transformers import AutoTokenizer, CLIPTokenizer, PreTrainedTokenizer, PretrainedConfig
 from transformers.tokenization_utils_base import BatchEncoding
+from rainbowneko.utils import BatchableDict
 
 class ComposeTokenizer(PreTrainedTokenizer):
-    def __init__(self, tokenizer_list: List[Tuple[str, CLIPTokenizer]], cat_dim=-1):
-        self.cat_dim = cat_dim
+    def __init__(self, tokenizers: Dict[str, CLIPTokenizer]):
 
         self.tokenizer_names = []
-        for name, tokenizer in tokenizer_list:
+        for name, tokenizer in tokenizers.items():
             setattr(self, name, tokenizer)
             self.tokenizer_names.append(name)
 
         super().__init__()
 
-        self.model_max_length = torch.tensor([tokenizer.model_max_length for name, tokenizer in tokenizer_list])
+        # self.model_max_length = torch.tensor([tokenizer.model_max_length for name, tokenizer in tokenizer_list])
+        self.model_max_length = {name: tokenizer.model_max_length for name, tokenizer in tokenizers.items()}
 
     @property
     def first_tokenizer(self):
@@ -57,15 +58,17 @@ class ComposeTokenizer(PreTrainedTokenizer):
         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)]
+        if isinstance(max_length, dict):
+            token_infos: Dict[str, BatchEncoding] = {name: getattr(self, name)(text, *args, max_length=max_length[name], **kwargs)
+                for name in self.tokenizer_names}
         else:
-            token_list: List[BatchEncoding] = [getattr(self, name)(text, *args, max_length=max_length, **kwargs) for name in self.tokenizer_names]
+            token_infos: Dict[str, BatchEncoding] = {name: getattr(self, name)(text, *args, max_length=max_length, **kwargs)
+                for name in self.tokenizer_names}
 
-        input_ids = torch.cat([token.input_ids for token in token_list], dim=-1)  # [N_tokenizer, N_token]
-        attention_mask = torch.cat([token.attention_mask for token in token_list], dim=-1)
-        return BatchEncoding({'input_ids':input_ids, 'attention_mask':attention_mask})
+        input_ids = BatchableDict({name: token.input_ids for name, token in token_infos.items()})  # [N_tokenizer, N_token]
+        attention_mask = BatchableDict({name: token.get('attention_mask', None) for name, token in token_infos.items()})
+        position_ids = BatchableDict({name: token.get('position_ids', None) for name, token in token_infos.items()})
+        return BatchEncoding({'input_ids':input_ids, 'attention_mask':attention_mask, 'position_ids':position_ids})
 
     @classmethod
     def from_pretrained(cls, pretrained_model_name_or_path: List[Tuple[str, str]], *args,
@@ -73,3 +76,37 @@ class ComposeTokenizer(PreTrainedTokenizer):
         tokenizer_list = [(name, AutoTokenizer.from_pretrained(path, subfolder=subfolder[name], **kwargs)) for name, path in pretrained_model_name_or_path]
         compose_tokenizer = cls(tokenizer_list)
         return compose_tokenizer
+
+    def __repr__(self):
+        return f'ComposeTokenizer(\n' + '\n'.join([f'  {name}: {repr(getattr(self, name))}' for name in self.tokenizer_names]) + ')'
+
+    @staticmethod
+    def tokenize_ex(tokenizer, *args, device='cpu', squeeze=False, **kwargs):
+        if isinstance(tokenizer, ComposeTokenizer):
+            max_length = {name: (tok := getattr(tokenizer, name)).model_max_length * getattr(tok, 'N_repeats', 1) for name in tokenizer.tokenizer_names}
+        else:
+            max_length = tokenizer.model_max_length * getattr(tokenizer, 'N_repeats', 1)
+
+        text_inputs = tokenizer(
+            *args,
+            max_length=max_length,
+            **kwargs
+        )
+
+        def proc_tensor(v):
+            if v is None:
+                return None
+            elif squeeze:
+                return v.squeeze().to(device)
+            else:
+                return v.to(device)
+
+        for k, v in text_inputs.items():
+            if torch.is_tensor(v):
+                text_inputs[k] = proc_tensor(v)
+            elif isinstance(v, (dict, BatchableDict)):
+                for name, vi in v.items():
+                    if torch.is_tensor(vi):
+                        v[name] = proc_tensor(vi)
+
+        return text_inputs

hcpdiff/models/compose/flux.py

@@ -0,0 +1,75 @@
+from .compose_textencoder import ComposeTextEncoder
+from .compose_tokenizer import ComposeTokenizer
+from transformers import CLIPTextModel, AutoTokenizer, CLIPTextModelWithProjection, T5EncoderModel
+from typing import Optional, Union, Tuple, Dict
+import torch
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+
+class T5EncoderModel_Align(T5EncoderModel):
+    # fxxk the transformers!
+    def forward(
+        self,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        head_mask: Optional[torch.FloatTensor] = None,
+        inputs_embeds: Optional[torch.FloatTensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple[torch.FloatTensor], BaseModelOutputWithPooling]:
+        text_outputs = super().forward(input_ids, attention_mask, head_mask, inputs_embeds, output_attentions, output_hidden_states, return_dict)
+        return BaseModelOutputWithPooling(
+            last_hidden_state=text_outputs.last_hidden_state,
+            pooler_output=None,
+            hidden_states=text_outputs.hidden_states,
+            attentions=text_outputs.attentions,
+        )
+
+class FluxTextEncoder(ComposeTextEncoder):
+    def forward(
+        self,
+        input_ids: Optional[Dict[str, torch.Tensor]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        output = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        if self.with_hook:
+            encoder_hidden_states_dict, pooled_output_dict = output
+            encoder_hidden_states = encoder_hidden_states_dict['T5']
+            pooled_output = pooled_output_dict['clip']
+        else:
+            last_hidden_state = output['last_hidden_state']['T5']
+            pooler_output = output['pooler_output']['clip']
+            attentions = output['attentions']['T5']
+            hidden_states = output['hidden_states']['T5']
+            return BaseModelOutputWithPooling(
+                last_hidden_state=last_hidden_state,
+                pooler_output=pooler_output,
+                hidden_states=hidden_states,
+                attentions=attentions,
+            )
+
+        return encoder_hidden_states, pooled_output
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, *args, subfolder=None, revision:str=None, **kwargs):
+        clip_L = CLIPTextModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder', **kwargs)
+        T5 = T5EncoderModel.from_pretrained(pretrained_model_name_or_path, subfolder='text_encoder_2', **kwargs)
+        return cls({'clip': clip_L, 'T5': T5})
+
+class FluxTokenizer(ComposeTokenizer):
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, *args, subfolder=None, revision:str=None, **kwargs):
+        clip_L = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer', **kwargs)
+        T5 = AutoTokenizer.from_pretrained(pretrained_model_name_or_path, subfolder='tokenizer_2', **kwargs)
+        return cls({'clip': clip_L, 'T5': T5})

hcpdiff/models/compose/sdxl.py

@@ -0,0 +1,86 @@
+from .compose_textencoder import ComposeTextEncoder
+from .compose_tokenizer import ComposeTokenizer
+from transformers import CLIPTextModel, AutoTokenizer, CLIPTextModelWithProjection
+from typing import Optional, Union, Tuple, Dict
+import torch
+from transformers.modeling_outputs import BaseModelOutputWithPooling
+
+class CLIPTextModelWithProjection_Align(CLIPTextModelWithProjection):
+    # fxxk the transformers!
+    def forward(
+        self,
+        input_ids: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        try: # old version of transformers
+            text_outputs = super().forward(input_ids, attention_mask, position_ids, output_attentions, output_hidden_states, return_dict)
+        except TypeError: # new version(like 4.53.1) of transformers removed 'return_dict'
+            text_outputs = super().forward(input_ids, attention_mask, position_ids, output_attentions, output_hidden_states)
+            
+        return BaseModelOutputWithPooling(
+            last_hidden_state=text_outputs.last_hidden_state,
+            pooler_output=text_outputs.text_embeds,
+            hidden_states=text_outputs.hidden_states,
+            attentions=text_outputs.attentions,
+        )
+
+class SDXLTextEncoder(ComposeTextEncoder):
+    def forward(
+        self,
+        input_ids: Optional[Dict[str, torch.Tensor]] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        position_ids: Optional[torch.Tensor] = None,
+        output_attentions: Optional[bool] = None,
+        output_hidden_states: Optional[bool] = None,
+        return_dict: Optional[bool] = None,
+    ) -> Union[Tuple, BaseModelOutputWithPooling]:
+        output = super().forward(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            position_ids=position_ids,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        if self.with_hook:
+            encoder_hidden_states_dict, pooled_output_dict = output
+            encoder_hidden_states = torch.cat([encoder_hidden_states_dict['clip_L'], encoder_hidden_states_dict['clip_bigG']], dim=-1)
+            pooled_output = pooled_output_dict['clip_bigG']
+        else:
+            last_hidden_state = torch.cat((output['last_hidden_state']['clip_L'], output['last_hidden_state']['clip_bigG']), dim=-1)
+            pooler_output = output['pooler_output']['clip_bigG']
+            attentions = output['attentions']['clip_bigG']
+            if output['hidden_states']['clip_L'] is None:
+                hidden_states = None
+            else:
+                hidden_states = [torch.cat(states, dim=self.cat_dim) for states in zip(output['hidden_states']['clip_L'], output['hidden_states']['clip_bigG'])]
+            return BaseModelOutputWithPooling(
+                last_hidden_state=last_hidden_state,
+                pooler_output=pooler_output,
+                hidden_states=hidden_states,
+                attentions=attentions,
+            )
+
+        return encoder_hidden_states, pooled_output
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, *args, subfolder=None, revision:str=None, **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_L': clip_L, 'clip_bigG': clip_bigG})
+
+class SDXLTokenizer(ComposeTokenizer):
+    def __call__(self, text, *args, max_length=None, **kwargs):
+        token_info = super().__call__(text, *args, max_length=max_length, **kwargs)
+        token_info['attention_mask'] = token_info['attention_mask']['clip_L']
+        return token_info
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path: str, *args, subfolder=None, revision:str=None, **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_L': clip_L, 'clip_bigG': clip_bigG})

hcpdiff/models/text_emb_ex.py

@@ -7,17 +7,18 @@ text_emb_ex.py
     :Created:     10/03/2023
     :Licence:     Apache-2.0
 """
+import os
+from pathlib import Path
 from typing import Tuple, Dict, Any
 
 import torch
-from torch import nn
-import os
-from rainbowneko import _share
-from einops import rearrange, repeat
 import torch.nn.functional as F
+from einops import rearrange, repeat
+from rainbowneko import _share
+from rainbowneko.models.plugin import SinglePluginBlock
+from torch import nn
 
 from ..utils.net_utils import load_emb
-from rainbowneko.models.plugin import SinglePluginBlock
 
 class EmbeddingPTHook(SinglePluginBlock):
     def __init__(self, token_embedding:nn.Embedding, N_word=75, N_repeats=3):
@@ -74,7 +75,7 @@ class EmbeddingPTHook(SinglePluginBlock):
         self.handle_pre.remove()
 
     @classmethod
-    def hook(cls, ex_words_emb, tokenizer, text_encoder, **kwargs):
+    def hook(cls, ex_words_emb:Dict[str, nn.Parameter], 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]
@@ -87,9 +88,12 @@ class EmbeddingPTHook(SinglePluginBlock):
         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')}
+    def hook_from_dir(cls, emb_dir:str|Path, tokenizer, text_encoder, device='cuda', **kwargs):
+        if emb_dir is None:
+            ex_words_emb = {}
+        else:
+            emb_dir = Path(emb_dir)
+            ex_words_emb = {file.stem: nn.Parameter(load_emb(file).to(device), requires_grad=False) for file in emb_dir.glob('*.pt')}
         return cls.hook(ex_words_emb, tokenizer, text_encoder, **kwargs), ex_words_emb
 
 class EmbeddingPTInterpHook(SinglePluginBlock):

hcpdiff/models/textencoder_ex.py

@@ -51,48 +51,18 @@ class TEEXHook:
         else:
             self.final_layer_norm = None
 
+    @property
+    def N_repeats(self):
+        return self.tokenizer.N_repeats
+
+    @N_repeats.setter
+    def N_repeats(self, value: int):
+        self.tokenizer.N_repeats = value
+
     @property
     def device(self):
         return self.text_enc.device
 
-    def encode_prompt_to_emb(self, prompt):
-        text_inputs = self.tokenizer(
-            prompt,
-            padding="max_length",
-            max_length=self.tokenizer.model_max_length*self.N_repeats,
-            truncation=True,
-            return_tensors="pt",
-        )
-        text_input_ids = text_inputs.input_ids
-        if self.use_attention_mask:
-            attention_mask = text_inputs.get('attention_mask', None)
-        else:
-            attention_mask = None
-        if attention_mask is not None:
-            attention_mask = attention_mask.to(self.device)
-        position_ids = text_inputs.get('position_ids', None)
-        if position_ids is not None:
-            position_ids = position_ids.to(self.device)
-
-        # 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):
         feat_re = rearrange(feat_in[0], 'b (r w) -> (b r) w', r=self.N_repeats)  # 使Attention mask的尺寸为N_word+2
         return (feat_re,) if len(feat_in) == 1 else (feat_re, *feat_in[1:])

hcpdiff/models/wrapper/__init__.py

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

hcpdiff/models/wrapper/flux.py

@@ -0,0 +1,75 @@
+import torch
+from diffusers import FluxTransformer2DModel, AutoencoderKL
+from einops import repeat, rearrange
+from hcpdiff.diffusion.sampler import BaseSampler
+from hcpdiff.utils import pad_attn_bias
+from rainbowneko.utils import add_dims
+
+from .sd import SD15Wrapper
+from .utils import TEHookCFG, SDXL_TEHookCFG
+from ..cfg_context import CFGContext
+
+class FluxWrapper(SD15Wrapper):
+    def __init__(self, denoiser: FluxTransformer2DModel, TE, vae: AutoencoderKL, noise_sampler: BaseSampler, tokenizer, min_attnmask=0,
+                 guidance=5.0, patch_size=2, 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, 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', 'neg_prompt -> neg_prompt_ids',
+            'neg_attn_mask -> neg_attn_mask', 'plugin_input -> plugin_input'))
+        self.guidance = guidance
+        self.patch_size = patch_size
+
+    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, H, W, prompt_ids, encoder_hidden_states, pooled_output, timesteps, attn_mask=None, plugin_input={}, **kwargs):
+        attn_mask = attn_mask['T5']
+        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)
+
+        img_ids = torch.zeros(H, W, 3)
+        img_ids[..., 1] = img_ids[..., 1]+torch.arange(H)[:, None]
+        img_ids[..., 2] = img_ids[..., 2]+torch.arange(W)[None, :]
+        img_ids = repeat(img_ids, "h w c -> b (h w) c", b=x_t.shape[0])
+
+        txt_ids = torch.zeros(x_t.shape[0], encoder_hidden_states.shape[1], 3)
+
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, attn_mask=attn_mask, img_ids=img_ids, txt_ids=txt_ids,
+                         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, self.guidance, pooled_output, encoder_hidden_states, img_ids=img_ids, txt_ids=txt_ids).sample
+        return model_pred
+
+    def model_forward(self, prompt_ids, image, attn_mask=None, neg_prompt_ids=None, neg_attn_mask=None, plugin_input={}, **kwargs):
+        # input prepare
+        x_0 = self.get_latents(image)
+        B, C, H, W = x_0.shape
+        x_0_patch = rearrange(x_0, "b c (h ph) (w pw) -> b (c ph pw) h w", ph=self.patch_size, pw=self.patch_size)
+        x_t, noise, timesteps = self.noise_sampler.add_noise_rand_t(x_0_patch)
+        x_t_in = x_t*add_dims(self.noise_sampler.sigma_scheduler.c_in(timesteps).to(dtype=x_t.dtype), x_t.ndim-1)
+        t_in = self.noise_sampler.sigma_scheduler.c_noise(timesteps)
+        x_t_in = rearrange(x_t_in, "b c h w -> b (h w) c")
+
+        if neg_prompt_ids:
+            prompt_ids = self.pn_cat(neg_prompt_ids, prompt_ids)
+            if neg_attn_mask:
+                attn_mask = self.pn_cat(neg_attn_mask, attn_mask)
+
+        # model forward
+        x_t_in, t_in = self.cfg_context.pre(x_t_in, t_in)
+        encoder_hidden_states, pooled_output = self.forward_TE(prompt_ids, t_in, attn_mask=attn_mask, plugin_input=plugin_input, **kwargs)
+        model_pred = self.forward_denoiser(x_t_in, H, W, prompt_ids, encoder_hidden_states, pooled_output, t_in, attn_mask=attn_mask,
+                                           plugin_input=plugin_input, **kwargs)
+        model_pred = rearrange(model_pred, "b (h w) (c ph pw) -> b c (h ph) (w pw)", ph=self.patch_size, pw=self.patch_size, h=H, w=W)
+        model_pred = self.cfg_context.post(model_pred)
+
+        return dict(model_pred=model_pred, noise=noise, timesteps=timesteps, x_0=x_0, x_t=x_t, noise_sampler=self.noise_sampler)

hcpdiff/models/wrapper/pixart.py

@@ -2,6 +2,16 @@ from .sd import SD15Wrapper
 from hcpdiff.utils import pad_attn_bias
 
 class PixArtWrapper(SD15Wrapper):
+    def forward_TE(self, prompt_ids, timesteps, attn_mask=None, plugin_input={}, **kwargs):
+        # T5Encoder do not need position_ids (It use relative position embedding for key and query)
+        input_all = dict(prompt_ids=prompt_ids, timesteps=timesteps, 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, 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, resolution=None, aspect_ratio=None,
                      plugin_input={}, **kwargs):
         if attn_mask is not None:
@@ -16,4 +26,6 @@ class PixArtWrapper(SD15Wrapper):
         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
+
+        # remove pred vars for pixart output (see DiT for more)
+        return model_pred.chunk(2, dim=1)[0]

hcpdiff/models/wrapper/sd.py

@@ -60,7 +60,10 @@ class SD15Wrapper(BaseWrapper):
         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
+                if shift_factor := getattr(self.vae.config, 'shift_factor', None) is not None:
+                    latents = (latents-shift_factor)*self.vae.config.scaling_factor
+                else:
+                    latents = latents*self.vae.config.scaling_factor
         else:
             latents = image  # Cached latents
         return latents
@@ -87,6 +90,12 @@ class SD15Wrapper(BaseWrapper):
         model_pred = self.denoiser(x_t, timesteps, encoder_hidden_states, encoder_attention_mask=attn_mask).sample  # Predict the noise residual
         return model_pred
 
+    def pn_cat(self, neg, pos, dim=0):
+        if isinstance(pos, dict): # ComposeTextEncoder
+            return {name:torch.cat([neg[name], pos_i], dim=dim) for name, pos_i in pos.items()}
+        else:
+            return torch.cat([neg, pos], dim=dim)
+
     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
@@ -96,11 +105,11 @@ class SD15Wrapper(BaseWrapper):
         t_in = self.noise_sampler.sigma_scheduler.c_noise(timesteps)
 
         if neg_prompt_ids:
-            prompt_ids = torch.cat([neg_prompt_ids, prompt_ids], dim=0)
+            prompt_ids = self.pn_cat(neg_prompt_ids, prompt_ids)
             if neg_attn_mask:
-                attn_mask = torch.cat([neg_attn_mask, attn_mask], dim=0)
+                attn_mask = self.pn_cat(neg_attn_mask, attn_mask)
             if neg_position_ids:
-                position_ids = torch.cat([neg_position_ids, position_ids], dim=0)
+                position_ids = self.pn_cat(neg_position_ids, position_ids)
 
         # model forward
         x_t_in, t_in = self.cfg_context.pre(x_t_in, t_in)
@@ -198,17 +207,17 @@ class SDXLWrapper(SD15Wrapper):
         t_in = self.noise_sampler.sigma_scheduler.c_noise(timesteps)
 
         if neg_prompt_ids:
-            prompt_ids = torch.cat([neg_prompt_ids, prompt_ids], dim=0)
+            prompt_ids = self.pn_cat(neg_prompt_ids, prompt_ids)
             if neg_attn_mask:
-                attn_mask = torch.cat([neg_attn_mask, attn_mask], dim=0)
+                attn_mask = self.pn_cat(neg_attn_mask, attn_mask)
             if neg_position_ids:
-                position_ids = torch.cat([neg_position_ids, position_ids], dim=0)
+                position_ids = self.pn_cat(neg_position_ids, position_ids)
 
         # model forward
         x_t_in, t_in = self.cfg_context.pre(x_t_in, t_in)
         encoder_hidden_states, pooled_output = self.forward_TE(prompt_ids, t_in, attn_mask=attn_mask, position_ids=position_ids,
                                                                plugin_input=plugin_input)
-        added_cond_kwargs = {"text_embeds":pooled_output[-1], "time_ids":crop_info}
+        added_cond_kwargs = {"text_embeds":pooled_output, "time_ids":crop_info}
         model_pred = self.forward_denoiser(x_t_in, prompt_ids, encoder_hidden_states, t_in, 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)

hcpdiff/parser/embpt.py

@@ -1,7 +1,7 @@
 from typing import Dict, Tuple, List
 from rainbowneko.utils import Path_Like
-from hcpdiff.models import EmbeddingPTHook
-from torch import Tensor
+from hcpdiff.models.compose import ComposeEmbPTHook
+from torch import Tensor, nn
 
 class CfgEmbPTParser:
     def __init__(self, emb_dir: Path_Like, cfg_pt: Dict[str, Dict], lr: float = 1e-5, weight_decay: float = 0):
@@ -11,22 +11,22 @@ class CfgEmbPTParser:
         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.embedding_hook, self.ex_words_emb = ComposeEmbPTHook.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]
+            word_emb: nn.Parameter | nn.ParameterDict = self.ex_words_emb[pt_name]
             train_pts[pt_name] = word_emb
-            word_emb.requires_grad = True
+            word_emb.requires_grad_(True)
             self.embedding_hook.emb_train.append(word_emb)
-            param_group = {'params':word_emb}
+            param_group = {'params':word_emb.parameters() if hasattr(word_emb, 'parameters') else [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
+        return train_params_emb, train_pts

hcpdiff/utils/net_utils.py

@@ -1,15 +1,16 @@
+import json
 import os
 from copy import deepcopy
+from functools import partial
 from typing import Optional, Union
 
 import torch
 from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION, Optimizer
+from huggingface_hub import hf_hub_download
 from torch import nn
 from torch.optim import lr_scheduler
 from transformers import PretrainedConfig, AutoTokenizer, T5EncoderModel, CLIPTextModel
-from functools import partial
-from huggingface_hub import hf_hub_download
-import json
+from transformers.models.auto.tokenization_auto import get_tokenizer_config
 
 dtype_dict = {'fp32':torch.float32, 'amp':torch.float32, 'fp16':torch.float16, 'bf16':torch.bfloat16}
 
@@ -91,19 +92,24 @@ def get_scheduler_with_name(
     return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, **scheduler_kwargs)
 
 def auto_tokenizer_cls(pretrained_model_name_or_path: str, revision: str = None):
-    from hcpdiff.models.compose import SDXLTokenizer
+    from hcpdiff.models.compose import SDXLTokenizer, FluxTokenizer
     try:
-        tokenizer = AutoTokenizer.from_pretrained(
-            pretrained_model_name_or_path, subfolder="tokenizer_2",
-            revision=revision, use_fast=False,
+        tokenizer_config = get_tokenizer_config(
+            pretrained_model_name_or_path,
+            subfolder="tokenizer_2",
+            revision=revision
         )
-        return SDXLTokenizer
+        class_name = tokenizer_config.get("tokenizer_class")
+        if class_name == 'T5Tokenizer' or class_name == 'T5TokenizerFast':
+            return FluxTokenizer
+        else:
+            return SDXLTokenizer
     except:
-        # not sdxl, only one tokenizer
+        # not composed, only one tokenizer
         return AutoTokenizer
 
 def auto_text_encoder_cls(pretrained_model_name_or_path: str, revision: str = None):
-    from hcpdiff.models.compose import SDXLTextEncoder
+    from hcpdiff.models.compose import SDXLTextEncoder, FluxTextEncoder
     try:
         text_encoder_config = PretrainedConfig.from_pretrained(
             pretrained_model_name_or_path,
@@ -112,7 +118,11 @@ def auto_text_encoder_cls(pretrained_model_name_or_path: str, revision: str = No
         )
         if text_encoder_config.architectures is None:
             raise ValueError()
-        return SDXLTextEncoder
+        model_class = text_encoder_config.architectures[0]
+        if model_class == "T5EncoderModel":
+            return FluxTextEncoder
+        else:
+            return SDXLTextEncoder
     except:
         text_encoder_config = PretrainedConfig.from_pretrained(
             pretrained_model_name_or_path,
@@ -248,4 +258,4 @@ def get_dtype(dtype):
     if isinstance(dtype, torch.dtype):
         return dtype
     else:
-        return dtype_dict.get(dtype, torch.float32)
+        return dtype_dict.get(dtype, torch.float32)

hcpdiff/workflow/__init__.py

@@ -1,4 +1,4 @@
-from .diffusion import InputFeederAction, MakeLatentAction, DenoiseAction, SampleAction, DiffusionStepAction, \
+from .diffusion import InputFeederAction, MakeLatentAction, SD15DenoiseAction, SDXLDenoiseAction, PixartDenoiseAction, FluxDenoiseAction, SampleAction, DiffusionStepAction, \
     X0PredAction, SeedAction, MakeTimestepsAction, PrepareDiffusionAction, time_iter, DiffusionActions
 from .text import TextEncodeAction, TextHookAction, AttnMultTextEncodeAction
 from .vae import EncodeAction, DecodeAction