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

hcpdiff/easy/model/loader.py

@@ -1,10 +1,12 @@
 import torch
-from hcpdiff.ckpt_manager import DiffusersSD15Format, DiffusersSDXLFormat, DiffusersPixArtFormat, OfficialSD15Format, OfficialSDXLFormat
+from diffusers import StableDiffusionPipeline, StableDiffusionXLPipeline, FluxPipeline
 from rainbowneko.ckpt_manager import NekoLoader, LocalCkptSource
-from hcpdiff.utils import auto_tokenizer_cls, auto_text_encoder_cls, get_pipe_name
-from hcpdiff.models.wrapper import SDXLWrapper, SD15Wrapper, PixArtWrapper
-from hcpdiff.models.compose import SDXLTextEncoder
-from diffusers import  StableDiffusionPipeline, StableDiffusionXLPipeline
+
+from hcpdiff.ckpt_manager import DiffusersSD15Format, DiffusersSDXLFormat, DiffusersPixArtFormat, OfficialSD15Format, OfficialSDXLFormat, \
+    DiffusersFluxFormat, OneFileFluxFormat
+from hcpdiff.models.compose import SDXLTextEncoder, FluxTextEncoder
+from hcpdiff.models.wrapper import SDXLWrapper, SD15Wrapper, PixArtWrapper, FluxWrapper
+from hcpdiff.utils import auto_text_encoder_cls, get_pipe_name
 
 def SD15_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=None,
                      tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
@@ -20,7 +22,7 @@ def SD15_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=
             source=LocalCkptSource(),
         )
     models = loader.load(ckpt_path, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
-                          dtype=dtype, **kwargs)
+                         dtype=dtype, **kwargs)
     return models
 
 def SDXL_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=None,
@@ -37,17 +39,34 @@ def SDXL_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=
             source=LocalCkptSource(),
         )
     models = loader.load(ckpt_path, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
-                          dtype=dtype, **kwargs)
+                         dtype=dtype, **kwargs)
     return models
 
 def PixArt_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=None,
-                     tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+                       tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
     loader = NekoLoader(
         format=DiffusersPixArtFormat(),
         source=LocalCkptSource(),
     )
     models = loader.load(ckpt_path, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
-                          dtype=dtype, **kwargs)
+                         dtype=dtype, **kwargs)
+    return models
+
+def Flux_auto_loader(ckpt_path, denoiser=None, TE=None, vae=None, noise_sampler=None,
+                     tokenizer=None, revision=None, dtype=torch.float32, **kwargs):
+    try:
+        try_diffusers = FluxPipeline.load_config(ckpt_path)
+        loader = NekoLoader(
+            format=DiffusersFluxFormat(),
+            source=LocalCkptSource(),
+        )
+    except EnvironmentError:
+        loader = NekoLoader(
+            format=OneFileFluxFormat(),
+            source=LocalCkptSource(),
+        )
+    models = loader.load(ckpt_path, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
+                         dtype=dtype, **kwargs)
     return models
 
 def auto_load_wrapper(pretrained_model, denoiser=None, TE=None, vae=None, noise_sampler=None, tokenizer=None, revision=None,
@@ -62,6 +81,9 @@ def auto_load_wrapper(pretrained_model, denoiser=None, TE=None, vae=None, noise_
     if text_encoder_cls == SDXLTextEncoder:
         wrapper_cls = SDXLWrapper
         format = DiffusersSDXLFormat()
+    elif text_encoder_cls == FluxTextEncoder:
+        wrapper_cls = FluxWrapper
+        format = DiffusersFluxFormat()
     elif 'PixArt' in pipe_name:
         wrapper_cls = PixArtWrapper
         format = DiffusersPixArtFormat()
@@ -74,6 +96,6 @@ def auto_load_wrapper(pretrained_model, denoiser=None, TE=None, vae=None, noise_
         source=LocalCkptSource(),
     )
     models = loader.load(pretrained_model, denoiser=denoiser, TE=TE, vae=vae, noise_sampler=noise_sampler, tokenizer=tokenizer, revision=revision,
-                          dtype=dtype)
+                         dtype=dtype)
 
-    return wrapper_cls.build_from_pretrained(models, **kwargs)
+    return wrapper_cls.build_from_pretrained(models, **kwargs)

hcpdiff/easy/sampler.py

@@ -1,5 +1,5 @@
 from hcpdiff.diffusion.sampler import DiffusersSampler
-from diffusers import DPMSolverMultistepScheduler, DDIMScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler
+from diffusers import DPMSolverMultistepScheduler, DDIMScheduler, EulerDiscreteScheduler, EulerAncestralDiscreteScheduler, FlowMatchEulerDiscreteScheduler
 
 class Diffusers_SD:
     dpmpp_2m = DiffusersSampler(
@@ -43,4 +43,11 @@ class Diffusers_SD:
             beta_end=0.012,
             beta_schedule='scaled_linear',
         )
+    )
+
+    euler_flow = DiffusersSampler(
+        FlowMatchEulerDiscreteScheduler(
+            shift=3.0,
+            use_dynamic_shifting=True,
+        )
     )

hcpdiff/loss/__init__.py

@@ -1,4 +1,5 @@
-from .weighting import MinSNRWeight, SNRWeight, EDMWeight, LossWeight
-from .ssim import SSIMLoss, MS_SSIMLoss
+from .base import DiffusionLossContainer
+from .charbonnier import CharbonnierLoss
 from .gw import GWLoss
-from .base import DiffusionLossContainer
+from .ssim import SSIMLoss, MS_SSIMLoss
+from .weighting import MinSNRWeight, SNRWeight, EDMWeight, LossWeight, LossMapWeight

hcpdiff/loss/charbonnier.py

@@ -0,0 +1,17 @@
+import torch
+from torch import nn
+
+class CharbonnierLoss(nn.Module):
+    """Charbonnier Loss (L1)"""
+
+    def __init__(self, eps=1e-3, size_average=True):
+        super(CharbonnierLoss, self).__init__()
+        self.eps = eps
+        self.size_average = size_average
+
+    def forward(self, x, y):
+        diff = x - y
+        loss = torch.sqrt((diff * diff) + (self.eps*self.eps))
+        if self.size_average:
+            loss = loss.mean()
+        return loss

hcpdiff/loss/vlb.py

@@ -25,10 +25,10 @@ class VLBLoss(nn.Module):
         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)
+        true_logvar = sampler.sigma_scheduler.get_post_log_var(timesteps, ndim=input.ndim)
 
         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)
+        pred_logvar = sampler.sigma_scheduler.get_post_log_var(timesteps, ndim=input.ndim, 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)

hcpdiff/loss/weighting.py

@@ -1,9 +1,10 @@
+from rainbowneko.utils import add_dims
+from rainbowneko.train.loss import FullInputLoss
 from torch import nn
+from typing import Callable
 
-from .base import DiffusionLossContainer
-
-class LossWeight(nn.Module):
-    def __init__(self, loss: DiffusionLossContainer):
+class LossWeight(nn.Module, FullInputLoss):
+    def __init__(self, loss: Callable):
         super().__init__()
         self.loss = loss
 
@@ -21,12 +22,29 @@ class LossWeight(nn.Module):
         '''
         raise NotImplementedError
 
-    def forward(self, pred, inputs):
+    def forward(self, pred, inputs, _full_pred, _full_inputs):
         '''
         weight: [B,1,1,1] or [B,C,H,W]
         loss: [B,*,*,*]
         '''
-        return self.get_weight(pred, inputs)*self.loss(pred, inputs)
+        return self.get_weight(_full_pred, _full_inputs)*self.loss(pred, inputs)
+
+class LossMapWeight(LossWeight):
+    def __init__(self, loss: Callable, normalize: bool = False):
+        super().__init__(loss)
+        self.normalize = normalize
+
+    def get_weight(self, pred, inputs):
+        ndim = pred['model_pred'].ndim
+        loss_map = inputs['loss_map'].float()
+        if ndim == 4:
+            if self.normalize:
+                loss_map /= loss_map.mean(dim=(1,2), keepdim=True)
+            return loss_map.unsqueeze(1)
+        elif ndim == 3:
+            if self.normalize:
+                loss_map /= loss_map.mean(dim=1, keepdim=True)
+            return loss_map.unsqueeze(-1)
 
 class SNRWeight(LossWeight):
     def get_weight(self, pred, inputs):
@@ -42,10 +60,10 @@ class SNRWeight(LossWeight):
         else:
             raise ValueError(f"{self.__class__.__name__} is not support for target_type {target_type}")
 
-        return w_snr.view(-1, 1, 1, 1)
+        return add_dims(w_snr, pred['model_pred'].ndim-1)
 
 class MinSNRWeight(LossWeight):
-    def __init__(self, loss: DiffusionLossContainer, gamma: float = 1.):
+    def __init__(self, loss: Callable, gamma: float = 1.):
         super().__init__(loss)
         self.gamma = gamma
 
@@ -63,10 +81,10 @@ class MinSNRWeight(LossWeight):
         else:
             raise ValueError(f"{self.__class__.__name__} is not support for target_type {target_type}")
 
-        return w_snr.view(-1, 1, 1, 1)
+        return add_dims(w_snr, pred['model_pred'].ndim-1)
 
 class EDMWeight(LossWeight):
-    def __init__(self, loss: DiffusionLossContainer, gamma: float = 1.):
+    def __init__(self, loss: Callable, gamma: float = 1.):
         super().__init__(loss)
         self.gamma = gamma
 
@@ -81,4 +99,4 @@ class EDMWeight(LossWeight):
         else:
             raise ValueError(f"{self.__class__.__name__} is not support for target_type {target_type}")
 
-        return w_snr.view(-1, 1, 1, 1)
+        return add_dims(w_snr, pred['model_pred'].ndim-1)

hcpdiff/models/__init__.py

@@ -6,5 +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 SD15Wrapper, SDXLWrapper, PixArtWrapper, TEHookCFG
+from .wrapper import SD15Wrapper, SDXLWrapper, PixArtWrapper, TEHookCFG, FluxWrapper
 from .controlnet import ControlNetPlugin

hcpdiff/models/cfg_context.py

@@ -1,8 +1,10 @@
-import torch
-from einops import repeat
 import math
 from typing import Union, Callable
 
+import torch
+from einops import repeat
+from rainbowneko.utils import add_dims
+
 class CFGContext:
     def pre(self, noisy_latents, timesteps):
         return noisy_latents, timesteps
@@ -35,7 +37,7 @@ class DreamArtistPTContext(CFGContext):
                 pass
             else:
                 rate = self.cfg_func(rate)
-            rate = rate.view(-1, 1, 1, 1)
+            rate = add_dims(rate, model_pred.ndim-1)
         else:
             rate = 1
         model_pred = e_t_uncond+((self.cfg_high-self.cfg_low)*rate+self.cfg_low)*(e_t-e_t_uncond)

hcpdiff/models/compose/__init__.py

@@ -1,4 +1,5 @@
 from .compose_tokenizer import ComposeTokenizer
 from .compose_textencoder import ComposeTextEncoder
 from .compose_hook import ComposeTEEXHook, ComposeEmbPTHook
-from .sdxl_composer import SDXLTokenizer, SDXLTextEncoder
+from .sdxl import SDXLTokenizer, SDXLTextEncoder
+from .flux import FluxTokenizer, FluxTextEncoder

hcpdiff/models/compose/compose_hook.py

@@ -1,128 +1,129 @@
-import os
-from typing import Dict, Union, Tuple, List
+from pathlib import Path
+from typing import Dict, Union, Tuple
 
 import torch
-from loguru import logger
 from torch import nn
 
+from hcpdiff.utils.net_utils import load_emb
 from .compose_textencoder import ComposeTextEncoder
 from ..text_emb_ex import EmbeddingPTHook
 from ..textencoder_ex import TEEXHook
-from ...utils.net_utils import load_emb
-from ..container import ParameterGroup
 
 class ComposeEmbPTHook(nn.Module):
-    def __init__(self, hook_list: List[Tuple[str, EmbeddingPTHook]]):
+    def __init__(self, hooks: Dict[str, EmbeddingPTHook]):
         super().__init__()
-        self.hook_list = hook_list
-        self.emb_train = nn.ParameterList()
+        self.hooks = hooks
+        self.emb_train = nn.ParameterList() # [ParameterDict{model_name:Parameter, ...}, ...]
 
     @property
     def N_repeats(self):
-        return self.hook_list[0][1].N_repeats
+        return {name:hook.N_repeats for name, hook in self.hooks.items()}
 
     @N_repeats.setter
     def N_repeats(self, value):
-        for name, hook in self.hook_list:
-            hook.N_repeats = value
+        for name, hook in self.hooks.items():
+            if isinstance(value, int):
+                hook.N_repeats = value
+            else:
+                hook.N_repeats = value[name]
 
-    def add_emb(self, emb: nn.Parameter, token_id_list: List[int]):
-        emb_len = 0
+    def add_emb(self, emb: Dict[str, nn.Parameter], token_ids: Dict[str, int]):
         # Same word in different tokenizer may have different token_id
-        for (name, hook), token_id in zip(self.hook_list, token_id_list):
-            hook.add_emb(emb[:, emb_len:emb_len+hook.embedding_dim], token_id)
-            emb_len += hook.embedding_dim
+        for name, hook in self.hooks.items():
+            hook.add_emb(emb[name], token_ids[name])
 
     def remove(self):
-        for name, hook in self.hook_list:
+        for name, hook in self.hooks.items():
             hook.remove()
 
     @classmethod
-    def hook(cls, ex_words_emb: Dict[str, ParameterGroup], tokenizer, text_encoder, **kwargs):
+    def hook(cls, ex_words_emb: Dict[str, nn.ParameterDict], tokenizer, text_encoder, **kwargs):
         if isinstance(text_encoder, ComposeTextEncoder):
-            hook_list = []
+            hooks = {}
 
             emb_len = 0
-            for i, name in enumerate(tokenizer.tokenizer_names):
+            for name in tokenizer.tokenizer_names:
                 text_encoder_i = getattr(text_encoder, 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()}
+                ex_words_emb_i = {k:v[name] for k, v in ex_words_emb.items()}  # {word_name:Parameter, ...}
                 emb_len += embedding_dim
-                hook_list.append((name, EmbeddingPTHook.hook(ex_words_emb_i, tokenizer_i, text_encoder_i, **kwargs)))
+                hooks[name] = EmbeddingPTHook.hook(ex_words_emb_i, tokenizer_i, text_encoder_i, **kwargs)
 
-            return cls(hook_list)
+            return cls(hooks)
         else:
             return EmbeddingPTHook.hook(ex_words_emb, tokenizer, text_encoder, **kwargs)
 
     @classmethod
-    def hook_from_dir(cls, emb_dir, tokenizer, text_encoder, device='cuda:0', **kwargs) -> Union[
-        Tuple['ComposeEmbPTHook', Dict], Tuple[EmbeddingPTHook, Dict]]:
+    def hook_from_dir(cls, emb_dir: str | Path, tokenizer, text_encoder, device='cuda', **kwargs) -> (
+            Tuple['ComposeEmbPTHook', Dict[str, nn.ParameterDict]] | Tuple[EmbeddingPTHook, Dict[str, nn.Parameter]]):
+        emb_dir = Path(emb_dir) if emb_dir is not None else None
         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')}
-
-            # slice of nn.Parameter cannot return grad. Split the tensor
-            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
+            ex_words_emb = {}  # {word_name:{model_name:Tensor, ...}, ...}
+            if emb_dir is not None and emb_dir.exists():
+                for file in emb_dir.glob('*.pt'):
+                    emb = load_emb(file)  # {model_name:Tensor, ...}
+                    emb = nn.ParameterDict({name:nn.Parameter(emb_i.to(device), requires_grad=False) for name, emb_i in emb.items()})
+                    ex_words_emb[file.stem] = 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, **kwargs)
 
 class ComposeTEEXHook:
-    def __init__(self, tehook_list: List[Tuple[str, TEEXHook]], cat_dim=-1):
-        self.tehook_list = tehook_list
-        self.cat_dim = cat_dim
+    def __init__(self, tehooks: Dict[str, TEEXHook]):
+        self.tehooks = tehooks
 
     @property
     def N_repeats(self):
-        return self.tehook_list[0][1].N_repeats
+        return {name:tehook.N_repeats for name, tehook in self.tehooks.items()}
 
     @N_repeats.setter
-    def N_repeats(self, value):
-        for name, tehook in self.tehook_list:
-            tehook.N_repeats = value
+    def N_repeats(self, value: int | Dict[str, int]):
+        for name, tehook in self.tehooks.items():
+            if isinstance(value, int):
+                tehook.N_repeats = value
+            else:
+                tehook.N_repeats = value[name]
 
     @property
     def clip_skip(self):
-        return self.tehook_list[0][1].clip_skip
+        return {name:tehook.clip_skip for name, tehook in self.tehooks.items()}
 
     @clip_skip.setter
-    def clip_skip(self, value):
-        for name, tehook in self.tehook_list:
-            tehook.clip_skip = value
+    def clip_skip(self, value: int | Dict[str, int]):
+        for name, tehook in self.tehooks.items():
+            if isinstance(value, int):
+                tehook.clip_skip = value
+            else:
+                tehook.clip_skip = value[name]
 
     @property
     def clip_final_norm(self):
-        return self.tehook_list[0][1].clip_final_norm
+        return {name:tehook.clip_final_norm for name, tehook in self.tehooks.items()}
 
     @clip_final_norm.setter
-    def clip_final_norm(self, value: bool):
-        for name, tehook in self.tehook_list:
-            tehook.clip_final_norm = value
+    def clip_final_norm(self, value: bool | Dict[str, bool]):
+        for name, tehook in self.tehooks.items():
+            if isinstance(value, bool):
+                tehook.clip_final_norm = value
+            else:
+                tehook.clip_final_norm = value[name]
 
     @property
     def use_attention_mask(self):
-        return self.tehook_list[0][1].use_attention_mask
+        return {name:tehook.use_attention_mask for name, tehook in self.tehooks.items()}
 
     @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, attention_mask = list(zip(*emb_list))
-        return torch.cat(encoder_hidden_states, dim=self.cat_dim), pooled_output, attention_mask[0]
+    def use_attention_mask(self, value: bool | Dict[str, bool]):
+        for name, tehook in self.tehooks.items():
+            if isinstance(value, bool):
+                tehook.use_attention_mask = value
+            else:
+                tehook.use_attention_mask = value[name]
 
     def enable_xformers(self):
-        for name, tehook in self.tehook_list:
+        for name, tehook in self.tehooks.items():
             tehook.enable_xformers()
 
     @staticmethod
@@ -130,19 +131,20 @@ 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, use_attention_mask=False) -> Union[
+    def hook(cls, text_enc: nn.Module, tokenizer, N_repeats=1, 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), 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)
+            get_data = lambda name, data:data[name] if isinstance(data, dict) else data
+            tehooks = {name:TEEXHook.hook(getattr(text_enc, name), getattr(tokenizer, name), get_data(name, N_repeats), get_data(name, clip_skip),
+                                          get_data(name, clip_final_norm), use_attention_mask=get_data(name, use_attention_mask))
+                for name in tokenizer.tokenizer_names}
+            return cls(tehooks)
         else:
             # single text encoder
             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):
+    def hook_pipe(cls, pipe, N_repeats=1, clip_skip=0, clip_final_norm=True, use_attention_mask=False):
         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_textencoder.py

@@ -13,24 +13,24 @@ from typing import Dict, Optional, Union, Tuple, List
 
 import torch
 from torch import nn
-from transformers import CLIPTextModel, PreTrainedModel, PretrainedConfig
+from transformers import CLIPTextModel, PreTrainedModel, PretrainedConfig, AutoModel
 from transformers.modeling_outputs import BaseModelOutputWithPooling
+from rainbowneko.utils import BatchableDict
 
 class ComposeTextEncoder(PreTrainedModel):
-    def __init__(self, model_list: List[Tuple[str, CLIPTextModel]], cat_dim=-1, with_hook=True):
-        super().__init__(PretrainedConfig(**{name:model.config for name, model in model_list}))
-        self.cat_dim = cat_dim
+    def __init__(self, models: Dict[str, PreTrainedModel], with_hook=True):
+        super().__init__(PretrainedConfig(**{name:model.config for name, model in models.items()}))
         self.with_hook = with_hook
 
         self.model_names = []
-        for name, model in model_list:
+        for name, model in models.items():
             setattr(self, name, model)
             self.model_names.append(name)
 
-    def get_input_embeddings(self) -> List[nn.Module]:
-        return [getattr(self, name).get_input_embeddings() for name in self.model_names]
+    def get_input_embeddings(self) -> Dict[str, nn.Module]:
+        return {name: getattr(self, name).get_input_embeddings() for name in self.model_names}
 
-    def set_input_embeddings(self, value_dict: Dict[str, int]):
+    def set_input_embeddings(self, value_dict: Dict[str, torch.Tensor]):
         for name, value in value_dict.items():
             getattr(self, name).set_input_embeddings(value)
 
@@ -60,7 +60,7 @@ class ComposeTextEncoder(PreTrainedModel):
         >>> tokenizer_B = CLIPTokenizer.from_pretrained("openai/clip-vit-base-patch32")
         >>> tokenizer_bigG = CLIPTokenizer.from_pretrained("laion/CLIP-ViT-bigG-14-laion2B-39B-b160k")
 
-        >>> clip_model = MultiTextEncoder([('clip_B', clip_B), ('clip_bigG', clip_bigG)])
+        >>> clip_model = ComposeTextEncoder({'clip_B': clip_B, 'clip_bigG': clip_bigG})
 
         >>> inputs = {
         >>>     'clip_B':tokenizer_B(["a photo of a cat", "a photo of a dog"], padding=True, return_tensors="pt").input_ids
@@ -72,28 +72,42 @@ class ComposeTextEncoder(PreTrainedModel):
         >>> pooled_output = outputs.pooler_output  # pooled (EOS token) states
         ```"""
 
-        input_ids_list = input_ids.chunk(len(self.model_names),dim=-1)
+        def get_data(name, data):
+            if data is None:
+                return None
+            elif isinstance(data, (dict, BatchableDict)):
+                return data[name]
+            else:
+                return data
 
         if self.with_hook:
-            encoder_hidden_states_list, pooled_output_list = [], []
-            for name, input_ids in zip(self.model_names, input_ids_list):
-                encoder_hidden_states, pooled_output = getattr(self, name)(
-                    input_ids,  # get token for model self.{name}
-                    attention_mask=attention_mask,
-                    position_ids=position_ids,
-                    output_attentions=output_attentions,
-                    output_hidden_states=output_hidden_states,
-                    return_dict=True,
-                )
-                encoder_hidden_states_list.append(encoder_hidden_states)
-                pooled_output_list.append(pooled_output)
-            encoder_hidden_states = torch.cat(encoder_hidden_states_list, dim=self.cat_dim)
-            return encoder_hidden_states, pooled_output_list
+            encoder_hidden_states_dict, pooled_output_dict = {}, {}
+            for name in self.model_names:
+                if position_ids_i := get_data(name, position_ids) is None:
+                    encoder_hidden_states, pooled_output = getattr(self, name)(
+                        get_data(name, input_ids),  # get token for model self.{name}
+                        attention_mask=get_data(name, attention_mask),
+                        output_attentions=get_data(name, output_attentions),
+                        output_hidden_states=get_data(name, output_hidden_states),
+                        return_dict=True,
+                    )
+                else:
+                    encoder_hidden_states, pooled_output = getattr(self, name)(
+                        get_data(name, input_ids),  # get token for model self.{name}
+                        attention_mask=get_data(name, attention_mask),
+                        position_ids=position_ids_i,
+                        output_attentions=get_data(name, output_attentions),
+                        output_hidden_states=get_data(name, output_hidden_states),
+                        return_dict=True,
+                    )
+                encoder_hidden_states_dict[name] = encoder_hidden_states
+                pooled_output_dict[name] = pooled_output
+            return encoder_hidden_states_dict, pooled_output_dict
         else:
             return_dict = return_dict if return_dict is not None else self.config.use_return_dict
 
-            text_feat_list = {'last_hidden_state':[], 'pooler_output':[], 'hidden_states':[], 'attentions':[]}
-            for name, input_ids in zip(self.model_names, input_ids_list):
+            text_feat_list = {'last_hidden_state':{}, 'pooler_output':{}, 'hidden_states':{}, 'attentions':{}}
+            for name in self.model_names:
                 text_feat: BaseModelOutputWithPooling = getattr(self, name)(
                     input_ids,  # get token for model self.{name}
                     attention_mask=attention_mask,
@@ -102,31 +116,31 @@ class ComposeTextEncoder(PreTrainedModel):
                     output_hidden_states=output_hidden_states,
                     return_dict=True,
                 )
-                text_feat_list['last_hidden_state'].append(text_feat.last_hidden_state)
-                text_feat_list['pooler_output'].append(text_feat.pooler_output)
-                text_feat_list['hidden_states'].append(text_feat.hidden_states)
-                text_feat_list['attentions'].append(text_feat.attentions)
-
-            last_hidden_state = torch.cat(text_feat_list['last_hidden_state'], dim=self.cat_dim)
-            # pooler_output = torch.cat(text_feat_list['pooler_output'], dim=self.cat_dim)
-            pooler_output = text_feat_list['pooler_output']
-            if text_feat_list['hidden_states'][0] is None:
-                hidden_states = None
-            else:
-                hidden_states = [torch.cat(states, dim=self.cat_dim) for states in zip(*text_feat_list['hidden_states'])]
+                text_feat_list['last_hidden_state'][name] = text_feat.last_hidden_state
+                text_feat_list['pooler_output'][name] = text_feat.pooler_output
+                text_feat_list['hidden_states'][name] = text_feat.hidden_states
+                text_feat_list['attentions'][name] = text_feat.attentions
+
+            # last_hidden_state = torch.cat(text_feat_list['last_hidden_state'], dim=self.cat_dim)
+            # # pooler_output = torch.cat(text_feat_list['pooler_output'], dim=self.cat_dim)
+            # pooler_output = text_feat_list['pooler_output']
+            # if text_feat_list['hidden_states'][0] is None:
+            #     hidden_states = None
+            # else:
+            #     hidden_states = [torch.cat(states, dim=self.cat_dim) for states in zip(*text_feat_list['hidden_states'])]
 
             if return_dict:
                 return BaseModelOutputWithPooling(
-                    last_hidden_state=last_hidden_state,
-                    pooler_output=pooler_output,
-                    hidden_states=hidden_states,
+                    last_hidden_state=text_feat_list['last_hidden_state'],
+                    pooler_output=text_feat_list['pooler_output'],
+                    hidden_states=text_feat_list['hidden_states'],
                     attentions=text_feat_list['attentions'],
                 )
             else:
-                return (last_hidden_state, pooler_output)+hidden_states
+                return text_feat_list['last_hidden_state'], text_feat_list['pooler_output'], text_feat_list['hidden_states']
 
     @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: List[Tuple[str, str]], *args,
+    def from_pretrained(cls, pretrained_model_name_or_path: Dict[str, str], *args,
                         subfolder: Dict[str, str] = None, revision: str = None, **kwargs):
         r"""
             Examples: sdxl text encoder
@@ -138,6 +152,6 @@ class ComposeTextEncoder(PreTrainedModel):
             >>>     ], subfolder={'clip_B':'text_encoder', 'clip_bigG':'text_encoder_2'})
             ```
         """
-        clip_list = [(name, CLIPTextModel.from_pretrained(path, subfolder=subfolder[name], **kwargs)) for name, path in pretrained_model_name_or_path]
-        compose_model = cls(clip_list)
+        models = {name: AutoModel.from_pretrained(path, subfolder=subfolder[name], **kwargs) for name, path in pretrained_model_name_or_path.items()}
+        compose_model = cls(models)
         return compose_model