Upload 5 files

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
+

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
+

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
+

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
+

pipeline.py

@@ -0,0 +1,550 @@
+# Copyright 2025 Tencent InstantX Team. All rights reserved.
+#
+
+from PIL import Image
+from einops import rearrange
+import torch
+from diffusers.pipelines.flux.pipeline_flux import *
+from transformers import SiglipVisionModel, SiglipImageProcessor, AutoModel, AutoImageProcessor
+
+from models.attn_processor import FluxIPAttnProcessor
+from models.resampler import CrossLayerCrossScaleProjector
+from models.utils import flux_load_lora
+
+
+# TODO
+EXAMPLE_DOC_STRING = """
+    Examples:
+        ```py
+        >>> import torch
+        >>> from diffusers import FluxPipeline
+
+        >>> pipe = FluxPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=torch.bfloat16)
+        >>> pipe.to("cuda")
+        >>> prompt = "A cat holding a sign that says hello world"
+        >>> # Depending on the variant being used, the pipeline call will slightly vary.
+        >>> # Refer to the pipeline documentation for more details.
+        >>> image = pipe(prompt, num_inference_steps=4, guidance_scale=0.0).images[0]
+        >>> image.save("flux.png")
+        ```
+"""
+
+
+class InstantCharacterFluxPipeline(FluxPipeline):
+
+
+    @torch.inference_mode()
+    def encode_siglip_image_emb(self, siglip_image, device, dtype):
+        siglip_image = siglip_image.to(device, dtype=dtype)
+        res = self.siglip_image_encoder(siglip_image, output_hidden_states=True)
+
+        siglip_image_embeds = res.last_hidden_state
+
+        siglip_image_shallow_embeds = torch.cat([res.hidden_states[i] for i in [7, 13, 26]], dim=1)
+        
+        return siglip_image_embeds, siglip_image_shallow_embeds
+
+
+    @torch.inference_mode()
+    def encode_dinov2_image_emb(self, dinov2_image, device, dtype):
+        dinov2_image = dinov2_image.to(device, dtype=dtype)
+        res = self.dino_image_encoder_2(dinov2_image, output_hidden_states=True)
+
+        dinov2_image_embeds = res.last_hidden_state[:, 1:]
+
+        dinov2_image_shallow_embeds = torch.cat([res.hidden_states[i][:, 1:] for i in [9, 19, 29]], dim=1)
+
+        return dinov2_image_embeds, dinov2_image_shallow_embeds
+
+
+    @torch.inference_mode()
+    def encode_image_emb(self, siglip_image, device, dtype):
+        object_image_pil = siglip_image
+        object_image_pil_low_res = [object_image_pil.resize((384, 384))]
+        object_image_pil_high_res = object_image_pil.resize((768, 768))
+        object_image_pil_high_res = [
+            object_image_pil_high_res.crop((0, 0, 384, 384)),
+            object_image_pil_high_res.crop((384, 0, 768, 384)),
+            object_image_pil_high_res.crop((0, 384, 384, 768)),
+            object_image_pil_high_res.crop((384, 384, 768, 768)),
+        ]
+        nb_split_image = len(object_image_pil_high_res)
+
+        siglip_image_embeds = self.encode_siglip_image_emb(
+            self.siglip_image_processor(images=object_image_pil_low_res, return_tensors="pt").pixel_values, 
+            device, 
+            dtype
+        )
+        dinov2_image_embeds = self.encode_dinov2_image_emb(
+            self.dino_image_processor_2(images=object_image_pil_low_res, return_tensors="pt").pixel_values, 
+            device, 
+            dtype
+        )
+
+        image_embeds_low_res_deep = torch.cat([siglip_image_embeds[0], dinov2_image_embeds[0]], dim=2)
+        image_embeds_low_res_shallow = torch.cat([siglip_image_embeds[1], dinov2_image_embeds[1]], dim=2)
+
+        siglip_image_high_res = self.siglip_image_processor(images=object_image_pil_high_res, return_tensors="pt").pixel_values
+        siglip_image_high_res = siglip_image_high_res[None]
+        siglip_image_high_res = rearrange(siglip_image_high_res, 'b n c h w -> (b n) c h w')
+        siglip_image_high_res_embeds = self.encode_siglip_image_emb(siglip_image_high_res, device, dtype)
+        siglip_image_high_res_deep = rearrange(siglip_image_high_res_embeds[0], '(b n) l c -> b (n l) c', n=nb_split_image)
+        dinov2_image_high_res = self.dino_image_processor_2(images=object_image_pil_high_res, return_tensors="pt").pixel_values
+        dinov2_image_high_res = dinov2_image_high_res[None]
+        dinov2_image_high_res = rearrange(dinov2_image_high_res, 'b n c h w -> (b n) c h w')
+        dinov2_image_high_res_embeds = self.encode_dinov2_image_emb(dinov2_image_high_res, device, dtype)
+        dinov2_image_high_res_deep = rearrange(dinov2_image_high_res_embeds[0], '(b n) l c -> b (n l) c', n=nb_split_image)
+        image_embeds_high_res_deep = torch.cat([siglip_image_high_res_deep, dinov2_image_high_res_deep], dim=2)
+
+        image_embeds_dict = dict(
+            image_embeds_low_res_shallow=image_embeds_low_res_shallow,
+            image_embeds_low_res_deep=image_embeds_low_res_deep,
+            image_embeds_high_res_deep=image_embeds_high_res_deep,
+        )
+        return image_embeds_dict
+
+
+    @torch.inference_mode()
+    def init_ccp_and_attn_processor(self, *args, **kwargs):
+        subject_ip_adapter_path = kwargs['subject_ip_adapter_path']
+        nb_token = kwargs['nb_token']
+        state_dict = torch.load(subject_ip_adapter_path, map_location="cpu")
+        device, dtype = self.transformer.device, self.transformer.dtype
+
+        print(f"=> init attn processor")
+        attn_procs = {}
+        for idx_attn, (name, v) in enumerate(self.transformer.attn_processors.items()):
+            attn_procs[name] = FluxIPAttnProcessor(
+                hidden_size=self.transformer.config.attention_head_dim * self.transformer.config.num_attention_heads,
+                ip_hidden_states_dim=self.text_encoder_2.config.d_model,
+            ).to(device, dtype=dtype)
+        self.transformer.set_attn_processor(attn_procs)
+        tmp_ip_layers = torch.nn.ModuleList(self.transformer.attn_processors.values())
+        key_name = tmp_ip_layers.load_state_dict(state_dict["ip_adapter"], strict=False)
+        print(f"=> load attn processor: {key_name}")
+
+        print(f"=> init project")
+        image_proj_model = CrossLayerCrossScaleProjector(
+            inner_dim=1152 + 1536,
+            num_attention_heads=42,
+            attention_head_dim=64,
+            cross_attention_dim=1152 + 1536,
+            num_layers=4,
+            dim=1280,
+            depth=4,
+            dim_head=64,
+            heads=20,
+            num_queries=nb_token,
+            embedding_dim=1152 + 1536,
+            output_dim=4096,
+            ff_mult=4,
+            timestep_in_dim=320,
+            timestep_flip_sin_to_cos=True,
+            timestep_freq_shift=0,
+        )
+        image_proj_model.eval()
+        image_proj_model.to(device, dtype=dtype)
+
+        key_name = image_proj_model.load_state_dict(state_dict["image_proj"], strict=False)
+        print(f"=> load project: {key_name}")
+        self.subject_image_proj_model = image_proj_model
+
+
+    @torch.inference_mode()
+    def init_adapter(
+        self, 
+        image_encoder_path=None, 
+        image_encoder_2_path=None, 
+        subject_ipadapter_cfg=None, 
+    ):
+        device, dtype = self.transformer.device, self.transformer.dtype
+
+        # image encoder
+        print(f"=> loading image_encoder_1: {image_encoder_path}")
+        image_encoder = SiglipVisionModel.from_pretrained(image_encoder_path)
+        image_processor = SiglipImageProcessor.from_pretrained(image_encoder_path)
+        image_encoder.eval()
+        image_encoder.to(device, dtype=dtype)
+        self.siglip_image_encoder = image_encoder
+        self.siglip_image_processor = image_processor
+
+        # image encoder 2
+        print(f"=> loading image_encoder_2: {image_encoder_2_path}")
+        image_encoder_2 = AutoModel.from_pretrained(image_encoder_2_path)
+        image_processor_2 = AutoImageProcessor.from_pretrained(image_encoder_2_path)
+        image_encoder_2.eval()
+        image_encoder_2.to(device, dtype=dtype)
+        image_processor_2.crop_size = dict(height=384, width=384)
+        image_processor_2.size = dict(shortest_edge=384)
+        self.dino_image_encoder_2 = image_encoder_2
+        self.dino_image_processor_2 = image_processor_2
+
+        # ccp and adapter
+        self.init_ccp_and_attn_processor(**subject_ipadapter_cfg)
+
+
+    @torch.no_grad()
+    @replace_example_docstring(EXAMPLE_DOC_STRING)
+    def __call__(
+        self,
+        prompt: Union[str, List[str]] = None,
+        prompt_2: Optional[Union[str, List[str]]] = None,
+        negative_prompt: Union[str, List[str]] = None,
+        negative_prompt_2: Optional[Union[str, List[str]]] = None,
+        true_cfg_scale: float = 1.0,
+        height: Optional[int] = None,
+        width: Optional[int] = None,
+        num_inference_steps: int = 28,
+        sigmas: Optional[List[float]] = None,
+        guidance_scale: float = 3.5,
+        num_images_per_prompt: Optional[int] = 1,
+        generator: Optional[Union[torch.Generator, List[torch.Generator]]] = None,
+        latents: Optional[torch.FloatTensor] = None,
+        prompt_embeds: Optional[torch.FloatTensor] = None,
+        pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        ip_adapter_image: Optional[PipelineImageInput] = None,
+        ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        negative_ip_adapter_image: Optional[PipelineImageInput] = None,
+        negative_ip_adapter_image_embeds: Optional[List[torch.Tensor]] = None,
+        negative_prompt_embeds: Optional[torch.FloatTensor] = None,
+        negative_pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+        output_type: Optional[str] = "pil",
+        return_dict: bool = True,
+        joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+        callback_on_step_end: Optional[Callable[[int, int, Dict], None]] = None,
+        callback_on_step_end_tensor_inputs: List[str] = ["latents"],
+        max_sequence_length: int = 512,
+        subject_image: Image.Image = None,
+        subject_scale: float = 0.8,
+
+    ):
+        r"""
+        Function invoked when calling the pipeline for generation.
+
+        Args:
+            prompt (`str` or `List[str]`, *optional*):
+                The prompt or prompts to guide the image generation. If not defined, one has to pass `prompt_embeds`.
+                instead.
+            prompt_2 (`str` or `List[str]`, *optional*):
+                The prompt or prompts to be sent to `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+                will be used instead
+            height (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The height in pixels of the generated image. This is set to 1024 by default for the best results.
+            width (`int`, *optional*, defaults to self.unet.config.sample_size * self.vae_scale_factor):
+                The width in pixels of the generated image. This is set to 1024 by default for the best results.
+            num_inference_steps (`int`, *optional*, defaults to 50):
+                The number of denoising steps. More denoising steps usually lead to a higher quality image at the
+                expense of slower inference.
+            sigmas (`List[float]`, *optional*):
+                Custom sigmas to use for the denoising process with schedulers which support a `sigmas` argument in
+                their `set_timesteps` method. If not defined, the default behavior when `num_inference_steps` is passed
+                will be used.
+            guidance_scale (`float`, *optional*, defaults to 7.0):
+                Guidance scale as defined in [Classifier-Free Diffusion Guidance](https://arxiv.org/abs/2207.12598).
+                `guidance_scale` is defined as `w` of equation 2. of [Imagen
+                Paper](https://arxiv.org/pdf/2205.11487.pdf). Guidance scale is enabled by setting `guidance_scale >
+                1`. Higher guidance scale encourages to generate images that are closely linked to the text `prompt`,
+                usually at the expense of lower image quality.
+            num_images_per_prompt (`int`, *optional*, defaults to 1):
+                The number of images to generate per prompt.
+            generator (`torch.Generator` or `List[torch.Generator]`, *optional*):
+                One or a list of [torch generator(s)](https://pytorch.org/docs/stable/generated/torch.Generator.html)
+                to make generation deterministic.
+            latents (`torch.FloatTensor`, *optional*):
+                Pre-generated noisy latents, sampled from a Gaussian distribution, to be used as inputs for image
+                generation. Can be used to tweak the same generation with different prompts. If not provided, a latents
+                tensor will ge generated by sampling using the supplied random `generator`.
+            prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting. If not
+                provided, text embeddings will be generated from `prompt` input argument.
+            pooled_prompt_embeds (`torch.FloatTensor`, *optional*):
+                Pre-generated pooled text embeddings. Can be used to easily tweak text inputs, *e.g.* prompt weighting.
+                If not provided, pooled text embeddings will be generated from `prompt` input argument.
+            ip_adapter_image: (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            negative_ip_adapter_image:
+                (`PipelineImageInput`, *optional*): Optional image input to work with IP Adapters.
+            negative_ip_adapter_image_embeds (`List[torch.Tensor]`, *optional*):
+                Pre-generated image embeddings for IP-Adapter. It should be a list of length same as number of
+                IP-adapters. Each element should be a tensor of shape `(batch_size, num_images, emb_dim)`. If not
+                provided, embeddings are computed from the `ip_adapter_image` input argument.
+            output_type (`str`, *optional*, defaults to `"pil"`):
+                The output format of the generate image. Choose between
+                [PIL](https://pillow.readthedocs.io/en/stable/): `PIL.Image.Image` or `np.array`.
+            return_dict (`bool`, *optional*, defaults to `True`):
+                Whether or not to return a [`~pipelines.flux.FluxPipelineOutput`] instead of a plain tuple.
+            joint_attention_kwargs (`dict`, *optional*):
+                A kwargs dictionary that if specified is passed along to the `AttentionProcessor` as defined under
+                `self.processor` in
+                [diffusers.models.attention_processor](https://github.com/huggingface/diffusers/blob/main/src/diffusers/models/attention_processor.py).
+            callback_on_step_end (`Callable`, *optional*):
+                A function that calls at the end of each denoising steps during the inference. The function is called
+                with the following arguments: `callback_on_step_end(self: DiffusionPipeline, step: int, timestep: int,
+                callback_kwargs: Dict)`. `callback_kwargs` will include a list of all tensors as specified by
+                `callback_on_step_end_tensor_inputs`.
+            callback_on_step_end_tensor_inputs (`List`, *optional*):
+                The list of tensor inputs for the `callback_on_step_end` function. The tensors specified in the list
+                will be passed as `callback_kwargs` argument. You will only be able to include variables listed in the
+                `._callback_tensor_inputs` attribute of your pipeline class.
+            max_sequence_length (`int` defaults to 512): Maximum sequence length to use with the `prompt`.
+
+        Examples:
+
+        Returns:
+            [`~pipelines.flux.FluxPipelineOutput`] or `tuple`: [`~pipelines.flux.FluxPipelineOutput`] if `return_dict`
+            is True, otherwise a `tuple`. When returning a tuple, the first element is a list with the generated
+            images.
+        """
+
+        height = height or self.default_sample_size * self.vae_scale_factor
+        width = width or self.default_sample_size * self.vae_scale_factor
+
+        # 1. Check inputs. Raise error if not correct
+        self.check_inputs(
+            prompt,
+            prompt_2,
+            height,
+            width,
+            negative_prompt=negative_prompt,
+            negative_prompt_2=negative_prompt_2,
+            prompt_embeds=prompt_embeds,
+            negative_prompt_embeds=negative_prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            negative_pooled_prompt_embeds=negative_pooled_prompt_embeds,
+            callback_on_step_end_tensor_inputs=callback_on_step_end_tensor_inputs,
+            max_sequence_length=max_sequence_length,
+        )
+
+        self._guidance_scale = guidance_scale
+        self._joint_attention_kwargs = joint_attention_kwargs
+        self._interrupt = False
+
+        # 2. Define call parameters
+        if prompt is not None and isinstance(prompt, str):
+            batch_size = 1
+        elif prompt is not None and isinstance(prompt, list):
+            batch_size = len(prompt)
+        else:
+            batch_size = prompt_embeds.shape[0]
+
+        device = self._execution_device
+        dtype = self.transformer.dtype
+
+        lora_scale = (
+            self.joint_attention_kwargs.get("scale", None) if self.joint_attention_kwargs is not None else None
+        )
+        do_true_cfg = true_cfg_scale > 1 and negative_prompt is not None
+        (
+            prompt_embeds,
+            pooled_prompt_embeds,
+            text_ids,
+        ) = self.encode_prompt(
+            prompt=prompt,
+            prompt_2=prompt_2,
+            prompt_embeds=prompt_embeds,
+            pooled_prompt_embeds=pooled_prompt_embeds,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+            max_sequence_length=max_sequence_length,
+            lora_scale=lora_scale,
+        )
+        if do_true_cfg:
+            (
+                negative_prompt_embeds,
+                negative_pooled_prompt_embeds,
+                _,
+            ) = self.encode_prompt(
+                prompt=negative_prompt,
+                prompt_2=negative_prompt_2,
+                prompt_embeds=negative_prompt_embeds,
+                pooled_prompt_embeds=negative_pooled_prompt_embeds,
+                device=device,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                lora_scale=lora_scale,
+            )
+
+        # 3.1 Prepare subject emb
+        if subject_image is not None:
+            subject_image = subject_image.resize((max(subject_image.size), max(subject_image.size)))
+            subject_image_embeds_dict = self.encode_image_emb(subject_image, device, dtype)
+
+        # 4. Prepare latent variables
+        num_channels_latents = self.transformer.config.in_channels // 4
+        latents, latent_image_ids = self.prepare_latents(
+            batch_size * num_images_per_prompt,
+            num_channels_latents,
+            height,
+            width,
+            prompt_embeds.dtype,
+            device,
+            generator,
+            latents,
+        )
+
+        # 5. Prepare timesteps
+        sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps) if sigmas is None else sigmas
+        image_seq_len = latents.shape[1]
+        mu = calculate_shift(
+            image_seq_len,
+            self.scheduler.config.base_image_seq_len,
+            self.scheduler.config.max_image_seq_len,
+            self.scheduler.config.base_shift,
+            self.scheduler.config.max_shift,
+        )
+        timesteps, num_inference_steps = retrieve_timesteps(
+            self.scheduler,
+            num_inference_steps,
+            device,
+            sigmas=sigmas,
+            mu=mu,
+        )
+        num_warmup_steps = max(len(timesteps) - num_inference_steps * self.scheduler.order, 0)
+        self._num_timesteps = len(timesteps)
+
+        # handle guidance
+        if self.transformer.config.guidance_embeds:
+            guidance = torch.full([1], guidance_scale, device=device, dtype=torch.float32)
+            guidance = guidance.expand(latents.shape[0])
+        else:
+            guidance = None
+
+        if (ip_adapter_image is not None or ip_adapter_image_embeds is not None) and (
+            negative_ip_adapter_image is None and negative_ip_adapter_image_embeds is None
+        ):
+            negative_ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+        elif (ip_adapter_image is None and ip_adapter_image_embeds is None) and (
+            negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None
+        ):
+            ip_adapter_image = np.zeros((width, height, 3), dtype=np.uint8)
+
+        if self.joint_attention_kwargs is None:
+            self._joint_attention_kwargs = {}
+
+        image_embeds = None
+        negative_image_embeds = None
+        if ip_adapter_image is not None or ip_adapter_image_embeds is not None:
+            image_embeds = self.prepare_ip_adapter_image_embeds(
+                ip_adapter_image,
+                ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+        if negative_ip_adapter_image is not None or negative_ip_adapter_image_embeds is not None:
+            negative_image_embeds = self.prepare_ip_adapter_image_embeds(
+                negative_ip_adapter_image,
+                negative_ip_adapter_image_embeds,
+                device,
+                batch_size * num_images_per_prompt,
+            )
+
+        # 6. Denoising loop
+        with self.progress_bar(total=num_inference_steps) as progress_bar:
+            for i, t in enumerate(timesteps):
+                if self.interrupt:
+                    continue
+
+                if image_embeds is not None:
+                    self._joint_attention_kwargs["ip_adapter_image_embeds"] = image_embeds
+                # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+                timestep = t.expand(latents.shape[0]).to(latents.dtype)
+
+
+                # subject adapter
+                if subject_image is not None:
+                    subject_image_prompt_embeds = self.subject_image_proj_model(
+                        low_res_shallow=subject_image_embeds_dict['image_embeds_low_res_shallow'],
+                        low_res_deep=subject_image_embeds_dict['image_embeds_low_res_deep'],
+                        high_res_deep=subject_image_embeds_dict['image_embeds_high_res_deep'],
+                        timesteps=timestep.to(dtype=latents.dtype), 
+                        need_temb=True
+                    )[0]
+                    self._joint_attention_kwargs['emb_dict'] = dict(
+                        length_encoder_hidden_states=prompt_embeds.shape[1]
+                    )
+                    self._joint_attention_kwargs['subject_emb_dict'] = dict(
+                        ip_hidden_states=subject_image_prompt_embeds,
+                        scale=subject_scale,
+                    )
+    
+                noise_pred = self.transformer(
+                    hidden_states=latents,
+                    timestep=timestep / 1000,
+                    guidance=guidance,
+                    pooled_projections=pooled_prompt_embeds,
+                    encoder_hidden_states=prompt_embeds,
+                    txt_ids=text_ids,
+                    img_ids=latent_image_ids,
+                    joint_attention_kwargs=self.joint_attention_kwargs,
+                    return_dict=False,
+                )[0]
+
+                if do_true_cfg:
+                    if negative_image_embeds is not None:
+                        self._joint_attention_kwargs["ip_adapter_image_embeds"] = negative_image_embeds
+                    neg_noise_pred = self.transformer(
+                        hidden_states=latents,
+                        timestep=timestep / 1000,
+                        guidance=guidance,
+                        pooled_projections=negative_pooled_prompt_embeds,
+                        encoder_hidden_states=negative_prompt_embeds,
+                        txt_ids=text_ids,
+                        img_ids=latent_image_ids,
+                        joint_attention_kwargs=self.joint_attention_kwargs,
+                        return_dict=False,
+                    )[0]
+                    noise_pred = neg_noise_pred + true_cfg_scale * (noise_pred - neg_noise_pred)
+
+                # compute the previous noisy sample x_t -> x_t-1
+                latents_dtype = latents.dtype
+                latents = self.scheduler.step(noise_pred, t, latents, return_dict=False)[0]
+
+                if latents.dtype != latents_dtype:
+                    if torch.backends.mps.is_available():
+                        # some platforms (eg. apple mps) misbehave due to a pytorch bug: https://github.com/pytorch/pytorch/pull/99272
+                        latents = latents.to(latents_dtype)
+
+                if callback_on_step_end is not None:
+                    callback_kwargs = {}
+                    for k in callback_on_step_end_tensor_inputs:
+                        callback_kwargs[k] = locals()[k]
+                    callback_outputs = callback_on_step_end(self, i, t, callback_kwargs)
+
+                    latents = callback_outputs.pop("latents", latents)
+                    prompt_embeds = callback_outputs.pop("prompt_embeds", prompt_embeds)
+
+                # call the callback, if provided
+                if i == len(timesteps) - 1 or ((i + 1) > num_warmup_steps and (i + 1) % self.scheduler.order == 0):
+                    progress_bar.update()
+
+                if XLA_AVAILABLE:
+                    xm.mark_step()
+
+        if output_type == "latent":
+            image = latents
+
+        else:
+            latents = self._unpack_latents(latents, height, width, self.vae_scale_factor)
+            latents = (latents / self.vae.config.scaling_factor) + self.vae.config.shift_factor
+            image = self.vae.decode(latents, return_dict=False)[0]
+            image = self.image_processor.postprocess(image, output_type=output_type)
+
+        # Offload all models
+        self.maybe_free_model_hooks()
+
+        if not return_dict:
+            return (image,)
+
+        return FluxPipelineOutput(images=image)
+
+
+    def with_style_lora(self, lora_file_path, lora_weight=1.0, trigger='', *args, **kwargs):
+        flux_load_lora(self, lora_file_path, lora_weight)
+        kwargs['prompt'] = f"{trigger}, {kwargs['prompt']}"
+        res = self.__call__(*args, **kwargs)
+        flux_load_lora(self, lora_file_path, -lora_weight)
+        return res
+