From 7495502924905ea73adf37f3df53a90ab117c7bb Mon Sep 17 00:00:00 2001
From: smthemex <138738845+smthemex@users.noreply.github.com>
Date: Mon, 21 Apr 2025 17:16:55 +0800
Subject: [PATCH] I

---
 InstantCharacter/models/attn_processor.py | 164 ++++++++++
 InstantCharacter/models/norm_layer.py     |  46 +++
 InstantCharacter/models/resampler.py      | 365 ++++++++++++++++++++++
 InstantCharacter/models/utils.py          | 139 ++++++++
 4 files changed, 714 insertions(+)
 create mode 100644 InstantCharacter/models/attn_processor.py
 create mode 100644 InstantCharacter/models/norm_layer.py
 create mode 100644 InstantCharacter/models/resampler.py
 create mode 100644 InstantCharacter/models/utils.py

diff --git a/InstantCharacter/models/attn_processor.py b/InstantCharacter/models/attn_processor.py
new file mode 100644
index 0000000..8682bce
--- /dev/null
+++ b/InstantCharacter/models/attn_processor.py
@@ -0,0 +1,164 @@
+from typing import Optional
+
+import torch.nn as nn
+import torch
+import torch.nn.functional as F
+from diffusers.models.embeddings import apply_rotary_emb
+from einops import rearrange
+
+from .norm_layer import RMSNorm
+
+
+class FluxIPAttnProcessor(nn.Module):
+    """Attention processor used typically in processing the SD3-like self-attention projections."""
+
+    def __init__(
+        self,
+        hidden_size=None,
+        ip_hidden_states_dim=None,
+    ):
+        super().__init__()
+        self.norm_ip_q = RMSNorm(128, eps=1e-6)
+        self.to_k_ip = nn.Linear(ip_hidden_states_dim, hidden_size)
+        self.norm_ip_k = RMSNorm(128, eps=1e-6)
+        self.to_v_ip = nn.Linear(ip_hidden_states_dim, hidden_size)
+
+
+    def __call__(
+        self,
+        attn,
+        hidden_states: torch.FloatTensor,
+        encoder_hidden_states: torch.FloatTensor = None,
+        attention_mask: Optional[torch.FloatTensor] = None,
+        image_rotary_emb: Optional[torch.Tensor] = None,
+        emb_dict={},
+        subject_emb_dict={},
+        *args,
+        **kwargs,
+    ) -> torch.FloatTensor:
+        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
+
+        # `sample` projections.
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+
+        # IPadapter
+        ip_hidden_states = self._get_ip_hidden_states(
+            attn, 
+            query if encoder_hidden_states is not None else query[:, emb_dict['length_encoder_hidden_states']:],
+            subject_emb_dict.get('ip_hidden_states', None)
+        )
+
+        inner_dim = key.shape[-1]
+        head_dim = inner_dim // attn.heads
+
+        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+
+        if attn.norm_q is not None:
+            query = attn.norm_q(query)
+        if attn.norm_k is not None:
+            key = attn.norm_k(key)
+
+        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
+        if encoder_hidden_states is not None:
+            # `context` projections.
+            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
+            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
+            encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)
+
+            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
+                batch_size, -1, attn.heads, head_dim
+            ).transpose(1, 2)
+
+            if attn.norm_added_q is not None:
+                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
+            if attn.norm_added_k is not None:
+                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)
+
+            # attention
+            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
+            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
+            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)
+
+        if image_rotary_emb is not None:
+            query = apply_rotary_emb(query, image_rotary_emb)
+            key = apply_rotary_emb(key, image_rotary_emb)
+
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
+        hidden_states = hidden_states.to(query.dtype)
+
+
+        if encoder_hidden_states is not None:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1] :],
+            )
+                
+            if ip_hidden_states is not None:
+                hidden_states = hidden_states + ip_hidden_states * subject_emb_dict.get('scale', 1.0)
+
+            # linear proj
+            hidden_states = attn.to_out[0](hidden_states)
+            # dropout
+            hidden_states = attn.to_out[1](hidden_states)
+
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+            return hidden_states, encoder_hidden_states
+        else:
+
+            if ip_hidden_states is not None:
+                hidden_states[:, emb_dict['length_encoder_hidden_states']:] = \
+                    hidden_states[:, emb_dict['length_encoder_hidden_states']:] + \
+                    ip_hidden_states * subject_emb_dict.get('scale', 1.0)
+
+            return hidden_states
+
+
+    def _scaled_dot_product_attention(self, query, key, value, attention_mask=None, heads=None):
+        query = rearrange(query, '(b h) l c -> b h l c', h=heads)
+        key = rearrange(key, '(b h) l c -> b h l c', h=heads)
+        value = rearrange(value, '(b h) l c -> b h l c', h=heads)
+        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False, attn_mask=None)
+        hidden_states = rearrange(hidden_states, 'b h l c -> (b h) l c', h=heads)
+        hidden_states = hidden_states.to(query)
+        return hidden_states
+
+
+    def _get_ip_hidden_states(
+            self, 
+            attn, 
+            img_query, 
+            ip_hidden_states, 
+        ):
+        if ip_hidden_states is None:
+            return None
+        
+        if not hasattr(self, 'to_k_ip') or not hasattr(self, 'to_v_ip'):
+            return None
+
+        ip_query = self.norm_ip_q(rearrange(img_query, 'b l (h d) -> b h l d', h=attn.heads))
+        ip_query = rearrange(ip_query, 'b h l d -> (b h) l d')
+        ip_key = self.to_k_ip(ip_hidden_states)
+        ip_key = self.norm_ip_k(rearrange(ip_key, 'b l (h d) -> b h l d', h=attn.heads))
+        ip_key = rearrange(ip_key, 'b h l d -> (b h) l d')
+        ip_value = self.to_v_ip(ip_hidden_states)
+        ip_value = attn.head_to_batch_dim(ip_value)
+        ip_hidden_states = self._scaled_dot_product_attention(
+            ip_query.to(ip_value.dtype), ip_key.to(ip_value.dtype), ip_value, None, attn.heads)
+        ip_hidden_states = ip_hidden_states.to(img_query.dtype)
+        ip_hidden_states = attn.batch_to_head_dim(ip_hidden_states)
+        return ip_hidden_states
+
diff --git a/InstantCharacter/models/norm_layer.py b/InstantCharacter/models/norm_layer.py
new file mode 100644
index 0000000..ee9ff4d
--- /dev/null
+++ b/InstantCharacter/models/norm_layer.py
@@ -0,0 +1,46 @@
+import torch.nn as nn
+import torch
+
+class RMSNorm(nn.Module):
+    def __init__(self, d, p=-1., eps=1e-8, bias=False):
+        """
+            Root Mean Square Layer Normalization
+        :param d: model size
+        :param p: partial RMSNorm, valid value [0, 1], default -1.0 (disabled)
+        :param eps:  epsilon value, default 1e-8
+        :param bias: whether use bias term for RMSNorm, disabled by
+            default because RMSNorm doesn't enforce re-centering invariance.
+        """
+        super(RMSNorm, self).__init__()
+
+        self.eps = eps
+        self.d = d
+        self.p = p
+        self.bias = bias
+
+        self.scale = nn.Parameter(torch.ones(d))
+        self.register_parameter("scale", self.scale)
+
+        if self.bias:
+            self.offset = nn.Parameter(torch.zeros(d))
+            self.register_parameter("offset", self.offset)
+
+    def forward(self, x):
+        if self.p < 0. or self.p > 1.:
+            norm_x = x.norm(2, dim=-1, keepdim=True)
+            d_x = self.d
+        else:
+            partial_size = int(self.d * self.p)
+            partial_x, _ = torch.split(x, [partial_size, self.d - partial_size], dim=-1)
+
+            norm_x = partial_x.norm(2, dim=-1, keepdim=True)
+            d_x = partial_size
+
+        rms_x = norm_x * d_x ** (-1. / 2)
+        x_normed = x / (rms_x + self.eps)
+
+        if self.bias:
+            return self.scale * x_normed + self.offset
+
+        return self.scale * x_normed
+
diff --git a/InstantCharacter/models/resampler.py b/InstantCharacter/models/resampler.py
new file mode 100644
index 0000000..a37a4ac
--- /dev/null
+++ b/InstantCharacter/models/resampler.py
@@ -0,0 +1,365 @@
+import torch.nn as nn
+import torch
+import math
+
+from diffusers.models.transformers.transformer_2d import BasicTransformerBlock
+from diffusers.models.embeddings import Timesteps, TimestepEmbedding
+from timm.models.vision_transformer import Mlp
+
+from .norm_layer import RMSNorm
+
+
+# FFN
+def FeedForward(dim, mult=4):
+    inner_dim = int(dim * mult)
+    return nn.Sequential(
+        nn.LayerNorm(dim),
+        nn.Linear(dim, inner_dim, bias=False),
+        nn.GELU(),
+        nn.Linear(inner_dim, dim, bias=False),
+    )
+
+    
+def reshape_tensor(x, heads):
+    bs, length, width = x.shape
+    #(bs, length, width) --> (bs, length, n_heads, dim_per_head)
+    x = x.view(bs, length, heads, -1)
+    # (bs, length, n_heads, dim_per_head) --> (bs, n_heads, length, dim_per_head)
+    x = x.transpose(1, 2)
+    # (bs, n_heads, length, dim_per_head) --> (bs*n_heads, length, dim_per_head)
+    x = x.reshape(bs, heads, length, -1)
+    return x
+
+
+class PerceiverAttention(nn.Module):
+    def __init__(self, *, dim, dim_head=64, heads=8):
+        super().__init__()
+        self.scale = dim_head**-0.5
+        self.dim_head = dim_head
+        self.heads = heads
+        inner_dim = dim_head * heads
+
+        self.norm1 = nn.LayerNorm(dim)
+        self.norm2 = nn.LayerNorm(dim)
+
+        self.to_q = nn.Linear(dim, inner_dim, bias=False)
+        self.to_kv = nn.Linear(dim, inner_dim * 2, bias=False)
+        self.to_out = nn.Linear(inner_dim, dim, bias=False)
+
+
+    def forward(self, x, latents, shift=None, scale=None):
+        """
+        Args:
+            x (torch.Tensor): image features
+                shape (b, n1, D)
+            latent (torch.Tensor): latent features
+                shape (b, n2, D)
+        """
+        x = self.norm1(x)
+        latents = self.norm2(latents)
+
+        if shift is not None and scale is not None:
+            latents = latents * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+        
+        b, l, _ = latents.shape
+
+        q = self.to_q(latents)
+        kv_input = torch.cat((x, latents), dim=-2)
+        k, v = self.to_kv(kv_input).chunk(2, dim=-1)
+        
+        q = reshape_tensor(q, self.heads)
+        k = reshape_tensor(k, self.heads)
+        v = reshape_tensor(v, self.heads)
+
+        # attention
+        scale = 1 / math.sqrt(math.sqrt(self.dim_head))
+        weight = (q * scale) @ (k * scale).transpose(-2, -1) # More stable with f16 than dividing afterwards
+        weight = torch.softmax(weight.float(), dim=-1).type(weight.dtype)
+        out = weight @ v
+        
+        out = out.permute(0, 2, 1, 3).reshape(b, l, -1)
+
+        return self.to_out(out)
+
+
+class ReshapeExpandToken(nn.Module):
+    def __init__(self, expand_token, token_dim):
+        super().__init__()
+        self.expand_token = expand_token
+        self.token_dim = token_dim
+
+    def forward(self, x):
+        x = x.reshape(-1, self.expand_token, self.token_dim)
+        return x
+
+
+class TimeResampler(nn.Module):
+    def __init__(
+        self,
+        dim=1024,
+        depth=8,
+        dim_head=64,
+        heads=16,
+        num_queries=8,
+        embedding_dim=768,
+        output_dim=1024,
+        ff_mult=4,
+        timestep_in_dim=320,
+        timestep_flip_sin_to_cos=True,
+        timestep_freq_shift=0,
+        expand_token=None,
+        extra_dim=None,
+    ):
+        super().__init__()
+        
+        self.latents = nn.Parameter(torch.randn(1, num_queries, dim) / dim**0.5)
+
+        self.expand_token = expand_token is not None
+        if expand_token:
+            self.expand_proj = torch.nn.Sequential(
+                torch.nn.Linear(embedding_dim, embedding_dim * 2),
+                torch.nn.GELU(),
+                torch.nn.Linear(embedding_dim * 2, embedding_dim * expand_token),
+                ReshapeExpandToken(expand_token, embedding_dim),
+                RMSNorm(embedding_dim, eps=1e-8),
+            )
+
+        self.proj_in = nn.Linear(embedding_dim, dim)
+        
+        self.extra_feature = extra_dim is not None
+        if self.extra_feature:
+            self.proj_in_norm = RMSNorm(dim, eps=1e-8)
+            self.extra_proj_in = torch.nn.Sequential(
+                nn.Linear(extra_dim, dim),
+                RMSNorm(dim, eps=1e-8),
+            )
+
+        self.proj_out = nn.Linear(dim, output_dim)
+        self.norm_out = nn.LayerNorm(output_dim)
+        
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(
+                nn.ModuleList(
+                    [
+                        # msa
+                        PerceiverAttention(dim=dim, dim_head=dim_head, heads=heads),
+                        # ff
+                        FeedForward(dim=dim, mult=ff_mult),
+                        # adaLN
+                        nn.Sequential(nn.SiLU(), nn.Linear(dim, 4 * dim, bias=True))
+                    ]
+                )
+            )
+
+        # time
+        self.time_proj = Timesteps(timestep_in_dim, timestep_flip_sin_to_cos, timestep_freq_shift)
+        self.time_embedding = TimestepEmbedding(timestep_in_dim, dim, act_fn="silu")
+
+
+    def forward(self, x, timestep, need_temb=False, extra_feature=None):
+        timestep_emb = self.embedding_time(x, timestep)  # bs, dim
+
+        latents = self.latents.repeat(x.size(0), 1, 1)
+        
+        if self.expand_token:
+            x = self.expand_proj(x)
+
+        x = self.proj_in(x)
+
+        if self.extra_feature:
+            extra_feature = self.extra_proj_in(extra_feature)
+            x = self.proj_in_norm(x)
+            x = torch.cat([x, extra_feature], dim=1)
+            
+        x = x + timestep_emb[:, None]
+
+        for attn, ff, adaLN_modulation in self.layers:
+            shift_msa, scale_msa, shift_mlp, scale_mlp = adaLN_modulation(timestep_emb).chunk(4, dim=1)
+            latents = attn(x, latents, shift_msa, scale_msa) + latents
+
+            res = latents
+            for idx_ff in range(len(ff)):
+                layer_ff = ff[idx_ff]
+                latents = layer_ff(latents)
+                if idx_ff == 0 and isinstance(layer_ff, nn.LayerNorm):  # adaLN
+                    latents = latents * (1 + scale_mlp.unsqueeze(1)) + shift_mlp.unsqueeze(1)
+            latents = latents + res
+
+            # latents = ff(latents) + latents
+            
+        latents = self.proj_out(latents)
+        latents = self.norm_out(latents)
+
+        if need_temb:
+            return latents, timestep_emb
+        else:
+            return latents
+
+
+    def embedding_time(self, sample, timestep):
+
+        # 1. time
+        timesteps = timestep
+        if not torch.is_tensor(timesteps):
+            # TODO: this requires sync between CPU and GPU. So try to pass timesteps as tensors if you can
+            # This would be a good case for the `match` statement (Python 3.10+)
+            is_mps = sample.device.type == "mps"
+            if isinstance(timestep, float):
+                dtype = torch.float32 if is_mps else torch.float64
+            else:
+                dtype = torch.int32 if is_mps else torch.int64
+            timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)
+        elif len(timesteps.shape) == 0:
+            timesteps = timesteps[None].to(sample.device)
+
+        # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+        timesteps = timesteps.expand(sample.shape[0])
+
+        t_emb = self.time_proj(timesteps)
+
+        # timesteps does not contain any weights and will always return f32 tensors
+        # but time_embedding might actually be running in fp16. so we need to cast here.
+        # there might be better ways to encapsulate this.
+        t_emb = t_emb.to(dtype=sample.dtype)
+
+        emb = self.time_embedding(t_emb, None)
+        return emb
+
+
+class CrossLayerCrossScaleProjector(nn.Module):
+    def __init__(
+        self,
+        inner_dim=2688,
+        num_attention_heads=42,
+        attention_head_dim=64,
+        cross_attention_dim=2688,
+        num_layers=4,
+
+        # resampler
+        dim=1280,
+        depth=4,
+        dim_head=64,
+        heads=20,
+        num_queries=1024,
+        embedding_dim=1152 + 1536,
+        output_dim=4096,
+        ff_mult=4,
+        timestep_in_dim=320,
+        timestep_flip_sin_to_cos=True,
+        timestep_freq_shift=0,
+    ):
+        super().__init__()
+
+        self.cross_layer_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    num_attention_heads,
+                    attention_head_dim,
+                    dropout=0,
+                    cross_attention_dim=cross_attention_dim,
+                    activation_fn="geglu",
+                    num_embeds_ada_norm=None,
+                    attention_bias=False,
+                    only_cross_attention=False,
+                    double_self_attention=False,
+                    upcast_attention=False,
+                    norm_type='layer_norm',
+                    norm_elementwise_affine=True,
+                    norm_eps=1e-6,
+                    attention_type="default",
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.cross_scale_blocks = nn.ModuleList(
+            [
+                BasicTransformerBlock(
+                    inner_dim,
+                    num_attention_heads,
+                    attention_head_dim,
+                    dropout=0,
+                    cross_attention_dim=cross_attention_dim,
+                    activation_fn="geglu",
+                    num_embeds_ada_norm=None,
+                    attention_bias=False,
+                    only_cross_attention=False,
+                    double_self_attention=False,
+                    upcast_attention=False,
+                    norm_type='layer_norm',
+                    norm_elementwise_affine=True,
+                    norm_eps=1e-6,
+                    attention_type="default",
+                )
+                for _ in range(num_layers)
+            ]
+        )
+
+        self.proj = Mlp(
+            in_features=inner_dim, 
+            hidden_features=int(inner_dim*2), 
+            act_layer=lambda: nn.GELU(approximate="tanh"), 
+            drop=0
+        )
+
+        self.proj_cross_layer = Mlp(
+            in_features=inner_dim, 
+            hidden_features=int(inner_dim*2), 
+            act_layer=lambda: nn.GELU(approximate="tanh"), 
+            drop=0
+        )
+
+        self.proj_cross_scale = Mlp(
+            in_features=inner_dim, 
+            hidden_features=int(inner_dim*2), 
+            act_layer=lambda: nn.GELU(approximate="tanh"), 
+            drop=0
+        )
+
+        self.resampler = TimeResampler(
+            dim=dim,
+            depth=depth,
+            dim_head=dim_head,
+            heads=heads,
+            num_queries=num_queries,
+            embedding_dim=embedding_dim,
+            output_dim=output_dim,
+            ff_mult=ff_mult,
+            timestep_in_dim=timestep_in_dim,
+            timestep_flip_sin_to_cos=timestep_flip_sin_to_cos,
+            timestep_freq_shift=timestep_freq_shift,
+        )
+
+    def forward(self, low_res_shallow, low_res_deep, high_res_deep, timesteps, cross_attention_kwargs=None, need_temb=True):
+        '''
+            low_res_shallow [bs, 729*l, c]
+            low_res_deep    [bs, 729, c]
+            high_res_deep   [bs, 729*4, c]
+        '''
+
+        cross_layer_hidden_states = low_res_deep
+        for block in self.cross_layer_blocks:
+            cross_layer_hidden_states = block(
+                cross_layer_hidden_states,
+                encoder_hidden_states=low_res_shallow,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        cross_layer_hidden_states = self.proj_cross_layer(cross_layer_hidden_states)
+
+        cross_scale_hidden_states = low_res_deep
+        for block in self.cross_scale_blocks:
+            cross_scale_hidden_states = block(
+                cross_scale_hidden_states,
+                encoder_hidden_states=high_res_deep,
+                cross_attention_kwargs=cross_attention_kwargs,
+            )
+        cross_scale_hidden_states = self.proj_cross_scale(cross_scale_hidden_states)
+        
+        hidden_states = self.proj(low_res_deep) + cross_scale_hidden_states
+        hidden_states = torch.cat([hidden_states, cross_layer_hidden_states], dim=1)
+
+        hidden_states, timestep_emb = self.resampler(hidden_states, timesteps, need_temb=True)
+        return hidden_states, timestep_emb
+
diff --git a/InstantCharacter/models/utils.py b/InstantCharacter/models/utils.py
new file mode 100644
index 0000000..36d57ec
--- /dev/null
+++ b/InstantCharacter/models/utils.py
@@ -0,0 +1,139 @@
+from safetensors.torch import load_file
+import torch
+from tqdm import tqdm
+
+__all__ = [
+    'flux_load_lora'
+]
+
+
+def is_int(d):
+    try:
+        d = int(d)
+        return True
+    except Exception as e:
+        return False
+
+
+def flux_load_lora(self, lora_file, lora_weight=1.0):
+    device = self.transformer.device
+
+    # DiT 部分
+    state_dict, network_alphas = self.lora_state_dict(lora_file, return_alphas=True)
+    state_dict = {k:v.to(device) for k,v in state_dict.items()}
+    
+    model = self.transformer
+    keys = list(state_dict.keys())
+    keys = [k for k in keys if k.startswith('transformer.')]
+
+    for k_lora in tqdm(keys, total=len(keys), desc=f"loading lora in transformer ..."):
+        v_lora = state_dict[k_lora]
+
+        # 非 up 的都跳过
+        if '.lora_A.weight' in k_lora:
+            continue
+        if '.alpha' in k_lora:
+            continue
+
+        k_lora_name = k_lora.replace("transformer.", "")
+        k_lora_name = k_lora_name.replace(".lora_B.weight", "")
+        attr_name_list = k_lora_name.split('.')
+
+        cur_attr = model
+        latest_attr_name = ''
+        for idx in range(0, len(attr_name_list)):
+            attr_name = attr_name_list[idx]
+            if is_int(attr_name):
+                cur_attr = cur_attr[int(attr_name)]
+                latest_attr_name = ''
+            else:
+                try:
+                    if latest_attr_name != '':
+                        cur_attr = cur_attr.__getattr__(f"{latest_attr_name}.{attr_name}")
+                    else:
+                        cur_attr = cur_attr.__getattr__(attr_name)
+                    latest_attr_name = ''
+                except Exception as e:
+                    if latest_attr_name != '':
+                        latest_attr_name = f"{latest_attr_name}.{attr_name}"
+                    else:
+                        latest_attr_name = attr_name
+
+        up_w = v_lora
+        down_w = state_dict[k_lora.replace('.lora_B.weight', '.lora_A.weight')]
+
+        # 赋值
+        einsum_a = f"ijabcdefg"
+        einsum_b = f"jkabcdefg"
+        einsum_res = f"ikabcdefg"
+        length_shape = len(up_w.shape)
+        einsum_str = f"{einsum_a[:length_shape]},{einsum_b[:length_shape]}->{einsum_res[:length_shape]}"
+        dtype = cur_attr.weight.data.dtype
+        d_w = torch.einsum(einsum_str, up_w.to(torch.float32), down_w.to(torch.float32)).to(dtype)
+        cur_attr.weight.data = cur_attr.weight.data + d_w * lora_weight
+
+
+
+    # text encoder 部分
+    raw_state_dict = load_file(lora_file)
+    raw_state_dict = {k:v.to(device) for k,v in raw_state_dict.items()}
+
+    # text encoder
+    state_dict = {k:v for k,v in raw_state_dict.items() if 'lora_te1_' in k}
+    model = self.text_encoder
+    keys = list(state_dict.keys())
+    keys = [k for k in keys if k.startswith('lora_te1_')]
+
+    for k_lora in tqdm(keys, total=len(keys), desc=f"loading lora in text_encoder ..."):
+        v_lora = state_dict[k_lora]
+
+        # 非 up 的都跳过
+        if '.lora_down.weight' in k_lora:
+            continue
+        if '.alpha' in k_lora:
+            continue
+
+        k_lora_name = k_lora.replace("lora_te1_", "")
+        k_lora_name = k_lora_name.replace(".lora_up.weight", "")
+        attr_name_list = k_lora_name.split('_')
+
+        cur_attr = model
+        latest_attr_name = ''
+        for idx in range(0, len(attr_name_list)):
+            attr_name = attr_name_list[idx]
+            if is_int(attr_name):
+                cur_attr = cur_attr[int(attr_name)]
+                latest_attr_name = ''
+            else:
+                try:
+                    if latest_attr_name != '':
+                        cur_attr = cur_attr.__getattr__(f"{latest_attr_name}_{attr_name}")
+                    else:
+                        cur_attr = cur_attr.__getattr__(attr_name)
+                    latest_attr_name = ''
+                except Exception as e:
+                    if latest_attr_name != '':
+                        latest_attr_name = f"{latest_attr_name}_{attr_name}"
+                    else:
+                        latest_attr_name = attr_name
+
+        up_w = v_lora
+        down_w = state_dict[k_lora.replace('.lora_up.weight', '.lora_down.weight')]
+        
+        alpha = state_dict.get(k_lora.replace('.lora_up.weight', '.alpha'), None)
+        if alpha is None:
+            lora_scale = 1
+        else:
+            rank = up_w.shape[1]
+            lora_scale = alpha / rank
+        
+        # 赋值
+        einsum_a = f"ijabcdefg"
+        einsum_b = f"jkabcdefg"
+        einsum_res = f"ikabcdefg"
+        length_shape = len(up_w.shape)
+        einsum_str = f"{einsum_a[:length_shape]},{einsum_b[:length_shape]}->{einsum_res[:length_shape]}"
+        dtype = cur_attr.weight.data.dtype
+        d_w = torch.einsum(einsum_str, up_w.to(torch.float32), down_w.to(torch.float32)).to(dtype)
+        cur_attr.weight.data = cur_attr.weight.data + d_w * lora_scale * lora_weight
+