Upload 10 files

src/adapters/mod_adapters.py

@@ -0,0 +1,243 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Dict, List, Optional, Set, Tuple, Union
+from dataclasses import dataclass
+from inspect import isfunction
+
+import os
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from einops import rearrange, repeat
+
+from diffusers.models.modeling_utils import ModelMixin
+from diffusers.configuration_utils import ConfigMixin, register_to_config
+from diffusers.models.embeddings import TimestepEmbedding, Timesteps
+
+from src.utils.data_utils import pad_to_square, pad_to_target
+
+from transformers import CLIPProcessor, CLIPModel, CLIPVisionModelWithProjection, CLIPVisionModel
+
+from collections import OrderedDict
+
+class SquaredReLU(nn.Module):
+    def forward(self, x: torch.Tensor):
+        return torch.square(torch.relu(x))
+
+class AdaLayerNorm(nn.Module):
+    def __init__(self, embedding_dim: int, time_embedding_dim: Optional[int] = None, ln_bias=True):
+        super().__init__()
+
+        if time_embedding_dim is None:
+            time_embedding_dim = embedding_dim
+
+        self.silu = nn.SiLU()
+        self.linear = nn.Linear(time_embedding_dim, 2 * embedding_dim, bias=True)
+        nn.init.zeros_(self.linear.weight)
+        nn.init.zeros_(self.linear.bias)
+
+        self.norm = nn.LayerNorm(embedding_dim, elementwise_affine=False, eps=1e-6, bias=ln_bias)
+
+    def forward(
+        self, x: torch.Tensor, timestep_embedding: torch.Tensor
+    ) -> tuple[torch.Tensor, torch.Tensor]:
+        emb = self.linear(self.silu(timestep_embedding))
+        shift, scale = emb.view(len(x), 1, -1).chunk(2, dim=-1)
+        x = self.norm(x) * (1 + scale) + shift
+        return x
+
+class PerceiverAttentionBlock(nn.Module):
+    def __init__(
+        self, d_model: int, n_heads: int, 
+        time_embedding_dim: Optional[int] = None,
+        double_kv: Optional[bool] = True,
+    ):
+        super().__init__()
+        self.attn = nn.MultiheadAttention(d_model, n_heads, batch_first=True)
+        self.n_heads = n_heads
+
+        self.mlp = nn.Sequential(
+            OrderedDict(
+                [
+                    ("c_fc", nn.Linear(d_model, d_model * 4)),
+                    ("sq_relu", SquaredReLU()),
+                    ("c_proj", nn.Linear(d_model * 4, d_model)),
+                ]
+            )
+        )
+        self.double_kv = double_kv
+        self.ln_1 = AdaLayerNorm(d_model, time_embedding_dim)
+        self.ln_2 = AdaLayerNorm(d_model, time_embedding_dim)
+        self.ln_ff = AdaLayerNorm(d_model, time_embedding_dim)
+
+    def attention(self, q: torch.Tensor, kv: torch.Tensor, attn_mask: torch.Tensor = None):
+        attn_output, attn_output_weights = self.attn(q, kv, kv, need_weights=False, key_padding_mask=attn_mask)
+        return attn_output
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        latents: torch.Tensor,
+        timestep_embedding: torch.Tensor = None,
+        attn_mask: torch.Tensor = None
+    ):
+        normed_latents = self.ln_1(latents, timestep_embedding)
+        normed_x = self.ln_2(x, timestep_embedding)
+        if self.double_kv:
+            kv = torch.cat([normed_latents, normed_x], dim=1)
+        else:
+            kv = normed_x
+        attn = self.attention(
+            q=normed_latents,
+            kv=kv,
+            attn_mask=attn_mask,
+        )
+        if attn_mask is not None:
+            query_padding_mask = attn_mask.chunk(2, -1)[0].unsqueeze(-1) # (B, 2S) -> (B, S, 1)
+            latents = latents + attn * (~query_padding_mask).to(attn)
+        else:
+            latents = latents + attn
+        latents = latents + self.mlp(self.ln_ff(latents, timestep_embedding))
+        return latents
+   
+
+class CLIPModAdapter(ModelMixin, ConfigMixin):
+    @register_to_config
+    def __init__(
+        self, 
+        out_dim=3072,
+        width=1024,
+        pblock_width=512,
+        layers=6,
+        pblock_layers=1,
+        heads=8,
+        input_text_dim=4096,
+        input_image_dim=1024,
+        pblock_single_blocks=0,
+    ):
+        super().__init__()
+        self.out_dim = out_dim
+
+        self.net = TextImageResampler(
+            width=width,
+            layers=layers,
+            heads=heads,
+            input_text_dim=input_text_dim,
+            input_image_dim=input_image_dim,
+            time_embedding_dim=64,
+            output_dim=out_dim,
+        )
+        self.net2 = TextImageResampler(
+            width=pblock_width,
+            layers=pblock_layers,
+            heads=heads,
+            input_text_dim=input_text_dim,
+            input_image_dim=input_image_dim,
+            time_embedding_dim=64,
+            output_dim=out_dim*(19+pblock_single_blocks),
+        )
+        
+    def enable_gradient_checkpointing(self):
+        self.gradient_checkpointing = True
+        self.net.enable_gradient_checkpointing()
+        self.net2.enable_gradient_checkpointing()
+
+
+    def forward(self, t_emb, llm_hidden_states, clip_outputs):
+        if len(llm_hidden_states.shape) > 3:
+            llm_hidden_states = llm_hidden_states[..., -1, :]
+        batch_size, seq_length = llm_hidden_states.shape[:2]
+
+        img_cls_feat = clip_outputs["image_embeds"] # (B, 768)
+        img_last_feat = clip_outputs["last_hidden_state"] # (B, 257, 1024)
+        img_layer_feats = clip_outputs["hidden_states"] # [(B, 257, 1024) * 25]
+        img_second_last_feat = img_layer_feats[-2] # (B, 257, 1024)
+
+        img_hidden_states = img_second_last_feat # (B, 257, 1024)
+
+        x = self.net(llm_hidden_states, img_hidden_states) # (B, S, 3072)
+        x2 = self.net2(llm_hidden_states, img_hidden_states).view(batch_size, seq_length, -1, self.out_dim) # (B, S, N, 3072)
+        return x, x2
+
+
+class TextImageResampler(nn.Module):
+    def __init__(
+        self,
+        width: int = 768,
+        layers: int = 6,
+        heads: int = 8,
+        output_dim: int = 3072,
+        input_text_dim: int = 4096,
+        input_image_dim: int = 1024,
+        time_embedding_dim: int = 64,
+    ):
+        super().__init__()
+        self.output_dim = output_dim
+        self.input_text_dim = input_text_dim
+        self.input_image_dim = input_image_dim
+        self.time_embedding_dim = time_embedding_dim
+
+        self.text_proj_in = nn.Linear(input_text_dim, width)
+        self.image_proj_in = nn.Linear(input_image_dim, width)
+
+        self.perceiver_blocks = nn.Sequential(
+            *[
+                PerceiverAttentionBlock(
+                    width, heads, time_embedding_dim=self.time_embedding_dim
+                )
+                for _ in range(layers)
+            ]
+        )
+
+        self.proj_out = nn.Sequential(
+            nn.Linear(width, output_dim), nn.LayerNorm(output_dim)
+        )
+        
+        self.gradient_checkpointing = False
+
+    def enable_gradient_checkpointing(self):
+        self.gradient_checkpointing = True
+        
+
+    def forward(
+        self, 
+        text_hidden_states: torch.Tensor, 
+        image_hidden_states: torch.Tensor,
+    ):
+        timestep_embedding = torch.zeros((text_hidden_states.shape[0], 1, self.time_embedding_dim)).to(text_hidden_states)
+
+        text_hidden_states = self.text_proj_in(text_hidden_states)
+        image_hidden_states = self.image_proj_in(image_hidden_states)
+
+        for p_block in self.perceiver_blocks:
+            if self.gradient_checkpointing:
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs)
+                    return custom_forward
+
+                text_hidden_states = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(p_block),
+                    image_hidden_states,
+                    text_hidden_states,
+                    timestep_embedding
+                )
+            else:
+                text_hidden_states = p_block(image_hidden_states, text_hidden_states, timestep_embedding=timestep_embedding)
+
+        text_hidden_states = self.proj_out(text_hidden_states)
+
+        return text_hidden_states

src/flux/block.py

@@ -0,0 +1,814 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from typing import List, Union, Optional, Tuple, Dict, Any, Callable
+from diffusers.models.attention_processor import Attention, F
+from .lora_controller import enable_lora
+from einops import rearrange
+import math
+from diffusers.models.embeddings import apply_rotary_emb
+
+
+def scaled_dot_product_attention(query, key, value, attn_mask=None, dropout_p=0.0, is_causal=False, scale=None) -> torch.Tensor:
+    # Efficient implementation equivalent to the following:
+    L, S = query.size(-2), key.size(-2)
+    B = query.size(0)
+    scale_factor = 1 / math.sqrt(query.size(-1)) if scale is None else scale
+    attn_bias = torch.zeros(B, 1, L, S, dtype=query.dtype, device=query.device)
+    if is_causal:
+        assert attn_mask is None
+        assert False
+        temp_mask = torch.ones(L, S, dtype=torch.bool).tril(diagonal=0)
+        attn_bias.masked_fill_(temp_mask.logical_not(), float("-inf"))
+        attn_bias.to(query.dtype)
+
+    if attn_mask is not None:
+        if attn_mask.dtype == torch.bool:
+            attn_bias.masked_fill_(attn_mask.logical_not(), float("-inf"))
+        else:
+            attn_bias += attn_mask
+    attn_weight = query @ key.transpose(-2, -1) * scale_factor
+    attn_weight += attn_bias.to(attn_weight.device)
+    attn_weight = torch.softmax(attn_weight, dim=-1)
+
+    return torch.dropout(attn_weight, dropout_p, train=True) @ value, attn_weight
+
+def attn_forward(
+    attn: Attention,
+    hidden_states: torch.FloatTensor,
+    encoder_hidden_states: torch.FloatTensor = None,
+    condition_latents: torch.FloatTensor = None,
+    text_cond_mask: Optional[torch.FloatTensor] = None,
+    attention_mask: Optional[torch.FloatTensor] = None,
+    image_rotary_emb: Optional[torch.Tensor] = None,
+    cond_rotary_emb: Optional[torch.Tensor] = None,
+    model_config: Optional[Dict[str, Any]] = {},
+    store_attn_map: bool = False,
+    latent_height: Optional[int] = None,
+    timestep: Optional[torch.Tensor] = None,
+    last_attn_map: Optional[torch.Tensor] = None,
+    condition_sblora_weight: Optional[float] = None,
+    latent_sblora_weight: Optional[float] = None,
+) -> torch.FloatTensor:
+    batch_size, _, _ = (
+        hidden_states.shape
+        if encoder_hidden_states is None
+        else encoder_hidden_states.shape
+    )
+    
+    is_sblock = encoder_hidden_states is None
+    is_dblock = not is_sblock
+    
+    with enable_lora(
+        (attn.to_q, attn.to_k, attn.to_v), 
+        (is_dblock and model_config["latent_lora"]) or (is_sblock and model_config["sblock_lora"]), latent_sblora_weight=latent_sblora_weight
+    ):
+        query = attn.to_q(hidden_states)
+        key = attn.to_k(hidden_states)
+        value = attn.to_v(hidden_states)
+         
+    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.
+        with enable_lora((attn.add_q_proj, attn.add_k_proj, attn.add_v_proj), model_config["text_lora"]):
+            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)
+
+    if condition_latents is not None:
+        assert condition_latents.shape[0] == batch_size
+        cond_length = condition_latents.shape[1]
+
+        cond_lora_activate = (is_dblock and model_config["use_condition_dblock_lora"]) or (is_sblock and model_config["use_condition_sblock_lora"])
+        with enable_lora(
+            (attn.to_q, attn.to_k, attn.to_v), 
+            dit_activated=not cond_lora_activate, cond_activated=cond_lora_activate, latent_sblora_weight=condition_sblora_weight  #TODO implementation for condition lora not share
+        ):
+            cond_query = attn.to_q(condition_latents)
+            cond_key = attn.to_k(condition_latents)
+            cond_value = attn.to_v(condition_latents)
+
+        cond_query = cond_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        cond_key = cond_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        cond_value = cond_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
+        if attn.norm_q is not None:
+            cond_query = attn.norm_q(cond_query)
+        if attn.norm_k is not None:
+            cond_key = attn.norm_k(cond_key)
+
+        if cond_rotary_emb is not None:
+            cond_query = apply_rotary_emb(cond_query, cond_rotary_emb)
+            cond_key = apply_rotary_emb(cond_key, cond_rotary_emb)
+
+        if model_config.get("text_cond_attn", False):
+            if encoder_hidden_states is not None:
+                assert text_cond_mask is not None
+                img_length = hidden_states.shape[1]
+                seq_length = encoder_hidden_states_query_proj.shape[2]
+                assert len(text_cond_mask.shape) == 2 or len(text_cond_mask.shape) == 3
+                if len(text_cond_mask.shape) == 2:
+                    text_cond_mask = text_cond_mask.unsqueeze(-1)
+                N = text_cond_mask.shape[-1] # num_condition
+            else:
+                raise NotImplementedError()
+
+            query = torch.cat([query, cond_query], dim=2) # (B, 24, S+HW+NC)
+            key = torch.cat([key, cond_key], dim=2)
+            value = torch.cat([value, cond_value], dim=2)
+            
+            assert query.shape[2] == key.shape[2]
+            assert query.shape[2] == cond_length + img_length + seq_length
+
+            attention_mask = torch.ones(batch_size, 1, query.shape[2], key.shape[2], device=query.device, dtype=torch.bool)
+            attention_mask[..., -cond_length:, :-cond_length] = False
+            attention_mask[..., :-cond_length, -cond_length:] = False
+
+            if encoder_hidden_states is not None:
+                tokens_per_cond = cond_length // N
+                for i in range(batch_size):
+                    for j in range(N):
+                        start = seq_length + img_length + tokens_per_cond * j
+                        attention_mask[i, 0, :seq_length, start:start+tokens_per_cond] = text_cond_mask[i, :, j].unsqueeze(-1)
+
+        elif model_config.get("union_cond_attn", False):
+            query = torch.cat([query, cond_query], dim=2) # (B, 24, S+HW+NC)
+            key = torch.cat([key, cond_key], dim=2)
+            value = torch.cat([value, cond_value], dim=2)
+            
+            attention_mask = torch.ones(batch_size, 1, query.shape[2], key.shape[2], device=query.device, dtype=torch.bool)
+            cond_length = condition_latents.shape[1]
+            assert len(text_cond_mask.shape) == 2 or len(text_cond_mask.shape) == 3
+            if len(text_cond_mask.shape) == 2:
+                text_cond_mask = text_cond_mask.unsqueeze(-1)
+            N = text_cond_mask.shape[-1] # num_condition
+            tokens_per_cond = cond_length // N
+            
+            seq_length = 0
+            if encoder_hidden_states is not None:
+                seq_length = encoder_hidden_states_query_proj.shape[2]
+                img_length = hidden_states.shape[1]
+            else:
+                seq_length = 128 # TODO, pass it here
+                img_length = hidden_states.shape[1] - seq_length
+            
+            if not model_config.get("cond_cond_cross_attn", True):
+                # no cross attention between different conds
+                cond_start = seq_length + img_length
+                attention_mask[:, :, cond_start:, cond_start:] = False
+                
+                for j in range(N):
+                    start = cond_start + tokens_per_cond * j
+                    end = cond_start + tokens_per_cond * (j + 1)
+                    attention_mask[..., start:end, start:end] = True
+            
+            # double block
+            if encoder_hidden_states is not None:
+                
+                # no cross attention
+                attention_mask[..., :-cond_length, -cond_length:] = False
+
+                if model_config.get("use_attention_double", False) and last_attn_map is not None:
+                    attention_mask = torch.zeros(batch_size, 1, query.shape[2], key.shape[2], device=query.device, dtype=torch.bfloat16)
+                    last_attn_map = last_attn_map.to(query.device)
+                    attention_mask[..., seq_length:-cond_length, :seq_length] = torch.log(last_attn_map/last_attn_map.mean()*model_config["use_atten_lambda"]).view(-1, seq_length)
+                
+            # single block
+            else:
+                # print(last_attn_map)
+                if model_config.get("use_attention_single", False) and last_attn_map is not None:
+                    attention_mask = torch.zeros(batch_size, 1, query.shape[2], key.shape[2], device=query.device, dtype=torch.bfloat16)
+                    attention_mask[..., :seq_length, -cond_length:] = float('-inf')
+                    # 确保 use_atten_lambda 是列表
+                    use_atten_lambdas = model_config["use_atten_lambda"] if len(model_config["use_atten_lambda"])!=1 else model_config["use_atten_lambda"] * (N+1)
+                    attention_mask[..., -cond_length:, seq_length:-cond_length] = math.log(use_atten_lambdas[0])
+                    last_attn_map = last_attn_map.to(query.device)
+                    
+                    cond2latents = []
+                    for i in range(batch_size):
+                        AM = last_attn_map[i] # (H, W, S)
+                        for j in range(N):
+                            start = seq_length + img_length + tokens_per_cond * j
+                            mask = text_cond_mask[i, :, j] # (S,)
+                            weighted_AM = AM * mask.unsqueeze(0).unsqueeze(0)  # 扩展 mask 维度以匹配 AM
+
+                            cond2latent = weighted_AM.mean(-1)
+                            if model_config.get("attention_norm", "mean") == "max":
+                                cond2latent = cond2latent / cond2latent.max()  # 归一化
+                            else:
+                                cond2latent = cond2latent / cond2latent.mean()  # 归一化
+                            cond2latent = cond2latent.view(-1,) # (WH,)
+
+                            # 使用对应 condition 的 lambda 值
+                            current_lambda = use_atten_lambdas[j+1]
+                            # 将 cond2latent 复制到 attention_mask[i, 0, :seq_length, start:start+tokens_per_cond]
+                            attention_mask[i, 0, seq_length:-cond_length, start:start+tokens_per_cond] = torch.log(current_lambda * cond2latent.unsqueeze(-1))
+                            
+                            # 将 text_cond_mask[i, :, j].unsqueeze(-1) 为 true 的位置设置为当前 lambda 值
+                            cond = mask.unsqueeze(-1).expand(-1, tokens_per_cond)
+                            sub_mask = attention_mask[i, 0, :seq_length, start:start+tokens_per_cond]
+                            attention_mask[i, 0, :seq_length, start:start+tokens_per_cond] = torch.where(cond, math.log(current_lambda), sub_mask)
+                            cond2latents.append(
+                                cond2latent.reshape(latent_height, -1).detach().cpu()
+                            )
+                    if store_attn_map:
+                        if not hasattr(attn, "cond2latents"):
+                            attn.cond2latents = []
+                            attn.cond_timesteps = []
+                        attn.cond2latents.append(torch.stack(cond2latents, dim=0)) # (N, H, W)
+                        attn.cond_timesteps.append(timestep.cpu())
+
+                pass
+        else:
+            raise NotImplementedError()
+        if hasattr(attn, "c_factor"):
+            assert False
+            attention_mask = torch.zeros(
+                query.shape[2], key.shape[2], device=query.device, dtype=query.dtype
+            )
+            bias = torch.log(attn.c_factor[0])
+            attention_mask[-cond_length:, :-cond_length] = bias
+            attention_mask[:-cond_length, -cond_length:] = bias
+
+    ####################################################################################################
+    if store_attn_map and encoder_hidden_states is not None:
+        seq_length = encoder_hidden_states_query_proj.shape[2]
+        img_length = hidden_states.shape[1]
+        hidden_states, attention_probs = scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+        # (B, 24, S+HW, S+HW) -> (B, 24, HW, S)
+        t2i_attention_probs = attention_probs[:, :, seq_length:seq_length+img_length, :seq_length]
+        # (B, 24, S+HW, S+HW) -> (B, 24, S, HW) -> (B, 24, HW, S)
+        i2t_attention_probs = attention_probs[:, :, :seq_length, seq_length:seq_length+img_length].transpose(-1, -2)
+        
+        if not hasattr(attn, "attn_maps"):
+            attn.attn_maps = []
+            attn.timestep = []
+
+        attn.attn_maps.append(
+            (
+                rearrange(t2i_attention_probs, 'B attn_head (H W) attn_dim -> B attn_head H W attn_dim', H=latent_height),
+                rearrange(i2t_attention_probs, 'B attn_head (H W) attn_dim -> B attn_head H W attn_dim', H=latent_height),
+            )
+        )
+
+        attn.timestep.append(timestep.cpu())
+        has_nan = torch.isnan(hidden_states).any().item()
+        if has_nan:
+            print("[attn_forward] detect nan hidden_states in store_attn_map")
+    else:
+        hidden_states = F.scaled_dot_product_attention(
+            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
+        )
+        has_nan = torch.isnan(hidden_states).any().item()
+        if has_nan:
+            print("[attn_forward] detect nan hidden_states")
+    ####################################################################################################
+    hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim).to(query.dtype)
+
+    if encoder_hidden_states is not None:
+        if condition_latents is not None:
+            encoder_hidden_states, hidden_states, condition_latents = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[
+                    :, encoder_hidden_states.shape[1] : -condition_latents.shape[1]
+                ],
+                hidden_states[:, -condition_latents.shape[1] :],
+            )
+            if model_config.get("latent_cond_by_text_attn", False):
+                # hidden_states += add_latent # (B, HW, D)
+                hidden_states = new_hidden_states # (B, HW, D)
+            
+        else:
+            encoder_hidden_states, hidden_states = (
+                hidden_states[:, : encoder_hidden_states.shape[1]],
+                hidden_states[:, encoder_hidden_states.shape[1] :],
+            )
+
+
+        with enable_lora((attn.to_out[0],), model_config["latent_lora"]):
+            hidden_states = attn.to_out[0](hidden_states) # linear proj
+            hidden_states = attn.to_out[1](hidden_states) # dropout
+        with enable_lora((attn.to_add_out,), model_config["text_lora"]):
+            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)
+
+        if condition_latents is not None:
+            cond_lora_activate = model_config["use_condition_dblock_lora"]
+            with enable_lora(
+                (attn.to_out[0],), 
+                dit_activated=not cond_lora_activate, cond_activated=cond_lora_activate,
+            ):
+                condition_latents = attn.to_out[0](condition_latents)
+                condition_latents = attn.to_out[1](condition_latents)
+
+
+        return (
+            (hidden_states, encoder_hidden_states, condition_latents)
+            if condition_latents is not None
+            else (hidden_states, encoder_hidden_states)
+        )
+    elif condition_latents is not None:
+        hidden_states, condition_latents = (
+            hidden_states[:, : -condition_latents.shape[1]],
+            hidden_states[:, -condition_latents.shape[1] :],
+        )
+        return hidden_states, condition_latents
+    else:
+        return hidden_states
+
+
+def set_delta_by_start_end(
+    start_ends, 
+    src_delta_emb, src_delta_emb_pblock, 
+    delta_emb, delta_emb_pblock, delta_emb_mask, 
+):
+    for (i, j, src_s, src_e, tar_s, tar_e) in start_ends:
+        if src_delta_emb is not None:
+            delta_emb[i, tar_s:tar_e] = src_delta_emb[j, src_s:src_e]
+        if src_delta_emb_pblock is not None:
+            delta_emb_pblock[i, tar_s:tar_e] = src_delta_emb_pblock[j, src_s:src_e]
+        delta_emb_mask[i, tar_s:tar_e] = True
+    return delta_emb, delta_emb_pblock, delta_emb_mask
+
+def norm1_context_forward(
+    self,
+    x: torch.Tensor,
+    condition_latents: Optional[torch.Tensor] = None, 
+    timestep: Optional[torch.Tensor] = None,
+    class_labels: Optional[torch.LongTensor] = None,
+    hidden_dtype: Optional[torch.dtype] = None,
+    emb: Optional[torch.Tensor] = None,
+    delta_emb: Optional[torch.Tensor] = None,
+    delta_emb_cblock: Optional[torch.Tensor] = None,
+    delta_emb_mask: Optional[torch.Tensor] = None,
+    delta_start_ends = None,
+    mod_adapter = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    batch_size, seq_length = x.shape[:2]
+    
+    if mod_adapter is not None:
+        assert False
+
+    if delta_emb is None:
+        emb = self.linear(self.silu(emb)) # (B, 3072) -> (B, 18432)
+        emb = emb.unsqueeze(1) # (B, 1, 18432)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1) # (B, 1, 3072)
+        x = self.norm(x) * (1 + scale_msa) + shift_msa # (B, 1, 3072)
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+    else:
+        # (B, 3072) > (B, 18432) -> (B, S, 18432)
+        emb_orig = self.linear(self.silu(emb)).unsqueeze(1).expand((-1, seq_length, -1))
+        # (B, 3072) -> (B, 1, 3072) -> (B, S, 3072) -> (B, S, 18432)
+        if delta_emb_cblock is None:
+            emb_new = self.linear(self.silu(emb.unsqueeze(1) + delta_emb))
+        else:
+            emb_new = self.linear(self.silu(emb.unsqueeze(1) + delta_emb + delta_emb_cblock))
+        emb = torch.where(delta_emb_mask.unsqueeze(-1), emb_new, emb_orig) # (B, S, 18432)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1) # (B, S, 3072)
+        x = self.norm(x) * (1 + scale_msa) + shift_msa # (B, S, 3072)
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+        
+
+def norm1_forward(
+    self,
+    x: torch.Tensor,
+    timestep: Optional[torch.Tensor] = None,
+    class_labels: Optional[torch.LongTensor] = None,
+    hidden_dtype: Optional[torch.dtype] = None,
+    emb: Optional[torch.Tensor] = None,
+    delta_emb: Optional[torch.Tensor] = None,
+    delta_emb_cblock: Optional[torch.Tensor] = None,
+    delta_emb_mask: Optional[torch.Tensor] = None,
+    t2i_attn_map: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    if delta_emb is None:
+        emb = self.linear(self.silu(emb)) # (B, 3072) -> (B, 18432)
+        emb = emb.unsqueeze(1) # (B, 1, 18432)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1) # (B, 1, 3072)
+        x = self.norm(x) * (1 + scale_msa) + shift_msa # (B, 1, 3072)
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+    else:
+        raise NotImplementedError()
+        batch_size, HW = x.shape[:2]
+        seq_length = t2i_attn_map.shape[-1]
+        # (B, 3072) > (B, 18432) -> (B, S, 18432)
+        emb_orig = self.linear(self.silu(emb)).unsqueeze(1).expand((-1, seq_length, -1))
+        # (B, 3072) -> (B, 1, 3072) -> (B, S, 3072) -> (B, S, 18432)
+        if delta_emb_cblock is None:
+            emb_new = self.linear(self.silu(emb.unsqueeze(1) + delta_emb))
+        else:
+            emb_new = self.linear(self.silu(emb.unsqueeze(1) + delta_emb + delta_emb_cblock))
+        # attn_weight (B, HW, S)
+        emb = torch.where(delta_emb_mask.unsqueeze(-1), emb_new, emb_orig) # (B, S, 18432)
+        emb = t2i_attn_map @ emb    # (B, HW, 18432)
+        shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = emb.chunk(6, dim=-1) # (B, HW, 3072)
+        x = self.norm(x) * (1 + scale_msa) + shift_msa # (B, HW, 3072)
+        return x, gate_msa, shift_mlp, scale_mlp, gate_mlp
+
+
+def block_forward(
+    self,
+    hidden_states: torch.FloatTensor,
+    encoder_hidden_states: torch.FloatTensor,
+    condition_latents: torch.FloatTensor,
+    temb: torch.FloatTensor,
+    cond_temb: torch.FloatTensor,
+    text_cond_mask: Optional[torch.FloatTensor] = None,
+    delta_emb: Optional[torch.FloatTensor] = None,
+    delta_emb_cblock: Optional[torch.FloatTensor] = None,
+    delta_emb_mask: Optional[torch.Tensor] = None,
+    delta_start_ends = None,
+    cond_rotary_emb=None,
+    image_rotary_emb=None,
+    model_config: Optional[Dict[str, Any]] = {},
+    store_attn_map: bool = False,
+    use_text_mod: bool = True,
+    use_img_mod: bool = False,
+    mod_adapter = None,
+    latent_height: Optional[int] = None,
+    timestep: Optional[torch.Tensor] = None,
+    last_attn_map: Optional[torch.Tensor] = None,
+):
+    batch_size = hidden_states.shape[0]
+    use_cond = condition_latents is not None
+
+    train_partial_latent_lora = model_config.get("train_partial_latent_lora", False)
+    train_partial_text_lora = model_config.get("train_partial_text_lora", False)
+    if train_partial_latent_lora:
+        train_partial_latent_lora_layers = model_config.get("train_partial_latent_lora_layers", "")
+        activate_norm1 = activate_ff = True
+        if "norm1" not in train_partial_latent_lora_layers:
+            activate_norm1 = False
+        if "ff" not in train_partial_latent_lora_layers:
+            activate_ff = False
+    
+    if train_partial_text_lora:
+        train_partial_text_lora_layers = model_config.get("train_partial_text_lora_layers", "")
+        activate_norm1_context = activate_ff_context = True
+        if "norm1" not in train_partial_text_lora_layers:
+            activate_norm1_context = False
+        if "ff" not in train_partial_text_lora_layers:
+            activate_ff_context = False
+    
+    if use_cond:
+        cond_lora_activate = model_config["use_condition_dblock_lora"]
+        with enable_lora(
+            (self.norm1.linear,), 
+            dit_activated=activate_norm1 if train_partial_latent_lora else not cond_lora_activate, cond_activated=cond_lora_activate,
+        ):
+            norm_condition_latents, cond_gate_msa, cond_shift_mlp, cond_scale_mlp, cond_gate_mlp = (
+                norm1_forward(
+                    self.norm1,
+                    condition_latents,
+                    emb=cond_temb, 
+                )
+            )
+    delta_emb_img = delta_emb_img_cblock = None
+    if use_img_mod and use_text_mod:
+        if delta_emb is not None:
+            delta_emb_img, delta_emb = delta_emb.chunk(2, dim=-1)
+        if delta_emb_cblock is not None:
+            delta_emb_img_cblock, delta_emb_cblock = delta_emb_cblock.chunk(2, dim=-1)
+
+    with enable_lora((self.norm1.linear,), activate_norm1 if train_partial_latent_lora else model_config["latent_lora"]):
+        if use_img_mod and encoder_hidden_states is not None:
+            with torch.no_grad():
+                attn = self.attn
+
+                norm_img = self.norm1(hidden_states, emb=temb)[0]
+                norm_text = self.norm1_context(encoder_hidden_states, emb=temb)[0]
+                
+                img_query = attn.to_q(norm_img)
+                img_key = attn.to_k(norm_img)
+                text_query = attn.add_q_proj(norm_text)
+                text_key = attn.add_k_proj(norm_text)
+
+                inner_dim = img_key.shape[-1]
+                head_dim = inner_dim // attn.heads
+
+                img_query = img_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)   # (B, N, HW, D)
+                img_key = img_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)       # (B, N, HW, D)
+                text_query = text_query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2) # (B, N, S, D)
+                text_key = text_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)     # (B, N, S, D)
+
+                if attn.norm_q is not None:
+                    img_query = attn.norm_q(img_query)
+                if attn.norm_added_q is not None:
+                    text_query = attn.norm_added_q(text_query)
+                if attn.norm_k is not None:
+                    img_key = attn.norm_k(img_key)
+                if attn.norm_added_k is not None:
+                    text_key = attn.norm_added_k(text_key)
+                
+                query = torch.cat([text_query, img_query], dim=2) # (B, N, S+HW, D)
+                key = torch.cat([text_key, img_key], dim=2)       # (B, N, S+HW, D)
+                if image_rotary_emb is not None:
+                    query = apply_rotary_emb(query, image_rotary_emb)
+                    key = apply_rotary_emb(key, image_rotary_emb)
+
+                seq_length = text_query.shape[2]
+
+                scale_factor = 1 / math.sqrt(query.size(-1))
+                t2i_attn_map = query @ key.transpose(-2, -1) * scale_factor # (B, N, S+HW, S+HW)
+                t2i_attn_map = t2i_attn_map.mean(1)[:, seq_length:, :seq_length] # (B, S+HW, S+HW) -> (B, HW, S)
+                t2i_attn_map = torch.softmax(t2i_attn_map, dim=-1) # (B, HW, S)
+
+        else:
+            t2i_attn_map = None
+
+        norm_hidden_states, gate_msa, shift_mlp, scale_mlp, gate_mlp = (
+            norm1_forward(
+                self.norm1,
+                hidden_states,
+                emb=temb,
+                delta_emb=delta_emb_img,
+                delta_emb_cblock=delta_emb_img_cblock,
+                delta_emb_mask=delta_emb_mask,
+                t2i_attn_map=t2i_attn_map,
+            )
+        )
+    # Modulation for double block
+    with enable_lora((self.norm1_context.linear,), activate_norm1_context if train_partial_text_lora else model_config["text_lora"]):
+        norm_encoder_hidden_states, c_gate_msa, c_shift_mlp, c_scale_mlp, c_gate_mlp = (
+            norm1_context_forward(
+                self.norm1_context, 
+                encoder_hidden_states, 
+                emb=temb, 
+                delta_emb=delta_emb if use_text_mod else None,
+                delta_emb_cblock=delta_emb_cblock if use_text_mod else None,
+                delta_emb_mask=delta_emb_mask if use_text_mod else None,
+                delta_start_ends=delta_start_ends if use_text_mod else None,
+                mod_adapter=mod_adapter,
+                condition_latents=condition_latents,
+            )
+        )
+
+    # Attention.
+    result = attn_forward(
+        self.attn,
+        model_config=model_config,
+        hidden_states=norm_hidden_states,
+        encoder_hidden_states=norm_encoder_hidden_states,
+        condition_latents=norm_condition_latents if use_cond else None,
+        text_cond_mask=text_cond_mask if use_cond else None,
+        image_rotary_emb=image_rotary_emb,
+        cond_rotary_emb=cond_rotary_emb if use_cond else None,
+        store_attn_map=store_attn_map,
+        latent_height=latent_height,
+        timestep=timestep,
+        last_attn_map=last_attn_map,
+    )
+    attn_output, context_attn_output = result[:2]
+    cond_attn_output = result[2] if use_cond else None
+
+    # Process attention outputs for the `hidden_states`.
+    # 1. hidden_states
+    attn_output = gate_msa * attn_output # NOTE: changed by img mod
+    hidden_states = hidden_states + attn_output
+    # 2. encoder_hidden_states
+    context_attn_output = c_gate_msa * context_attn_output # NOTE: changed by delta_temb
+    encoder_hidden_states = encoder_hidden_states + context_attn_output
+    # 3. condition_latents
+    if use_cond:
+        cond_attn_output = cond_gate_msa * cond_attn_output  # NOTE: changed by img mod
+        condition_latents = condition_latents + cond_attn_output
+        if model_config.get("add_cond_attn", False):
+            hidden_states += cond_attn_output
+
+    # LayerNorm + MLP.
+    # 1. hidden_states
+    norm_hidden_states = self.norm2(hidden_states)
+    norm_hidden_states = (
+        norm_hidden_states * (1 + scale_mlp) + shift_mlp  # NOTE: changed by img mod
+    )
+    # 2. encoder_hidden_states
+    norm_encoder_hidden_states = self.norm2_context(encoder_hidden_states)
+    norm_encoder_hidden_states = (
+        norm_encoder_hidden_states * (1 + c_scale_mlp) + c_shift_mlp  # NOTE: changed by delta_temb
+    )
+    # 3. condition_latents
+    if use_cond:
+        norm_condition_latents = self.norm2(condition_latents)
+        norm_condition_latents = (
+            norm_condition_latents * (1 + cond_scale_mlp) + cond_shift_mlp # NOTE: changed by img mod
+        )
+
+    # Feed-forward.
+    with enable_lora((self.ff.net[2],), activate_ff if train_partial_latent_lora else model_config["latent_lora"]):
+        # 1. hidden_states
+        ff_output = self.ff(norm_hidden_states)
+        ff_output = gate_mlp * ff_output  # NOTE: changed by img mod
+    # 2. encoder_hidden_states
+    with enable_lora((self.ff_context.net[2],), activate_ff_context if train_partial_text_lora else model_config["text_lora"]):
+        context_ff_output = self.ff_context(norm_encoder_hidden_states)
+        context_ff_output = c_gate_mlp * context_ff_output  # NOTE: changed by delta_temb
+    # 3. condition_latents
+    if use_cond:
+        cond_lora_activate = model_config["use_condition_dblock_lora"]
+        with enable_lora(
+            (self.ff.net[2],), 
+            dit_activated=activate_ff if train_partial_latent_lora else not cond_lora_activate, cond_activated=cond_lora_activate,
+        ):
+            cond_ff_output = self.ff(norm_condition_latents)
+            cond_ff_output = cond_gate_mlp * cond_ff_output  # NOTE: changed by img mod
+
+    # Process feed-forward outputs.
+    hidden_states = hidden_states + ff_output
+    encoder_hidden_states = encoder_hidden_states + context_ff_output
+    if use_cond:
+        condition_latents = condition_latents + cond_ff_output
+
+    # Clip to avoid overflow.
+    if encoder_hidden_states.dtype == torch.float16:
+        encoder_hidden_states = encoder_hidden_states.clip(-65504, 65504)
+
+    return encoder_hidden_states, hidden_states, condition_latents if use_cond else None
+
+def single_norm_forward(
+    self,
+    x: torch.Tensor,
+    timestep: Optional[torch.Tensor] = None,
+    class_labels: Optional[torch.LongTensor] = None,
+    hidden_dtype: Optional[torch.dtype] = None,
+    emb: Optional[torch.Tensor] = None,
+    delta_emb: Optional[torch.Tensor] = None,
+    delta_emb_cblock: Optional[torch.Tensor] = None,
+    delta_emb_mask: Optional[torch.Tensor] = None,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    if delta_emb is None:
+        emb = self.linear(self.silu(emb)) # (B, 3072) -> (B, 9216)
+        emb = emb.unsqueeze(1) # (B, 1, 9216)
+        shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=-1) # (B, 1, 3072)
+        x = self.norm(x) * (1 + scale_msa) + shift_msa # (B, S, 3072) * (B, 1, 3072)
+        return x, gate_msa
+    else:
+        img_text_seq_length = x.shape[1] # S+
+        text_seq_length = delta_emb_mask.shape[1] # S
+        # (B, 3072) -> (B, 9216) -> (B, S+, 9216)
+        emb_orig = self.linear(self.silu(emb)).unsqueeze(1).expand((-1, img_text_seq_length, -1))
+        # (B, 3072) -> (B, 1, 3072) -> (B, S, 3072) -> (B, S, 9216)
+        if delta_emb_cblock is None:
+            emb_new = self.linear(self.silu(emb.unsqueeze(1) + delta_emb))
+        else:
+            emb_new = self.linear(self.silu(emb.unsqueeze(1) + delta_emb + delta_emb_cblock))
+
+        emb_text = torch.where(delta_emb_mask.unsqueeze(-1), emb_new, emb_orig[:, :text_seq_length]) # (B, S, 9216)
+        emb_img = emb_orig[:, text_seq_length:] # (B, s, 9216)
+        emb = torch.cat([emb_text, emb_img], dim=1) # (B, S+, 9216)
+
+        shift_msa, scale_msa, gate_msa = emb.chunk(3, dim=-1) # (B, S+, 3072)
+        x = self.norm(x) * (1 + scale_msa) + shift_msa # (B, S+, 3072)
+        return x, gate_msa
+
+
+def single_block_forward(
+    self,
+    hidden_states: torch.FloatTensor,
+    temb: torch.FloatTensor,
+    image_rotary_emb=None,
+    condition_latents: torch.FloatTensor = None,
+    text_cond_mask: torch.FloatTensor = None,
+    cond_temb: torch.FloatTensor = None,
+    delta_emb: Optional[torch.FloatTensor] = None,
+    delta_emb_cblock: Optional[torch.FloatTensor] = None,
+    delta_emb_mask: Optional[torch.Tensor] = None,
+    use_text_mod: bool = True,
+    use_img_mod: bool = False,
+    cond_rotary_emb=None,
+    latent_height: Optional[int] = None,
+    timestep: Optional[torch.Tensor] = None,
+    store_attn_map: bool = False,
+    model_config: Optional[Dict[str, Any]] = {},
+    last_attn_map: Optional[torch.Tensor] = None,
+    latent_sblora_weight=None,
+    condition_sblora_weight=None,
+):
+
+    using_cond = condition_latents is not None
+    residual = hidden_states
+    
+    train_partial_lora = model_config.get("train_partial_lora", False)
+    if train_partial_lora:
+        train_partial_lora_layers = model_config.get("train_partial_lora_layers", "")
+        activate_norm = activate_projmlp = activate_projout = True
+        
+        if "norm" not in train_partial_lora_layers:
+            activate_norm = False
+        if "projmlp" not in train_partial_lora_layers:
+            activate_projmlp = False
+        if "projout" not in train_partial_lora_layers:
+            activate_projout = False
+
+    with enable_lora((self.norm.linear,), activate_norm if train_partial_lora else model_config["sblock_lora"], latent_sblora_weight=latent_sblora_weight):
+        # Modulation for single block
+        norm_hidden_states, gate = single_norm_forward(
+            self.norm, 
+            hidden_states, 
+            emb=temb,
+            delta_emb=delta_emb if use_text_mod else None,
+            delta_emb_cblock=delta_emb_cblock if use_text_mod else None,
+            delta_emb_mask=delta_emb_mask if use_text_mod else None,
+        )
+    with enable_lora((self.proj_mlp,), activate_projmlp if train_partial_lora else model_config["sblock_lora"], latent_sblora_weight=latent_sblora_weight):
+        mlp_hidden_states = self.act_mlp(self.proj_mlp(norm_hidden_states))
+    if using_cond:
+        cond_lora_activate = model_config["use_condition_sblock_lora"]
+        with enable_lora(
+            (self.norm.linear,), 
+            dit_activated=activate_norm if train_partial_lora else not cond_lora_activate, cond_activated=cond_lora_activate, latent_sblora_weight=condition_sblora_weight
+        ):
+            residual_cond = condition_latents
+            norm_condition_latents, cond_gate = self.norm(condition_latents, emb=cond_temb)
+        with enable_lora(
+            (self.proj_mlp,), 
+            dit_activated=activate_projmlp if train_partial_lora else not cond_lora_activate, cond_activated=cond_lora_activate, latent_sblora_weight=condition_sblora_weight
+        ):
+            mlp_cond_hidden_states = self.act_mlp(self.proj_mlp(norm_condition_latents))
+
+    attn_output = attn_forward(
+        self.attn,
+        model_config=model_config,
+        hidden_states=norm_hidden_states,
+        image_rotary_emb=image_rotary_emb,
+        last_attn_map=last_attn_map,
+        latent_height=latent_height,
+        store_attn_map=store_attn_map,
+        timestep=timestep,
+        latent_sblora_weight=latent_sblora_weight,
+        condition_sblora_weight=condition_sblora_weight,
+        **(
+            {
+                "condition_latents": norm_condition_latents,
+                "cond_rotary_emb": cond_rotary_emb if using_cond else None,
+                "text_cond_mask": text_cond_mask if using_cond else None,
+            }
+            if using_cond
+            else {}
+        ),
+    )
+    if using_cond:
+        attn_output, cond_attn_output = attn_output
+
+    with enable_lora((self.proj_out,), activate_projout if train_partial_lora else model_config["sblock_lora"], latent_sblora_weight=latent_sblora_weight):
+        hidden_states = torch.cat([attn_output, mlp_hidden_states], dim=2)
+        # gate = (B, 1, 3072) or (B, S+, 3072)
+        hidden_states = gate * self.proj_out(hidden_states)
+        hidden_states = residual + hidden_states
+    if using_cond:
+        cond_lora_activate = model_config["use_condition_sblock_lora"]
+        with enable_lora(
+            (self.proj_out,), 
+            dit_activated=activate_projout if train_partial_lora else not cond_lora_activate, cond_activated=cond_lora_activate, latent_sblora_weight=condition_sblora_weight
+        ):
+            condition_latents = torch.cat([cond_attn_output, mlp_cond_hidden_states], dim=2)
+            cond_gate = cond_gate.unsqueeze(1)
+            condition_latents = cond_gate * self.proj_out(condition_latents)
+            condition_latents = residual_cond + condition_latents
+
+    if hidden_states.dtype == torch.float16:
+        hidden_states = hidden_states.clip(-65504, 65504)
+
+    return hidden_states if not using_cond else (hidden_states, condition_latents)

src/flux/condition.py

@@ -0,0 +1,133 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from typing import Optional, Union, List, Tuple
+from diffusers.pipelines import FluxPipeline
+from PIL import Image, ImageFilter
+import numpy as np
+import cv2
+
+from .pipeline_tools import encode_vae_images
+
+condition_dict = {
+    "depth": 0,
+    "canny": 1,
+    "subject": 4,
+    "coloring": 6,
+    "deblurring": 7,
+    "depth_pred": 8,
+    "fill": 9,
+    "sr": 10,
+}
+
+
+class Condition(object):
+    def __init__(
+        self,
+        condition_type: str,
+        raw_img: Union[Image.Image, torch.Tensor] = None,
+        condition: Union[Image.Image, torch.Tensor] = None,
+        mask=None,
+        position_delta=None,
+    ) -> None:
+        self.condition_type = condition_type
+        assert raw_img is not None or condition is not None
+        if raw_img is not None:
+            self.condition = self.get_condition(condition_type, raw_img)
+        else:
+            self.condition = condition
+        self.position_delta = position_delta
+        # TODO: Add mask support
+        assert mask is None, "Mask not supported yet"
+
+    def get_condition(
+        self, condition_type: str, raw_img: Union[Image.Image, torch.Tensor]
+    ) -> Union[Image.Image, torch.Tensor]:
+        """
+        Returns the condition image.
+        """
+        if condition_type == "depth":
+            from transformers import pipeline
+
+            depth_pipe = pipeline(
+                task="depth-estimation",
+                model="LiheYoung/depth-anything-small-hf",
+                device="cuda",
+            )
+            source_image = raw_img.convert("RGB")
+            condition_img = depth_pipe(source_image)["depth"].convert("RGB")
+            return condition_img
+        elif condition_type == "canny":
+            img = np.array(raw_img)
+            edges = cv2.Canny(img, 100, 200)
+            edges = Image.fromarray(edges).convert("RGB")
+            return edges
+        elif condition_type == "subject":
+            return raw_img
+        elif condition_type == "coloring":
+            return raw_img.convert("L").convert("RGB")
+        elif condition_type == "deblurring":
+            condition_image = (
+                raw_img.convert("RGB")
+                .filter(ImageFilter.GaussianBlur(10))
+                .convert("RGB")
+            )
+            return condition_image
+        elif condition_type == "fill":
+            return raw_img.convert("RGB")
+        return self.condition
+
+    @property
+    def type_id(self) -> int:
+        """
+        Returns the type id of the condition.
+        """
+        return condition_dict[self.condition_type]
+
+    @classmethod
+    def get_type_id(cls, condition_type: str) -> int:
+        """
+        Returns the type id of the condition.
+        """
+        return condition_dict[condition_type]
+
+    def encode(self, pipe: FluxPipeline) -> Tuple[torch.Tensor, torch.Tensor, int]:
+        """
+        Encodes the condition into tokens, ids and type_id.
+        """
+        if self.condition_type in [
+            "depth",
+            "canny",
+            "subject",
+            "coloring",
+            "deblurring",
+            "depth_pred",
+            "fill",
+            "sr",
+        ]:
+            tokens, ids = encode_vae_images(pipe, self.condition)
+        else:
+            raise NotImplementedError(
+                f"Condition type {self.condition_type} not implemented"
+            )
+        if self.position_delta is None and self.condition_type == "subject":
+            self.position_delta = [0, -self.condition.size[0] // 16]
+        if self.position_delta is not None:
+            ids[:, 1] += self.position_delta[0]
+            ids[:, 2] += self.position_delta[1]
+        print(f"[Condition.encode] position_delta={self.position_delta}")
+        type_id = torch.ones_like(ids[:, :1]) * self.type_id
+        return tokens, ids, type_id

src/flux/generate.py

@@ -0,0 +1,838 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+import yaml, os
+from PIL import Image
+from diffusers.pipelines import FluxPipeline
+from typing import List, Union, Optional, Dict, Any, Callable
+from src.flux.transformer import tranformer_forward
+from src.flux.condition import Condition
+
+from diffusers.pipelines.flux.pipeline_flux import (
+    FluxPipelineOutput,
+    calculate_shift,
+    retrieve_timesteps,
+    np,
+)
+from src.flux.pipeline_tools import (
+    encode_prompt_with_clip_t5, tokenize_t5_prompt, clear_attn_maps, encode_vae_images
+)
+
+from src.flux.pipeline_tools import CustomFluxPipeline, load_modulation_adapter, decode_vae_images, \
+    save_attention_maps, gather_attn_maps, clear_attn_maps, load_dit_lora, quantization
+
+from src.utils.data_utils import pad_to_square, pad_to_target, pil2tensor, get_closest_ratio, get_aspect_ratios
+from src.utils.modulation_utils import get_word_index, unpad_input_ids
+
+def get_config(config_path: str = None):
+    config_path = config_path or os.environ.get("XFL_CONFIG")
+    if not config_path:
+        return {}
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+    return config
+
+
+def prepare_params(
+    prompt: Union[str, List[str]] = None,
+    prompt_2: Optional[Union[str, List[str]]] = None,
+    height: Optional[int] = 512,
+    width: Optional[int] = 512,
+    num_inference_steps: int = 28,
+    timesteps: List[int] = 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,
+    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,
+    verbose: bool = False,
+    **kwargs: dict,
+):
+    return (
+        prompt,
+        prompt_2,
+        height,
+        width,
+        num_inference_steps,
+        timesteps,
+        guidance_scale,
+        num_images_per_prompt,
+        generator,
+        latents,
+        prompt_embeds,
+        pooled_prompt_embeds,
+        output_type,
+        return_dict,
+        joint_attention_kwargs,
+        callback_on_step_end,
+        callback_on_step_end_tensor_inputs,
+        max_sequence_length,
+        verbose,
+    )
+
+
+def seed_everything(seed: int = 42):
+    torch.backends.cudnn.deterministic = True
+    torch.manual_seed(seed)
+    np.random.seed(seed)
+
+
+@torch.no_grad()
+def generate(
+    pipeline: FluxPipeline,
+    vae_conditions: List[Condition] = None,
+    config_path: str = None,
+    model_config: Optional[Dict[str, Any]] = {},
+    vae_condition_scale: float = 1.0,
+    default_lora: bool = False,
+    condition_pad_to: str = "square",
+    condition_size: int = 512,
+    text_cond_mask: Optional[torch.FloatTensor] = None,
+    delta_emb: Optional[torch.FloatTensor] = None,
+    delta_emb_pblock: Optional[torch.FloatTensor] = None,
+    delta_emb_mask: Optional[torch.FloatTensor] = None,
+    delta_start_ends = None,
+    condition_latents = None,
+    condition_ids = None,
+    mod_adapter = None,
+    store_attn_map: bool = False,
+    vae_skip_iter: str = None,
+    control_weight_lambda: str = None,
+    double_attention: bool = False,
+    single_attention: bool = False,
+    ip_scale: str = None,
+    use_latent_sblora_control: bool = False,
+    latent_sblora_scale: str = None,
+    use_condition_sblora_control: bool = False,
+    condition_sblora_scale: str = None,
+    idips = None,
+    **params: dict,
+):
+    model_config = model_config or get_config(config_path).get("model", {})
+
+    vae_skip_iter = model_config.get("vae_skip_iter", vae_skip_iter)
+    double_attention = model_config.get("double_attention", double_attention)
+    single_attention = model_config.get("single_attention", single_attention)
+    control_weight_lambda = model_config.get("control_weight_lambda", control_weight_lambda)
+    ip_scale = model_config.get("ip_scale", ip_scale)
+    use_latent_sblora_control = model_config.get("use_latent_sblora_control", use_latent_sblora_control)
+    use_condition_sblora_control = model_config.get("use_condition_sblora_control", use_condition_sblora_control)
+
+    latent_sblora_scale = model_config.get("latent_sblora_scale", latent_sblora_scale)
+    condition_sblora_scale = model_config.get("condition_sblora_scale", condition_sblora_scale)
+
+    model_config["use_attention_double"] = False
+    model_config["use_attention_single"] = False
+    use_attention = False
+    
+    if idips is not None:
+        if control_weight_lambda != "no":
+            parts = control_weight_lambda.split(',')
+            new_parts = []
+            for part in parts:
+                if ':' in part:
+                    left, right = part.split(':')
+                    values = right.split('/')
+                    # 保存整体值
+                    global_value = values[0]
+                    id_value = values[1]
+                    ip_value = values[2]
+                    new_values = [global_value]
+                    for is_id in idips:
+                        if is_id:
+                            new_values.append(id_value)
+                        else:
+                            new_values.append(ip_value)
+                    new_part = f"{left}:{('/'.join(new_values))}"
+                    new_parts.append(new_part)
+                else:
+                    new_parts.append(part)
+            control_weight_lambda = ','.join(new_parts)
+
+    if vae_condition_scale != 1:
+        for name, module in pipeline.transformer.named_modules():
+            if not name.endswith(".attn"):
+                continue
+            module.c_factor = torch.ones(1, 1) * vae_condition_scale
+
+    self = pipeline
+    (
+        prompt,
+        prompt_2,
+        height,
+        width,
+        num_inference_steps,
+        timesteps,
+        guidance_scale,
+        num_images_per_prompt,
+        generator,
+        latents,
+        prompt_embeds,
+        pooled_prompt_embeds,
+        output_type,
+        return_dict,
+        joint_attention_kwargs,
+        callback_on_step_end,
+        callback_on_step_end_tensor_inputs,
+        max_sequence_length,
+        verbose,
+    ) = prepare_params(**params)
+
+    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,
+        prompt_embeds=prompt_embeds,
+        pooled_prompt_embeds=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
+
+    lora_scale = (
+        self.joint_attention_kwargs.get("scale", None)
+        if self.joint_attention_kwargs is not None
+        else None
+    )
+    (
+        t5_prompt_embeds,
+        pooled_prompt_embeds,
+        text_ids,
+    ) = encode_prompt_with_clip_t5(
+        self=self,
+        prompt="" if self.text_encoder_2 is None else prompt,
+        prompt_2=None,
+        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,
+    )
+
+    # 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,
+        pooled_prompt_embeds.dtype,
+        device,
+        generator,
+        latents,
+    )
+
+    latent_height = height // 16
+
+    # 5. Prepare timesteps
+    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+    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,
+        timesteps,
+        sigmas,
+        mu=mu,
+    )
+    num_warmup_steps = max(
+        len(timesteps) - num_inference_steps * self.scheduler.order, 0
+    )
+    self._num_timesteps = len(timesteps)
+
+    attn_map = None
+
+    # 6. Denoising loop
+    with self.progress_bar(total=num_inference_steps) as progress_bar:
+        totalsteps = timesteps[0]
+        if control_weight_lambda is not None:
+            print("control_weight_lambda", control_weight_lambda)
+            control_weight_lambda_schedule = []
+            for scale_str in control_weight_lambda.split(','):
+                time_region, scale = scale_str.split(':')
+                start, end = time_region.split('-')
+                scales = [float(s) for s in scale.split('/')]
+                control_weight_lambda_schedule.append([(1-float(start))*totalsteps, (1-float(end))*totalsteps, scales])
+
+        if ip_scale is not None:
+            print("ip_scale", ip_scale)
+            ip_scale_schedule = []
+            for scale_str in ip_scale.split(','):
+                time_region, scale = scale_str.split(':')
+                start, end = time_region.split('-')
+                ip_scale_schedule.append([(1-float(start))*totalsteps, (1-float(end))*totalsteps, float(scale)])
+
+        if use_latent_sblora_control:
+            if latent_sblora_scale is not None:
+                print("latent_sblora_scale", latent_sblora_scale)
+                latent_sblora_scale_schedule = []
+                for scale_str in latent_sblora_scale.split(','):
+                    time_region, scale = scale_str.split(':')
+                    start, end = time_region.split('-')
+                    latent_sblora_scale_schedule.append([(1-float(start))*totalsteps, (1-float(end))*totalsteps, float(scale)])
+        
+        if use_condition_sblora_control:
+            if condition_sblora_scale is not None:
+                print("condition_sblora_scale", condition_sblora_scale)
+                condition_sblora_scale_schedule = []
+                for scale_str in condition_sblora_scale.split(','):
+                    time_region, scale = scale_str.split(':')
+                    start, end = time_region.split('-')
+                    condition_sblora_scale_schedule.append([(1-float(start))*totalsteps, (1-float(end))*totalsteps, float(scale)])
+
+
+        if vae_skip_iter is not None:
+            print("vae_skip_iter", vae_skip_iter)
+            vae_skip_iter_schedule = []
+            for scale_str in vae_skip_iter.split(','):
+                time_region, scale = scale_str.split(':')
+                start, end = time_region.split('-')
+                vae_skip_iter_schedule.append([(1-float(start))*totalsteps, (1-float(end))*totalsteps, float(scale)])
+
+        if control_weight_lambda is not None and attn_map is None:
+            batch_size = latents.shape[0]
+            latent_width = latents.shape[1]//latent_height
+            attn_map = torch.ones(batch_size, latent_height, latent_width, 128, device=latents.device, dtype=torch.bfloat16)
+            print("contol_weight_only", attn_map.shape)
+
+        self.scheduler.set_begin_index(0)
+        self.scheduler._init_step_index(0)
+        for i, t in enumerate(timesteps):
+            
+            if control_weight_lambda is not None:
+                cur_control_weight_lambda = []
+                for start, end, scale in control_weight_lambda_schedule:
+                    if t <= start and t >= end:
+                        cur_control_weight_lambda = scale
+                        break
+                print(f"timestep:{t}, cur_control_weight_lambda:{cur_control_weight_lambda}")
+               
+                if cur_control_weight_lambda:
+                    model_config["use_attention_single"] = True
+                    use_attention = True
+                    model_config["use_atten_lambda"] = cur_control_weight_lambda  
+                else:
+                    model_config["use_attention_single"] = False
+                    use_attention = False
+                     
+            if self.interrupt:
+                continue
+
+            if isinstance(delta_emb, list):
+                cur_delta_emb = delta_emb[i]
+                cur_delta_emb_pblock = delta_emb_pblock[i]
+                cur_delta_emb_mask = delta_emb_mask[i]
+            else:
+                cur_delta_emb = delta_emb
+                cur_delta_emb_pblock = delta_emb_pblock
+                cur_delta_emb_mask = delta_emb_mask
+
+
+            # broadcast to batch dimension in a way that's compatible with ONNX/Core ML
+            timestep = t.expand(latents.shape[0]).to(latents.dtype) / 1000
+            prompt_embeds = t5_prompt_embeds
+            text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=prompt_embeds.dtype)
+
+            # handle guidance
+            if self.transformer.config.guidance_embeds:
+                guidance = torch.tensor([guidance_scale], device=device)
+                guidance = guidance.expand(latents.shape[0])
+            else:
+                guidance = None
+            self.transformer.enable_lora()
+            
+            lora_weight  = 1
+            if ip_scale is not None:
+                lora_weight = 0
+                for start, end, scale in ip_scale_schedule:
+                    if t <= start and t >= end:
+                        lora_weight = scale
+                        break
+                if lora_weight != 1: print(f"timestep:{t}, lora_weights:{lora_weight}")
+            
+            latent_sblora_weight = None
+            if use_latent_sblora_control:
+                if latent_sblora_scale is not None:
+                    latent_sblora_weight = 0
+                    for start, end, scale in latent_sblora_scale_schedule:
+                        if t <= start and t >= end:
+                            latent_sblora_weight = scale
+                            break
+                    if latent_sblora_weight != 1: print(f"timestep:{t}, latent_sblora_weight:{latent_sblora_weight}")
+            
+            condition_sblora_weight = None
+            if use_condition_sblora_control:
+                if condition_sblora_scale is not None:
+                    condition_sblora_weight = 0
+                    for start, end, scale in condition_sblora_scale_schedule:
+                        if t <= start and t >= end:
+                            condition_sblora_weight = scale
+                            break
+                    if condition_sblora_weight !=1: print(f"timestep:{t}, condition_sblora_weight:{condition_sblora_weight}")
+
+            vae_skip_iter_t = False
+            if vae_skip_iter is not None:
+                for start, end, scale in vae_skip_iter_schedule:
+                    if t <= start and t >= end:
+                        vae_skip_iter_t = bool(scale)
+                        break
+                if vae_skip_iter_t:
+                    print(f"timestep:{t}, skip vae:{vae_skip_iter_t}")               
+
+            noise_pred = tranformer_forward(
+                self.transformer,
+                model_config=model_config,
+                # Inputs of the condition (new feature)
+                text_cond_mask=text_cond_mask,
+                delta_emb=cur_delta_emb,
+                delta_emb_pblock=cur_delta_emb_pblock,
+                delta_emb_mask=cur_delta_emb_mask,
+                delta_start_ends=delta_start_ends,
+                condition_latents=None if vae_skip_iter_t else condition_latents,
+                condition_ids=None if vae_skip_iter_t else condition_ids,
+                condition_type_ids=None,
+                # Inputs to the original transformer
+                hidden_states=latents,
+                # YiYi notes: divide it by 1000 for now because we scale it by 1000 in the transforme rmodel (we should not keep it but I want to keep the inputs same for the model for testing)
+                timestep=timestep,
+                guidance=guidance,
+                pooled_projections=pooled_prompt_embeds,
+                encoder_hidden_states=prompt_embeds,
+                txt_ids=text_ids,
+                img_ids=latent_image_ids,
+                joint_attention_kwargs={'scale': lora_weight, "latent_sblora_weight": latent_sblora_weight, "condition_sblora_weight": condition_sblora_weight}, 
+                store_attn_map=use_attention,
+                last_attn_map=attn_map if cur_control_weight_lambda else None,
+                use_text_mod=model_config["modulation"]["use_text_mod"],
+                use_img_mod=model_config["modulation"]["use_img_mod"],
+                mod_adapter=mod_adapter,
+                latent_height=latent_height,
+                return_dict=False,
+            )[0]
+
+            if use_attention:
+                attn_maps, _ = gather_attn_maps(self.transformer, clear=True)
+
+            # 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 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()
+
+    self.transformer.enable_lora()
+
+    if vae_condition_scale != 1:
+        for name, module in pipeline.transformer.named_modules():
+            if not name.endswith(".attn"):
+                continue
+            del module.c_factor
+
+    if not return_dict:
+        return (image,)
+
+    return FluxPipelineOutput(images=image)
+
+
+@torch.no_grad()
+def generate_from_test_sample(
+    test_sample, pipe, config, 
+    num_images=1, 
+    vae_skip_iter: str = None, 
+    target_height: int = None,
+    target_width: int = None,
+    seed: int = 42,
+    control_weight_lambda: str = None,
+    double_attention: bool = False,
+    single_attention: bool = False,
+    ip_scale: str = None,
+    use_latent_sblora_control: bool = False,
+    latent_sblora_scale: str = None,
+    use_condition_sblora_control: bool = False,
+    condition_sblora_scale: str = None,
+    use_idip = False,
+    **kargs
+):
+    target_size = config["train"]["dataset"]["val_target_size"]
+    condition_size = config["train"]["dataset"].get("val_condition_size", target_size//2)
+    condition_pad_to = config["train"]["dataset"]["condition_pad_to"]
+    pos_offset_type = config["model"].get("pos_offset_type", "width")
+    seed = config["model"].get("seed", seed)
+
+    device = pipe._execution_device
+
+    condition_imgs = test_sample['input_images']
+    position_delta = test_sample['position_delta']
+    prompt = test_sample['prompt']
+    original_image = test_sample.get('original_image', None)
+    condition_type = test_sample.get('condition_type', "subject")
+    modulation_input = test_sample.get('modulation', None)
+
+    delta_start_ends = None
+    condition_latents = condition_ids = None
+    text_cond_mask = None
+    
+    delta_embs = None
+    delta_embs_pblock = None
+    delta_embs_mask = None
+
+    try:
+        max_length = config["model"]["modulation"]["max_text_len"]
+    except Exception as e:
+        print(e)
+        max_length = 512
+
+    if modulation_input is None or len(modulation_input) == 0:
+        delta_emb = delta_emb_pblock = delta_emb_mask = None
+    else:
+        dtype = torch.bfloat16
+        batch_size = 1
+        N = config["model"]["modulation"].get("per_block_adapter_single_blocks", 0) + 19
+        guidance = torch.tensor([3.5]).to(device).expand(batch_size)
+        out_dim = config["model"]["modulation"]["out_dim"]
+
+        tar_text_inputs = tokenize_t5_prompt(pipe, prompt, max_length)
+        tar_padding_mask = tar_text_inputs.attention_mask.to(device).bool()
+        tar_tokens = tar_text_inputs.input_ids.to(device)
+        if config["model"]["modulation"]["eos_exclude"]:
+            tar_padding_mask[tar_tokens == 1] = False
+
+        def get_start_end_by_pompt_matching(src_prompts, tar_prompts):
+            text_cond_mask = torch.zeros(batch_size, max_length, device=device, dtype=torch.bool)
+            tar_prompt_input_ids = tokenize_t5_prompt(pipe, tar_prompts, max_length).input_ids
+            src_prompt_count = 1
+            start_ends = []
+            for i, (src_prompt, tar_prompt, tar_prompt_tokens) in enumerate(zip(src_prompts, tar_prompts, tar_prompt_input_ids)):
+                try:
+                    tar_start, tar_end = get_word_index(pipe, tar_prompt, tar_prompt_tokens, src_prompt, src_prompt_count, max_length, verbose=False)
+                    start_ends.append([tar_start, tar_end])
+                    text_cond_mask[i, tar_start:tar_end] = True
+                except Exception as e:
+                    print(e)
+            return start_ends, text_cond_mask
+
+        def encode_mod_image(pil_images):
+            if config["model"]["modulation"]["use_dit"]:
+                raise NotImplementedError()
+            else:
+                pil_images = [pad_to_square(img).resize((224, 224)) for img in pil_images]
+                if config["model"]["modulation"]["use_vae"]:
+                    raise NotImplementedError()
+                else:
+                    clip_pixel_values = pipe.clip_processor(
+                        text=None, images=pil_images, do_resize=False, do_center_crop=False, return_tensors="pt",
+                    ).pixel_values.to(dtype=dtype, device=device)
+                    clip_outputs = pipe.clip_model(clip_pixel_values, output_hidden_states=True, interpolate_pos_encoding=True, return_dict=True)
+                    return clip_outputs
+
+        def rgba_to_white_background(input_path, background=(255,255,255)):
+            with Image.open(input_path).convert("RGBA") as img:
+                img_np = np.array(img)
+                alpha = img_np[:, :, 3] / 255.0  # 归一化Alpha通道[3](@ref)
+                rgb = img_np[:, :, :3].astype(float)  # 提取RGB通道
+                
+                background_np = np.full_like(rgb, background, dtype=float)  # 根据参数生成背景[7](@ref)
+                
+                # 混合计算：前景色*alpha + 背景色*(1-alpha)
+                result_np = rgb * alpha[..., np.newaxis] + \
+                            background_np * (1 - alpha[..., np.newaxis])
+                
+                result = Image.fromarray(result_np.astype(np.uint8), "RGB")
+                return result
+        def get_mod_emb(modulation_input, timestep):
+            delta_emb = torch.zeros((batch_size, max_length, out_dim), dtype=dtype, device=device)
+            delta_emb_pblock = torch.zeros((batch_size, max_length, N, out_dim), dtype=dtype, device=device)
+            delta_emb_mask = torch.zeros((batch_size, max_length), dtype=torch.bool, device=device)
+            delta_start_ends = None
+            condition_latents = condition_ids = None
+            text_cond_mask = None
+
+            if modulation_input[0]["type"] == "adapter":
+                num_inputs = len(modulation_input[0]["src_inputs"])
+                src_prompts = [x["caption"] for x in modulation_input[0]["src_inputs"]]
+                src_text_inputs = tokenize_t5_prompt(pipe, src_prompts, max_length)
+                src_input_ids = unpad_input_ids(src_text_inputs.input_ids, src_text_inputs.attention_mask)
+                tar_input_ids = unpad_input_ids(tar_text_inputs.input_ids, tar_text_inputs.attention_mask)
+                src_prompt_embeds = pipe._get_t5_prompt_embeds(prompt=src_prompts, max_sequence_length=max_length, device=device) # (M, 512, 4096)
+                
+                pil_images = [rgba_to_white_background(x["image_path"]) for x in modulation_input[0]["src_inputs"]]
+
+                src_ds_scales = [x.get("downsample_scale", 1.0) for x in modulation_input[0]["src_inputs"]]
+                resized_pil_images = []
+                for img, ds_scale in zip(pil_images, src_ds_scales):
+                    img = pad_to_square(img)
+                    if ds_scale < 1.0:
+                        assert ds_scale > 0
+                        img = img.resize((int(224 * ds_scale), int(224 * ds_scale))).resize((224, 224))
+                    resized_pil_images.append(img)
+                pil_images = resized_pil_images
+                
+                img_encoded = encode_mod_image(pil_images)
+                delta_start_ends = []
+                text_cond_mask = torch.zeros(num_inputs, max_length, device=device, dtype=torch.bool)
+                if config["model"]["modulation"]["pass_vae"]:
+                    pil_images = [pad_to_square(img).resize((condition_size, condition_size)) for img in pil_images]
+                    with torch.no_grad():
+                        batch_tensor = torch.stack([pil2tensor(x) for x in pil_images])
+                        x_0, img_ids = encode_vae_images(pipe, batch_tensor) # (N, 256, 64)
+
+                    condition_latents = x_0.clone().detach().reshape(1, -1, 64) # (1, N256, 64)
+                    condition_ids = img_ids.clone().detach()
+                    condition_ids = condition_ids.unsqueeze(0).repeat_interleave(num_inputs, dim=0) # (N, 256, 3)
+                    for i in range(num_inputs):
+                        condition_ids[i, :, 1] += 0 if pos_offset_type == "width" else -(batch_tensor.shape[-1]//16) * (i + 1)
+                        condition_ids[i, :, 2] += -(batch_tensor.shape[-1]//16) * (i + 1)
+                    condition_ids = condition_ids.reshape(-1, 3) # (N256, 3)
+
+                if config["model"]["modulation"]["use_dit"]:
+                    raise NotImplementedError()
+                else:
+                    src_delta_embs = [] # [(512, 3072)]
+                    src_delta_emb_pblock = []
+                    for i in range(num_inputs):
+                        if isinstance(img_encoded, dict):
+                            _src_clip_outputs = {}
+                            for key in img_encoded:
+                                if torch.is_tensor(img_encoded[key]):
+                                    _src_clip_outputs[key] = img_encoded[key][i:i+1]
+                                else:
+                                    _src_clip_outputs[key] = [x[i:i+1] for x in img_encoded[key]]
+                            _img_encoded = _src_clip_outputs
+                        else:
+                            _img_encoded = img_encoded[i:i+1]
+                    
+                        x1, x2 = pipe.modulation_adapters[0](timestep, src_prompt_embeds[i:i+1], _img_encoded)
+                        src_delta_embs.append(x1[0]) # (512, 3072)
+                        src_delta_emb_pblock.append(x2[0]) # (512, N, 3072)
+
+                for input_args in modulation_input[0]["use_words"]:
+                    src_word_count = 1
+                    if len(input_args) == 3:
+                        src_input_index, src_word, tar_word = input_args
+                        tar_word_count = 1
+                    else:
+                        src_input_index, src_word, tar_word, tar_word_count = input_args[:4]
+                    src_prompt = src_prompts[src_input_index]
+                    tar_prompt = prompt
+
+                    src_start, src_end = get_word_index(pipe, src_prompt, src_input_ids[src_input_index], src_word, src_word_count, max_length, verbose=False)
+                    tar_start, tar_end = get_word_index(pipe, tar_prompt, tar_input_ids[0], tar_word, tar_word_count, max_length, verbose=False)
+                    if delta_emb is not None:
+                        delta_emb[:, tar_start:tar_end] = src_delta_embs[src_input_index][src_start:src_end] # (B, 512, 3072)
+                    if delta_emb_pblock is not None:
+                        delta_emb_pblock[:, tar_start:tar_end] = src_delta_emb_pblock[src_input_index][src_start:src_end] # (B, 512, N, 3072)
+                    delta_emb_mask[:, tar_start:tar_end] = True
+                    text_cond_mask[src_input_index, tar_start:tar_end] = True
+                    delta_start_ends.append([0, src_input_index, src_start, src_end, tar_start, tar_end])
+                text_cond_mask = text_cond_mask.transpose(0, 1).unsqueeze(0)
+
+            else:
+                raise NotImplementedError()
+            return delta_emb, delta_emb_pblock, delta_emb_mask, \
+                text_cond_mask, delta_start_ends, condition_latents, condition_ids
+    
+    num_inference_steps = 28 # FIXME: harcoded here
+    num_channels_latents = pipe.transformer.config.in_channels // 4
+
+    # set timesteps
+    sigmas = np.linspace(1.0, 1 / num_inference_steps, num_inference_steps)
+    mu = calculate_shift(
+        num_channels_latents,
+        pipe.scheduler.config.base_image_seq_len,
+        pipe.scheduler.config.max_image_seq_len,
+        pipe.scheduler.config.base_shift,
+        pipe.scheduler.config.max_shift,
+    )
+    timesteps, num_inference_steps = retrieve_timesteps(
+        pipe.scheduler,
+        num_inference_steps,
+        device,
+        None,
+        sigmas,
+        mu=mu,
+    )
+
+    if modulation_input is not None:
+        delta_embs = []
+        delta_embs_pblock = []
+        delta_embs_mask = []
+        for i, t in enumerate(timesteps):
+            t = t.expand(1).to(torch.bfloat16) / 1000
+            (
+                delta_emb, delta_emb_pblock, delta_emb_mask, 
+                text_cond_mask, delta_start_ends, 
+                condition_latents, condition_ids
+            ) = get_mod_emb(modulation_input, t)
+            delta_embs.append(delta_emb)
+            delta_embs_pblock.append(delta_emb_pblock)
+            delta_embs_mask.append(delta_emb_mask)
+
+    if original_image is not None:
+        raise NotImplementedError()
+        (target_height, target_width), closest_ratio = get_closest_ratio(original_image.height, original_image.width, train_aspect_ratios)
+    elif modulation_input is None or len(modulation_input) == 0:
+        delta_emb = delta_emb_pblock = delta_emb_mask = None
+    else:
+        for i, t in enumerate(timesteps):
+            t = t.expand(1).to(torch.bfloat16) / 1000
+            (
+                delta_emb, delta_emb_pblock, delta_emb_mask, 
+                text_cond_mask, delta_start_ends, 
+                condition_latents, condition_ids
+            ) = get_mod_emb(modulation_input, t)
+            delta_embs.append(delta_emb)
+            delta_embs_pblock.append(delta_emb_pblock)
+            delta_embs_mask.append(delta_emb_mask)
+
+    if target_height is None or target_width is None:
+        target_height = target_width = target_size
+
+    if condition_pad_to == "square":
+        condition_imgs = [pad_to_square(x) for x in condition_imgs]
+    elif condition_pad_to == "target":
+        condition_imgs = [pad_to_target(x, (target_size, target_size)) for x in condition_imgs]
+    condition_imgs = [x.resize((condition_size, condition_size)).convert("RGB") for x in condition_imgs]
+    # TODO: fix position_delta
+    conditions = [
+        Condition(
+            condition_type=condition_type,
+            condition=x,
+            position_delta=position_delta,
+        ) for x in condition_imgs
+    ]
+    # vlm_images = condition_imgs if config["model"]["use_vlm"] else []
+
+    use_perblock_adapter = False
+    try:
+        if config["model"]["modulation"]["use_perblock_adapter"]:
+            use_perblock_adapter = True
+    except Exception as e:
+        pass
+
+    results = []
+    for i in range(num_images):
+        clear_attn_maps(pipe.transformer)
+        generator = torch.Generator(device=device)
+        generator.manual_seed(seed + i)
+        if modulation_input is None or len(modulation_input) == 0:
+            idips = None
+        else:
+            idips = ["human" in p["image_path"] for p in modulation_input[0]["src_inputs"]]
+            if len(modulation_input[0]["use_words"][0])==5:
+                print("use idips in use_words")
+                idips = [x[-1] for x in modulation_input[0]["use_words"]]
+        result_img = generate(
+            pipe,
+            prompt=prompt,
+            max_sequence_length=max_length,
+            vae_conditions=conditions,
+            generator=generator,
+            model_config=config["model"],
+            height=target_height,
+            width=target_width,
+            condition_pad_to=condition_pad_to,
+            condition_size=condition_size,
+            text_cond_mask=text_cond_mask,
+            delta_emb=delta_embs,
+            delta_emb_pblock=delta_embs_pblock if use_perblock_adapter else None,
+            delta_emb_mask=delta_embs_mask,
+            delta_start_ends=delta_start_ends,
+            condition_latents=condition_latents,
+            condition_ids=condition_ids,
+            mod_adapter=pipe.modulation_adapters[0] if config["model"]["modulation"]["use_dit"] else None,
+            vae_skip_iter=vae_skip_iter,
+            control_weight_lambda=control_weight_lambda,
+            double_attention=double_attention,
+            single_attention=single_attention,
+            ip_scale=ip_scale,
+            use_latent_sblora_control=use_latent_sblora_control,
+            latent_sblora_scale=latent_sblora_scale,
+            use_condition_sblora_control=use_condition_sblora_control,
+            condition_sblora_scale=condition_sblora_scale,
+            idips=idips if use_idip else None,
+            **kargs,
+        ).images[0]
+
+        final_image = result_img
+        results.append(final_image)
+
+    if num_images == 1:
+        return results[0]
+    return results

src/flux/lora_controller.py

@@ -0,0 +1,99 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from peft.tuners.tuners_utils import BaseTunerLayer
+from typing import List, Any, Optional, Type
+
+
+class enable_lora:
+    def __init__(self, lora_modules: List[BaseTunerLayer], dit_activated: bool, cond_activated: bool=False, latent_sblora_weight: float=None, condition_sblora_weight: float=None) -> None:
+        self.dit_activated = dit_activated
+        self.cond_activated = cond_activated
+        self.latent_sblora_weight = latent_sblora_weight
+        self.condition_sblora_weight = condition_sblora_weight
+        # assert not (dit_activated and cond_activated)
+        
+        self.lora_modules: List[BaseTunerLayer] = [
+            each for each in lora_modules if isinstance(each, BaseTunerLayer)
+        ]
+
+        self.scales = [
+            {
+                active_adapter: lora_module.scaling[active_adapter] if active_adapter in lora_module.scaling else 1
+                for active_adapter in lora_module.active_adapters
+            } for lora_module in self.lora_modules
+        ]
+
+
+    def __enter__(self) -> None:
+        for i, lora_module in enumerate(self.lora_modules):
+            if not isinstance(lora_module, BaseTunerLayer):
+                continue
+            for active_adapter in lora_module.active_adapters:
+                if active_adapter == "default":
+                    if self.dit_activated:
+                        lora_module.scaling[active_adapter] = self.scales[0]["default"] if self.latent_sblora_weight is None else self.latent_sblora_weight
+                    else:
+                        lora_module.scaling[active_adapter] = 0
+                else:
+                    assert active_adapter == "condition"
+                    if self.cond_activated:
+                        lora_module.scaling[active_adapter] = self.scales[0]["condition"] if self.condition_sblora_weight is None else self.condition_sblora_weight
+                    else:
+                        lora_module.scaling[active_adapter] = 0
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[Any],
+    ) -> None:
+        for i, lora_module in enumerate(self.lora_modules):
+            if not isinstance(lora_module, BaseTunerLayer):
+                continue
+            for active_adapter in lora_module.active_adapters:
+                lora_module.scaling[active_adapter] = self.scales[i][active_adapter]
+        
+class set_lora_scale:
+    def __init__(self, lora_modules: List[BaseTunerLayer], scale: float) -> None:
+        self.lora_modules: List[BaseTunerLayer] = [
+            each for each in lora_modules if isinstance(each, BaseTunerLayer)
+        ]
+        self.scales = [
+            {
+                active_adapter: lora_module.scaling[active_adapter]
+                for active_adapter in lora_module.active_adapters
+            }
+            for lora_module in self.lora_modules
+        ]
+        self.scale = scale
+
+    def __enter__(self) -> None:
+        for lora_module in self.lora_modules:
+            if not isinstance(lora_module, BaseTunerLayer):
+                continue
+            lora_module.scale_layer(self.scale)
+
+    def __exit__(
+        self,
+        exc_type: Optional[Type[BaseException]],
+        exc_val: Optional[BaseException],
+        exc_tb: Optional[Any],
+    ) -> None:
+        for i, lora_module in enumerate(self.lora_modules):
+            if not isinstance(lora_module, BaseTunerLayer):
+                continue
+            for active_adapter in lora_module.active_adapters:
+                lora_module.scaling[active_adapter] = self.scales[i][active_adapter]

src/flux/pipeline_tools.py

@@ -0,0 +1,685 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import inspect
+from typing import Any, Callable, Dict, List, Optional, Union
+import os
+import torch
+from torch import Tensor
+import torch.nn.functional as F
+from diffusers.pipelines import FluxPipeline
+from diffusers.utils import logging
+from diffusers.loaders import TextualInversionLoaderMixin
+from diffusers.pipelines.flux.pipeline_flux import FluxLoraLoaderMixin
+from diffusers.models.transformers.transformer_flux import (
+    USE_PEFT_BACKEND,
+    scale_lora_layers,
+    unscale_lora_layers,
+    logger,
+)
+from torchvision.transforms import ToPILImage
+from peft.tuners.tuners_utils import BaseTunerLayer
+from optimum.quanto import (
+    freeze, quantize, QTensor, qfloat8, qint8, qint4, qint2,
+)
+import re
+import safetensors
+from src.adapters.mod_adapters import CLIPModAdapter
+from peft import LoraConfig, set_peft_model_state_dict
+from transformers import CLIPProcessor, CLIPModel, CLIPVisionModelWithProjection, CLIPVisionModel
+
+
+def encode_vae_images(pipeline: FluxPipeline, images: Tensor):
+    images = pipeline.image_processor.preprocess(images)
+    images = images.to(pipeline.device).to(pipeline.dtype)
+    images = pipeline.vae.encode(images).latent_dist.sample()
+    images = (
+        images - pipeline.vae.config.shift_factor
+    ) * pipeline.vae.config.scaling_factor
+    images_tokens = pipeline._pack_latents(images, *images.shape)
+    images_ids = pipeline._prepare_latent_image_ids(
+        images.shape[0],
+        images.shape[2],
+        images.shape[3],
+        pipeline.device,
+        pipeline.dtype,
+    )
+    if images_tokens.shape[1] != images_ids.shape[0]:
+        images_ids = pipeline._prepare_latent_image_ids(
+            images.shape[0],
+            images.shape[2] // 2,
+            images.shape[3] // 2,
+            pipeline.device,
+            pipeline.dtype,
+        )
+    return images_tokens, images_ids
+
+def decode_vae_images(pipeline: FluxPipeline, latents: Tensor, height, width, output_type: Optional[str] = "pil"):
+    latents = pipeline._unpack_latents(latents, height, width, pipeline.vae_scale_factor)
+    latents = (latents / pipeline.vae.config.scaling_factor) + pipeline.vae.config.shift_factor
+    image = pipeline.vae.decode(latents, return_dict=False)[0]
+    return pipeline.image_processor.postprocess(image, output_type=output_type)
+
+
+def _get_clip_prompt_embeds(
+    self,
+    prompt: Union[str, List[str]],
+    num_images_per_prompt: int = 1,
+    device: Optional[torch.device] = None,
+):
+    device = device or self._execution_device
+
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+    batch_size = len(prompt)
+
+    if isinstance(self, TextualInversionLoaderMixin):
+        prompt = self.maybe_convert_prompt(prompt, self.tokenizer)
+
+    text_inputs = self.tokenizer(
+        prompt,
+        padding="max_length",
+        max_length=self.tokenizer_max_length,
+        truncation=True,
+        return_overflowing_tokens=False,
+        return_length=False,
+        return_tensors="pt",
+    )
+
+    text_input_ids = text_inputs.input_ids
+
+    prompt_embeds = self.text_encoder(text_input_ids.to(device), output_hidden_states=False)
+
+    # Use pooled output of CLIPTextModel
+    prompt_embeds = prompt_embeds.pooler_output
+    prompt_embeds = prompt_embeds.to(dtype=self.text_encoder.dtype, device=device)
+
+    # duplicate text embeddings for each generation per prompt, using mps friendly method
+    prompt_embeds = prompt_embeds.repeat(1, num_images_per_prompt)
+    prompt_embeds = prompt_embeds.view(batch_size * num_images_per_prompt, -1)
+
+    return prompt_embeds
+
+def encode_prompt_with_clip_t5(
+    self,
+    prompt: Union[str, List[str]],
+    prompt_2: Union[str, List[str]],
+    device: Optional[torch.device] = None,
+    num_images_per_prompt: int = 1,
+    prompt_embeds: Optional[torch.FloatTensor] = None,
+    pooled_prompt_embeds: Optional[torch.FloatTensor] = None,
+    max_sequence_length: int = 512,
+    lora_scale: Optional[float] = None,
+):
+    r"""
+
+    Args:
+        prompt (`str` or `List[str]`, *optional*):
+            prompt to be encoded
+        prompt_2 (`str` or `List[str]`, *optional*):
+            The prompt or prompts to be sent to the `tokenizer_2` and `text_encoder_2`. If not defined, `prompt` is
+            used in all text-encoders
+        device: (`torch.device`):
+            torch device
+        num_images_per_prompt (`int`):
+            number of images that should be generated per prompt
+        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.
+        lora_scale (`float`, *optional*):
+            A lora scale that will be applied to all LoRA layers of the text encoder if LoRA layers are loaded.
+    """
+    device = device or self._execution_device
+
+    # set lora scale so that monkey patched LoRA
+    # function of text encoder can correctly access it
+    if lora_scale is not None and isinstance(self, FluxLoraLoaderMixin):
+        self._lora_scale = lora_scale
+
+        # dynamically adjust the LoRA scale
+        if self.text_encoder is not None and USE_PEFT_BACKEND:
+            scale_lora_layers(self.text_encoder, lora_scale)
+        if self.text_encoder_2 is not None and USE_PEFT_BACKEND:
+            scale_lora_layers(self.text_encoder_2, lora_scale)
+
+    prompt = [prompt] if isinstance(prompt, str) else prompt
+
+    if prompt_embeds is None:
+        prompt_2 = prompt_2 or prompt
+        prompt_2 = [prompt_2] if isinstance(prompt_2, str) else prompt_2
+
+        # We only use the pooled prompt output from the CLIPTextModel
+        pooled_prompt_embeds = _get_clip_prompt_embeds(
+            self=self,
+            prompt=prompt,
+            device=device,
+            num_images_per_prompt=num_images_per_prompt,
+        )
+        if self.text_encoder_2 is not None:
+            prompt_embeds = self._get_t5_prompt_embeds(
+                prompt=prompt_2,
+                num_images_per_prompt=num_images_per_prompt,
+                max_sequence_length=max_sequence_length,
+                device=device,
+            )
+            
+    if self.text_encoder is not None:
+        if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+            # Retrieve the original scale by scaling back the LoRA layers
+            unscale_lora_layers(self.text_encoder, lora_scale)
+
+    if self.text_encoder_2 is not None:
+        if isinstance(self, FluxLoraLoaderMixin) and USE_PEFT_BACKEND:
+            # Retrieve the original scale by scaling back the LoRA layers
+            unscale_lora_layers(self.text_encoder_2, lora_scale)
+
+    dtype = self.text_encoder.dtype if self.text_encoder is not None else self.transformer.dtype
+    if self.text_encoder_2 is not None:
+        text_ids = torch.zeros(prompt_embeds.shape[1], 3).to(device=device, dtype=dtype)
+    else:
+        text_ids = None
+
+    return prompt_embeds, pooled_prompt_embeds, text_ids
+
+
+
+def prepare_text_input(
+    pipeline: FluxPipeline, 
+    prompts, 
+    max_sequence_length=512,
+):
+    # Turn off warnings (CLIP overflow)
+    logger.setLevel(logging.ERROR)
+    (
+        t5_prompt_embeds,
+        pooled_prompt_embeds,
+        text_ids,
+    ) = encode_prompt_with_clip_t5(
+        self=pipeline,
+        prompt=prompts,
+        prompt_2=None,
+        prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        device=pipeline.device,
+        num_images_per_prompt=1,
+        max_sequence_length=max_sequence_length,
+        lora_scale=None,
+    )
+    # Turn on warnings
+    logger.setLevel(logging.WARNING)
+    return t5_prompt_embeds, pooled_prompt_embeds, text_ids
+
+def prepare_t5_input(
+    pipeline: FluxPipeline, 
+    prompts, 
+    max_sequence_length=512,
+):
+    # Turn off warnings (CLIP overflow)
+    logger.setLevel(logging.ERROR)
+    (
+        t5_prompt_embeds,
+        pooled_prompt_embeds,
+        text_ids,
+    ) = encode_prompt_with_clip_t5(
+        self=pipeline,
+        prompt=prompts,
+        prompt_2=None,
+        prompt_embeds=None,
+        pooled_prompt_embeds=None,
+        device=pipeline.device,
+        num_images_per_prompt=1,
+        max_sequence_length=max_sequence_length,
+        lora_scale=None,
+    )
+    # Turn on warnings
+    logger.setLevel(logging.WARNING)
+    return t5_prompt_embeds, pooled_prompt_embeds, text_ids
+
+def tokenize_t5_prompt(pipe, input_prompt, max_length, **kargs):
+    return pipe.tokenizer_2(
+        input_prompt,
+        padding="max_length",
+        max_length=max_length,
+        truncation=True,
+        return_length=False,
+        return_overflowing_tokens=False,
+        return_tensors="pt",
+        **kargs,
+    )
+
+def clear_attn_maps(transformer):
+    for i, block in enumerate(transformer.transformer_blocks):
+        if hasattr(block.attn, "attn_maps"):
+            del block.attn.attn_maps
+            del block.attn.timestep
+    for i, block in enumerate(transformer.single_transformer_blocks):
+        if hasattr(block.attn, "cond2latents"):
+            del block.attn.cond2latents
+
+def gather_attn_maps(transformer, clear=False):
+    t2i_attn_maps = {}
+    i2t_attn_maps = {}
+    for i, block in enumerate(transformer.transformer_blocks):
+        name = f"block_{i}"
+        if hasattr(block.attn, "attn_maps"):
+            attention_maps = block.attn.attn_maps
+            timesteps = block.attn.timestep # (B,)
+            for (timestep, (t2i_attn_map, i2t_attn_map)) in zip(timesteps, attention_maps):
+                timestep = str(timestep.item())
+
+                t2i_attn_maps[timestep] = t2i_attn_maps.get(timestep, dict())
+                t2i_attn_maps[timestep][name] = t2i_attn_maps[timestep].get(name, [])
+                t2i_attn_maps[timestep][name].append(t2i_attn_map.cpu())
+
+                i2t_attn_maps[timestep] = i2t_attn_maps.get(timestep, dict())
+                i2t_attn_maps[timestep][name] = i2t_attn_maps[timestep].get(name, [])
+                i2t_attn_maps[timestep][name].append(i2t_attn_map.cpu())
+
+        if clear:
+            del block.attn.attn_maps
+
+    for timestep in t2i_attn_maps:
+        for name in t2i_attn_maps[timestep]:
+            t2i_attn_maps[timestep][name] = torch.cat(t2i_attn_maps[timestep][name], dim=0)
+            i2t_attn_maps[timestep][name] = torch.cat(i2t_attn_maps[timestep][name], dim=0)
+
+    return t2i_attn_maps, i2t_attn_maps
+
+def process_token(token, startofword):
+    if '</w>' in token:
+        token = token.replace('</w>', '')
+        if startofword:
+            token = '<' + token + '>'
+        else:
+            token = '-' + token + '>'
+            startofword = True
+    elif token not in ['<|startoftext|>', '<|endoftext|>']:
+        if startofword:
+            token = '<' + token + '-'
+            startofword = False
+        else:
+            token = '-' + token + '-'
+    return token, startofword
+
+def save_attention_image(attn_map, tokens, batch_dir, to_pil):
+    startofword = True
+    for i, (token, a) in enumerate(zip(tokens, attn_map[:len(tokens)])):
+        token, startofword = process_token(token, startofword)
+        token = token.replace("/", "-")
+        if token == '-<pad>-':
+            continue
+        a = a.to(torch.float32)
+        a = a / a.max() * 255 / 256
+        to_pil(a).save(os.path.join(batch_dir, f'{i}-{token}.png'))
+
+def save_attention_maps(attn_maps, pipe, prompts, base_dir='attn_maps'):
+    to_pil = ToPILImage()
+    
+    token_ids = tokenize_t5_prompt(pipe, prompts, 512).input_ids # (B, 512)
+    token_ids = [x for x in token_ids]
+    total_tokens = [pipe.tokenizer_2.convert_ids_to_tokens(token_id) for token_id in token_ids]
+    
+    os.makedirs(base_dir, exist_ok=True)
+    
+    total_attn_map_shape = (256, 256)
+    total_attn_map_number = 0
+
+    # (B, 24, H, W, 512) -> (B, H, W, 512) -> (B, 512, H, W)
+    print(attn_maps.keys())
+    total_attn_map = list(list(attn_maps.values())[0].values())[0].sum(1)
+    total_attn_map = total_attn_map.permute(0, 3, 1, 2)
+    total_attn_map = torch.zeros_like(total_attn_map)
+    total_attn_map = F.interpolate(total_attn_map, size=total_attn_map_shape, mode='bilinear', align_corners=False)
+    
+    for timestep, layers in attn_maps.items():
+        timestep_dir = os.path.join(base_dir, f'{timestep}')
+        os.makedirs(timestep_dir, exist_ok=True)
+        
+        for layer, attn_map in layers.items():
+            layer_dir = os.path.join(timestep_dir, f'{layer}')
+            os.makedirs(layer_dir, exist_ok=True)
+            
+            attn_map = attn_map.sum(1).squeeze(1).permute(0, 3, 1, 2)
+
+            resized_attn_map = F.interpolate(attn_map, size=total_attn_map_shape, mode='bilinear', align_corners=False)
+            total_attn_map += resized_attn_map
+            total_attn_map_number += 1
+
+            for batch, (attn_map, tokens) in enumerate(zip(resized_attn_map, total_tokens)):
+                save_attention_image(attn_map, tokens, layer_dir, to_pil)
+            
+            # for batch, (tokens, attn) in enumerate(zip(total_tokens, attn_map)):
+            #     batch_dir = os.path.join(layer_dir, f'batch-{batch}')
+            #     os.makedirs(batch_dir, exist_ok=True)
+            #     save_attention_image(attn, tokens, batch_dir, to_pil)
+    
+    total_attn_map /= total_attn_map_number
+    for batch, (attn_map, tokens) in enumerate(zip(total_attn_map, total_tokens)):
+        batch_dir = os.path.join(base_dir, f'batch-{batch}')
+        os.makedirs(batch_dir, exist_ok=True)
+        save_attention_image(attn_map, tokens, batch_dir, to_pil)
+
+def gather_cond2latents(transformer, clear=False):
+    c2l_attn_maps = {}
+    # for i, block in enumerate(transformer.transformer_blocks):
+    for i, block in enumerate(transformer.single_transformer_blocks):
+        name = f"block_{i}"
+        if hasattr(block.attn, "cond2latents"):
+            attention_maps = block.attn.cond2latents
+            timesteps = block.attn.cond_timesteps # (B,)
+            for (timestep, c2l_attn_map) in zip(timesteps, attention_maps):
+                timestep = str(timestep.item())
+
+                c2l_attn_maps[timestep] = c2l_attn_maps.get(timestep, dict())
+                c2l_attn_maps[timestep][name] = c2l_attn_maps[timestep].get(name, [])
+                c2l_attn_maps[timestep][name].append(c2l_attn_map.cpu())
+
+            if clear:
+                # del block.attn.attn_maps
+                del block.attn.cond2latents
+                del block.attn.cond_timesteps
+
+    for timestep in c2l_attn_maps:
+        for name in c2l_attn_maps[timestep]:
+            c2l_attn_maps[timestep][name] = torch.cat(c2l_attn_maps[timestep][name], dim=0)
+
+    return c2l_attn_maps
+
+def save_cond2latent_image(attn_map, batch_dir, to_pil):
+    for i, a in enumerate(attn_map): # (N, H, W)
+        a = a.to(torch.float32)
+        a = a / a.max() * 255 / 256
+        to_pil(a).save(os.path.join(batch_dir, f'{i}.png'))
+
+def save_cond2latent(attn_maps, base_dir='attn_maps'):
+    to_pil = ToPILImage()
+    
+    os.makedirs(base_dir, exist_ok=True)
+    
+    total_attn_map_shape = (256, 256)
+    total_attn_map_number = 0
+
+    # (N, H, W) -> (1, N, H, W)
+    total_attn_map = list(list(attn_maps.values())[0].values())[0].unsqueeze(0)
+    total_attn_map = torch.zeros_like(total_attn_map)
+    total_attn_map = F.interpolate(total_attn_map, size=total_attn_map_shape, mode='bilinear', align_corners=False)
+    
+    for timestep, layers in attn_maps.items():
+        cur_ts_attn_map = torch.zeros_like(total_attn_map)
+        cur_ts_attn_map_number = 0
+
+        timestep_dir = os.path.join(base_dir, f'{timestep}')
+        os.makedirs(timestep_dir, exist_ok=True)
+        
+        for layer, attn_map in layers.items():
+            # layer_dir = os.path.join(timestep_dir, f'{layer}')
+            # os.makedirs(layer_dir, exist_ok=True)
+            
+            attn_map = attn_map.unsqueeze(0) # (1, N, H, W)
+            resized_attn_map = F.interpolate(attn_map, size=total_attn_map_shape, mode='bilinear', align_corners=False)
+
+            cur_ts_attn_map += resized_attn_map
+            cur_ts_attn_map_number += 1
+
+        for batch, attn_map in enumerate(cur_ts_attn_map / cur_ts_attn_map_number):
+            save_cond2latent_image(attn_map, timestep_dir, to_pil)
+        
+        total_attn_map += cur_ts_attn_map
+        total_attn_map_number += cur_ts_attn_map_number
+            
+    total_attn_map /= total_attn_map_number
+    for batch, attn_map in enumerate(total_attn_map):
+        batch_dir = os.path.join(base_dir, f'batch-{batch}')
+        os.makedirs(batch_dir, exist_ok=True)
+        save_cond2latent_image(attn_map, batch_dir, to_pil)
+
+def quantization(pipe, qtype):
+    if qtype != "None" and qtype != "":
+        if qtype.endswith("quanto"):
+            if qtype == "int2-quanto":
+                quant_level = qint2
+            elif qtype == "int4-quanto":
+                quant_level = qint4
+            elif qtype == "int8-quanto":
+                quant_level = qint8
+            elif qtype == "fp8-quanto":
+                quant_level = qfloat8
+            else:
+                raise ValueError(f"Invalid quantisation level: {qtype}")
+
+            extra_quanto_args = {}
+            extra_quanto_args["exclude"] = [
+                "*.norm",
+                "*.norm1",
+                "*.norm2",
+                "*.norm2_context",
+                "proj_out",
+                "x_embedder",
+                "norm_out",
+                "context_embedder",
+            ]
+            try:
+                quantize(pipe.transformer, weights=quant_level, **extra_quanto_args)
+                quantize(pipe.text_encoder_2, weights=quant_level, **extra_quanto_args)
+                print("[Quantization] Start freezing")
+                freeze(pipe.transformer)
+                freeze(pipe.text_encoder_2)
+                print("[Quantization] Finished")
+            except Exception as e:
+                if "out of memory" in str(e).lower():
+                    print(
+                        "GPU ran out of memory during quantisation. Use --quantize_via=cpu to use the slower CPU method."
+                    )
+                raise e
+        else:
+            assert qtype == "fp8-ao"
+            from torchao.float8 import convert_to_float8_training, Float8LinearConfig
+            def module_filter_fn(mod: torch.nn.Module, fqn: str):
+                # don't convert the output module
+                if fqn == "proj_out":
+                    return False
+                # don't convert linear modules with weight dimensions not divisible by 16
+                if isinstance(mod, torch.nn.Linear):
+                    if mod.in_features % 16 != 0 or mod.out_features % 16 != 0:
+                        return False
+                return True
+            convert_to_float8_training(
+                pipe.transformer, module_filter_fn=module_filter_fn, config=Float8LinearConfig(pad_inner_dim=True)
+            )
+
+class CustomFluxPipeline:
+    def __init__(
+        self,
+        config,
+        device="cuda",
+        ckpt_root=None,
+        ckpt_root_condition=None,
+        torch_dtype=torch.bfloat16,
+    ):
+        model_path = os.getenv("FLUX_MODEL_PATH", "black-forest-labs/FLUX.1-dev")
+        print("[CustomFluxPipeline] Loading FLUX Pipeline")
+        self.pipe = FluxPipeline.from_pretrained(model_path, torch_dtype=torch_dtype).to(device)
+
+        self.config = config
+        self.device = device
+        self.dtype = torch_dtype
+        if config["model"].get("dit_quant", "None") != "None":
+            quantization(self.pipe, config["model"]["dit_quant"])
+
+        self.modulation_adapters = []
+        self.pipe.modulation_adapters = []
+        
+        try:
+            if config["model"]["modulation"]["use_clip"]:
+                load_clip(self, config, torch_dtype, device, None, is_training=False)
+        except Exception as e:
+            print(e)
+        
+        if config["model"]["use_dit_lora"] or config["model"]["use_condition_dblock_lora"] or config["model"]["use_condition_sblock_lora"]:
+            if ckpt_root_condition is None and (config["model"]["use_condition_dblock_lora"] or config["model"]["use_condition_sblock_lora"]):
+                ckpt_root_condition = ckpt_root
+            load_dit_lora(self, self.pipe, config, torch_dtype, device, f"{ckpt_root}", f"{ckpt_root_condition}", is_training=False)
+
+    def add_modulation_adapter(self, modulation_adapter):
+        self.modulation_adapters.append(modulation_adapter)
+        self.pipe.modulation_adapters.append(modulation_adapter)
+
+    def clear_modulation_adapters(self):
+        self.modulation_adapters = []
+        self.pipe.modulation_adapters = []
+        torch.cuda.empty_cache()
+
+def load_clip(self, config, torch_dtype, device, ckpt_dir=None, is_training=False):
+    model_path = os.getenv("CLIP_MODEL_PATH", "openai/clip-vit-large-patch14")
+    clip_model = CLIPVisionModelWithProjection.from_pretrained(model_path).to(device, dtype=torch_dtype)
+    clip_processor = CLIPProcessor.from_pretrained(model_path)
+    self.pipe.clip_model = clip_model
+    self.pipe.clip_processor = clip_processor
+
+def load_dit_lora(self, pipe, config, torch_dtype, device, ckpt_dir=None, condition_ckpt_dir=None, is_training=False):
+
+    if not config["model"]["use_condition_dblock_lora"] and not config["model"]["use_condition_sblock_lora"] and not config["model"]["use_dit_lora"]:
+        print("[load_dit_lora] no dit lora, no condition lora")
+        return []
+    
+    adapter_names = ["default", "condition"]
+    
+    if condition_ckpt_dir is None:
+        condition_ckpt_dir = ckpt_dir
+    
+    if not config["model"]["use_condition_dblock_lora"] and not config["model"]["use_condition_sblock_lora"]:
+        print("[load_dit_lora] no condition lora")
+        adapter_names.pop(1)
+    elif condition_ckpt_dir is not None and os.path.exists(os.path.join(condition_ckpt_dir, "pytorch_lora_weights_condition.safetensors")):
+        assert "condition" in adapter_names
+        print(f"[load_dit_lora] load condition lora from {condition_ckpt_dir}")
+        pipe.transformer.load_lora_adapter(condition_ckpt_dir, use_safetensors=True, adapter_name="condition", weight_name="pytorch_lora_weights_condition.safetensors") # TODO: check if they are trainable
+    else:
+        assert is_training
+        assert "condition" in adapter_names
+        print("[load_dit_lora] init new condition lora")
+        pipe.transformer.add_adapter(LoraConfig(**config["model"]["condition_lora_config"]), adapter_name="condition")
+
+    if not config["model"]["use_dit_lora"]:
+        print("[load_dit_lora] no dit lora")
+        adapter_names.pop(0)
+    elif ckpt_dir is not None and os.path.exists(os.path.join(ckpt_dir, "pytorch_lora_weights.safetensors")):
+        assert "default" in adapter_names
+        print(f"[load_dit_lora] load dit lora from {ckpt_dir}")
+        lora_file = os.path.join(ckpt_dir, "pytorch_lora_weights.safetensors")
+        lora_state_dict = safetensors.torch.load_file(lora_file, device="cpu")
+        
+        single_lora_pattern = "(.*single_transformer_blocks\\.[0-9]+\\.norm\\.linear|.*single_transformer_blocks\\.[0-9]+\\.proj_mlp|.*single_transformer_blocks\\.[0-9]+\\.proj_out|.*single_transformer_blocks\\.[0-9]+\\.attn.to_k|.*single_transformer_blocks\\.[0-9]+\\.attn.to_q|.*single_transformer_blocks\\.[0-9]+\\.attn.to_v|.*single_transformer_blocks\\.[0-9]+\\.attn.to_out)"
+        latent_lora_pattern = "(.*(?<!single_)transformer_blocks\\.[0-9]+\\.norm1\\.linear|.*(?<!single_)transformer_blocks\\.[0-9]+\\.attn\\.to_k|.*(?<!single_)transformer_blocks\\.[0-9]+\\.attn\\.to_q|.*(?<!single_)transformer_blocks\\.[0-9]+\\.attn\\.to_v|.*(?<!single_)transformer_blocks\\.[0-9]+\\.attn\\.to_out\\.0|.*(?<!single_)transformer_blocks\\.[0-9]+\\.ff\\.net\\.2)"
+        use_pretrained_dit_single_lora = config["model"].get("use_pretrained_dit_single_lora", True)
+        use_pretrained_dit_latent_lora = config["model"].get("use_pretrained_dit_latent_lora", True)
+        if not use_pretrained_dit_single_lora or not use_pretrained_dit_latent_lora:
+            lora_state_dict_keys = list(lora_state_dict.keys())
+            for layer_name in lora_state_dict_keys:
+                if not use_pretrained_dit_single_lora:
+                    if re.search(single_lora_pattern, layer_name):
+                        del lora_state_dict[layer_name]
+                if not use_pretrained_dit_latent_lora:
+                    if re.search(latent_lora_pattern, layer_name):
+                        del lora_state_dict[layer_name]
+            pipe.transformer.add_adapter(LoraConfig(**config["model"]["dit_lora_config"]), adapter_name="default")
+            set_peft_model_state_dict(pipe.transformer, lora_state_dict, adapter_name="default")
+        else:
+            pipe.transformer.load_lora_adapter(ckpt_dir, use_safetensors=True, adapter_name="default", weight_name="pytorch_lora_weights.safetensors") # TODO: check if they are trainable
+    else:
+        assert is_training
+        assert "default" in adapter_names
+        print("[load_dit_lora] init new dit lora")
+        pipe.transformer.add_adapter(LoraConfig(**config["model"]["dit_lora_config"]), adapter_name="default")
+
+    assert len(adapter_names) <= 2 and len(adapter_names) > 0
+    for name, module in pipe.transformer.named_modules():
+        if isinstance(module, BaseTunerLayer):
+            module.set_adapter(adapter_names)
+    
+    if "default" in adapter_names: assert config["model"]["use_dit_lora"]
+    if "condition" in adapter_names: assert config["model"]["use_condition_dblock_lora"] or config["model"]["use_condition_sblock_lora"]
+    
+    lora_layers = list(filter(
+        lambda p: p[1].requires_grad, pipe.transformer.named_parameters()
+    ))
+    
+    lora_layers = [l[1] for l in lora_layers]
+    return lora_layers
+
+def load_modulation_adapter(self, config, torch_dtype, device, ckpt_dir=None, is_training=False):
+    adapter_type = config["model"]["modulation"]["adapter_type"]
+
+    if ckpt_dir is not None and os.path.exists(ckpt_dir):
+        print(f"loading modulation adapter from {ckpt_dir}")
+        modulation_adapter = CLIPModAdapter.from_pretrained(
+            ckpt_dir, subfolder="modulation_adapter", strict=False, 
+            low_cpu_mem_usage=False, device_map=None,
+        ).to(device)
+    else:
+        print(f"Init new modulation adapter")
+        adapter_layers = config["model"]["modulation"]["adapter_layers"]
+        adapter_width = config["model"]["modulation"]["adapter_width"]
+        pblock_adapter_layers = config["model"]["modulation"]["per_block_adapter_layers"]
+        pblock_adapter_width = config["model"]["modulation"]["per_block_adapter_width"]
+        pblock_adapter_single_blocks = config["model"]["modulation"]["per_block_adapter_single_blocks"]
+        use_text_mod = config["model"]["modulation"]["use_text_mod"]
+        use_img_mod = config["model"]["modulation"]["use_img_mod"]
+        
+        out_dim = config["model"]["modulation"]["out_dim"]
+        if adapter_type == "clip_adapter":
+            modulation_adapter = CLIPModAdapter(
+                out_dim=out_dim,
+                width=adapter_width,
+                pblock_width=pblock_adapter_width,
+                layers=adapter_layers,
+                pblock_layers=pblock_adapter_layers,
+                heads=8,
+                input_text_dim=4096,
+                input_image_dim=1024,
+                pblock_single_blocks=pblock_adapter_single_blocks,
+            )
+        else:
+            raise NotImplementedError()
+
+    if is_training:
+        modulation_adapter.train()
+        try:
+            modulation_adapter.enable_gradient_checkpointing()
+        except Exception as e:
+            print(e)
+        if not config["model"]["modulation"]["use_perblock_adapter"]:
+            try:
+                modulation_adapter.net2.requires_grad_(False)
+            except Exception as e:
+                print(e)
+    else:
+        modulation_adapter.requires_grad_(False)
+
+    modulation_adapter.to(device, dtype=torch_dtype)
+    return modulation_adapter
+
+
+def load_ckpt(self, ckpt_dir, is_training=False):
+    if self.config["model"]["use_dit_lora"]:
+        self.pipe.transformer.delete_adapters(["subject"])
+        lora_path = f"{ckpt_dir}/pytorch_lora_weights.safetensors"
+        print(f"Loading DIT Lora from {lora_path}")
+        self.pipe.load_lora_weights(lora_path, adapter_name="subject")
+
+        
+

src/flux/transformer.py

@@ -0,0 +1,361 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# Copyright 2024 Black Forest Labs and The HuggingFace Team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from diffusers.pipelines import FluxPipeline
+from typing import List, Union, Optional, Dict, Any, Callable
+from .block import block_forward, single_block_forward
+from .lora_controller import enable_lora
+from diffusers.models.transformers.transformer_flux import (
+    FluxTransformer2DModel,
+    Transformer2DModelOutput,
+    USE_PEFT_BACKEND,
+    is_torch_version,
+    scale_lora_layers,
+    unscale_lora_layers,
+    logger,
+)
+import numpy as np
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+def prepare_params(
+    hidden_states: torch.Tensor,
+    encoder_hidden_states: torch.Tensor = None,
+    pooled_projections: torch.Tensor = None,
+    timestep: torch.LongTensor = None,
+    img_ids: torch.Tensor = None,
+    txt_ids: torch.Tensor = None,
+    guidance: torch.Tensor = None,
+    joint_attention_kwargs: Optional[Dict[str, Any]] = None,
+    controlnet_block_samples=None,
+    controlnet_single_block_samples=None,
+    return_dict: bool = True,
+    **kwargs: dict,
+):
+    return (
+        hidden_states,
+        encoder_hidden_states,
+        pooled_projections,
+        timestep,
+        img_ids,
+        txt_ids,
+        guidance,
+        joint_attention_kwargs,
+        controlnet_block_samples,
+        controlnet_single_block_samples,
+        return_dict,
+    )
+
+
+def tranformer_forward(
+    transformer: FluxTransformer2DModel,
+    condition_latents: torch.Tensor,
+    condition_ids: torch.Tensor,
+    condition_type_ids: torch.Tensor,
+    model_config: Optional[Dict[str, Any]] = {},
+    c_t=0,
+    text_cond_mask: Optional[torch.FloatTensor] = None,
+    delta_emb: Optional[torch.FloatTensor] = None,
+    delta_emb_pblock: Optional[torch.FloatTensor] = None,
+    delta_emb_mask: Optional[torch.FloatTensor] = None,
+    delta_start_ends = None,
+    store_attn_map: bool = False,
+    use_text_mod: bool = True,
+    use_img_mod: bool = False,
+    mod_adapter = None,
+    latent_height: Optional[int] = None,
+    last_attn_map = None,
+    **params: dict,
+):
+    self = transformer
+    use_condition = condition_latents is not None
+
+    (
+        hidden_states,
+        encoder_hidden_states,
+        pooled_projections,
+        timestep,
+        img_ids,
+        txt_ids,
+        guidance,
+        joint_attention_kwargs,
+        controlnet_block_samples,
+        controlnet_single_block_samples,
+        return_dict,
+    ) = prepare_params(**params)
+
+    if joint_attention_kwargs is not None:
+        joint_attention_kwargs = joint_attention_kwargs.copy()
+        lora_scale = joint_attention_kwargs.pop("scale", 1.0)
+        latent_sblora_weight = joint_attention_kwargs.pop("latent_sblora_weight", None)
+        condition_sblora_weight = joint_attention_kwargs.pop("condition_sblora_weight", None)
+    else:
+        lora_scale = 1.0
+        latent_sblora_weight = None
+        condition_sblora_weight = None
+    if USE_PEFT_BACKEND:
+        # weight the lora layers by setting `lora_scale` for each PEFT layer
+        scale_lora_layers(self, lora_scale)
+    else:
+        if (
+            joint_attention_kwargs is not None
+            and joint_attention_kwargs.get("scale", None) is not None
+        ):
+            logger.warning(
+                "Passing `scale` via `joint_attention_kwargs` when not using the PEFT backend is ineffective."
+            )
+
+    train_partial_text_lora = model_config.get("train_partial_text_lora", False)
+    train_partial_latent_lora = model_config.get("train_partial_latent_lora", False)
+
+    if train_partial_text_lora or train_partial_latent_lora:
+        train_partial_text_lora_layers = model_config.get("train_partial_text_lora_layers", "")
+        train_partial_latent_lora_layers = model_config.get("train_partial_latent_lora_layers", "")
+        activate_x_embedder = True
+        if "x_embedder" not in train_partial_text_lora_layers or "x_embedder" not in train_partial_latent_lora_layers:
+            activate_x_embedder = False
+    if train_partial_text_lora or train_partial_latent_lora:
+        activate_x_embedder_ = activate_x_embedder
+    else:
+        activate_x_embedder_ = model_config["latent_lora"] or model_config["text_lora"]
+    
+    with enable_lora((self.x_embedder,), activate_x_embedder_):
+        hidden_states = self.x_embedder(hidden_states)
+    cond_lora_activate = model_config["use_condition_dblock_lora"] or model_config["use_condition_sblock_lora"]
+    with enable_lora(
+        (self.x_embedder,), 
+        dit_activated=activate_x_embedder if train_partial_text_lora or train_partial_latent_lora else not cond_lora_activate, cond_activated=cond_lora_activate,
+    ):
+        condition_latents = self.x_embedder(condition_latents) if use_condition else None
+
+    timestep = timestep.to(hidden_states.dtype) * 1000
+
+    if guidance is not None:
+        guidance = guidance.to(hidden_states.dtype) * 1000
+    else:
+        guidance = None
+
+    temb = (
+        self.time_text_embed(timestep, pooled_projections)
+        if guidance is None
+        else self.time_text_embed(timestep, guidance, pooled_projections)
+    ) # (B, 3072)
+
+    cond_temb = (
+        self.time_text_embed(torch.ones_like(timestep) * c_t * 1000, pooled_projections)
+        if guidance is None
+        else self.time_text_embed(
+            torch.ones_like(timestep) * c_t * 1000, guidance, pooled_projections
+        )
+    )
+    encoder_hidden_states = self.context_embedder(encoder_hidden_states)
+
+    if txt_ids.ndim == 3:
+        logger.warning(
+            "Passing `txt_ids` 3d torch.Tensor is deprecated."
+            "Please remove the batch dimension and pass it as a 2d torch Tensor"
+        )
+        txt_ids = txt_ids[0]
+    if img_ids.ndim == 3:
+        logger.warning(
+            "Passing `img_ids` 3d torch.Tensor is deprecated."
+            "Please remove the batch dimension and pass it as a 2d torch Tensor"
+        )
+        img_ids = img_ids[0]
+
+    ids = torch.cat((txt_ids, img_ids), dim=0)
+    image_rotary_emb = self.pos_embed(ids)
+    if use_condition:
+        cond_rotary_emb = self.pos_embed(condition_ids)
+
+    for index_block, block in enumerate(self.transformer_blocks):
+        if delta_emb_pblock is None:
+            delta_emb_cblock = None
+        else:
+            delta_emb_cblock = delta_emb_pblock[:, :, index_block]
+        condition_pass_to_double = use_condition and (model_config["double_use_condition"] or model_config["single_use_condition"])
+        if self.training and self.gradient_checkpointing:
+            ckpt_kwargs: Dict[str, Any] = (
+                {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+            )
+            
+            encoder_hidden_states, hidden_states, condition_latents = (
+                torch.utils.checkpoint.checkpoint(
+                    block_forward,
+                    self=block,
+                    model_config=model_config,
+                    hidden_states=hidden_states,
+                    encoder_hidden_states=encoder_hidden_states,
+                    condition_latents=condition_latents if condition_pass_to_double else None,
+                    cond_temb=cond_temb if condition_pass_to_double else None,
+                    cond_rotary_emb=cond_rotary_emb if condition_pass_to_double else None,
+                    temb=temb,
+                    text_cond_mask=text_cond_mask,
+                    delta_emb=delta_emb,
+                    delta_emb_cblock=delta_emb_cblock,
+                    delta_emb_mask=delta_emb_mask,
+                    delta_start_ends=delta_start_ends,
+                    image_rotary_emb=image_rotary_emb,
+                    store_attn_map=store_attn_map,
+                    use_text_mod=use_text_mod,
+                    use_img_mod=use_img_mod,
+                    mod_adapter=mod_adapter,
+                    latent_height=latent_height,
+                    timestep=timestep,
+                    last_attn_map=last_attn_map,
+                    **ckpt_kwargs,
+                )
+            )
+
+        else:
+            encoder_hidden_states, hidden_states, condition_latents = block_forward(
+                block,
+                model_config=model_config,
+                hidden_states=hidden_states,
+                encoder_hidden_states=encoder_hidden_states,
+                condition_latents=condition_latents if condition_pass_to_double else None,
+                cond_temb=cond_temb if condition_pass_to_double else None,
+                cond_rotary_emb=cond_rotary_emb if condition_pass_to_double else None,
+                temb=temb,
+                text_cond_mask=text_cond_mask,
+                delta_emb=delta_emb,
+                delta_emb_cblock=delta_emb_cblock,
+                delta_emb_mask=delta_emb_mask,
+                delta_start_ends=delta_start_ends,
+                image_rotary_emb=image_rotary_emb,
+                store_attn_map=store_attn_map,
+                use_text_mod=use_text_mod,
+                use_img_mod=use_img_mod,
+                mod_adapter=mod_adapter,
+                latent_height=latent_height,
+                timestep=timestep,
+                last_attn_map=last_attn_map,
+            )
+
+        # controlnet residual
+        if controlnet_block_samples is not None:
+            interval_control = len(self.transformer_blocks) / len(
+                controlnet_block_samples
+            )
+            interval_control = int(np.ceil(interval_control))
+            hidden_states = (
+                hidden_states
+                + controlnet_block_samples[index_block // interval_control]
+            )
+    hidden_states = torch.cat([encoder_hidden_states, hidden_states], dim=1)
+
+    for index_block, block in enumerate(self.single_transformer_blocks):
+        if delta_emb_pblock is not None and delta_emb_pblock.shape[2] > 19+index_block:
+            delta_emb_single = delta_emb
+            delta_emb_cblock = delta_emb_pblock[:, :, index_block+19]
+        else:
+            delta_emb_single = None
+            delta_emb_cblock = None
+        if self.training and self.gradient_checkpointing:
+            ckpt_kwargs: Dict[str, Any] = (
+                {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
+            )
+            result = torch.utils.checkpoint.checkpoint(
+                single_block_forward,
+                self=block,
+                model_config=model_config,
+                hidden_states=hidden_states,
+                temb=temb,
+                delta_emb=delta_emb_single,
+                delta_emb_cblock=delta_emb_cblock,
+                delta_emb_mask=delta_emb_mask,
+                use_text_mod=use_text_mod,
+                use_img_mod=use_img_mod,
+                image_rotary_emb=image_rotary_emb,
+                last_attn_map=last_attn_map,
+                latent_height=latent_height,
+                timestep=timestep,
+                store_attn_map=store_attn_map,
+                **(
+                    {
+                        "condition_latents": condition_latents,
+                        "cond_temb": cond_temb,
+                        "cond_rotary_emb": cond_rotary_emb,
+                        "text_cond_mask": text_cond_mask,
+                    }
+                    if use_condition and model_config["single_use_condition"]
+                    else {}
+                ),
+                **ckpt_kwargs,
+            )
+
+        else:
+            result = single_block_forward(
+                block,
+                model_config=model_config,
+                hidden_states=hidden_states,
+                temb=temb,
+                delta_emb=delta_emb_single,
+                delta_emb_cblock=delta_emb_cblock,
+                delta_emb_mask=delta_emb_mask,
+                use_text_mod=use_text_mod,
+                use_img_mod=use_img_mod,
+                image_rotary_emb=image_rotary_emb,
+                last_attn_map=last_attn_map,
+                latent_height=latent_height,
+                timestep=timestep,
+                store_attn_map=store_attn_map,
+                latent_sblora_weight=latent_sblora_weight,
+                condition_sblora_weight=condition_sblora_weight,
+                **(
+                    {
+                        "condition_latents": condition_latents,
+                        "cond_temb": cond_temb,
+                        "cond_rotary_emb": cond_rotary_emb,
+                        "text_cond_mask": text_cond_mask,
+                    }
+                    if use_condition and model_config["single_use_condition"]
+                    else {}
+                ),
+            )
+        if use_condition and model_config["single_use_condition"]:
+            hidden_states, condition_latents = result
+        else:
+            hidden_states = result
+
+        # controlnet residual
+        if controlnet_single_block_samples is not None:
+            interval_control = len(self.single_transformer_blocks) / len(
+                controlnet_single_block_samples
+            )
+            interval_control = int(np.ceil(interval_control))
+            hidden_states[:, encoder_hidden_states.shape[1] :, ...] = (
+                hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+                + controlnet_single_block_samples[index_block // interval_control]
+            )
+
+    hidden_states = hidden_states[:, encoder_hidden_states.shape[1] :, ...]
+
+    hidden_states = self.norm_out(hidden_states, temb)
+    output = self.proj_out(hidden_states)
+
+    if USE_PEFT_BACKEND:
+        # remove `lora_scale` from each PEFT layer
+        unscale_lora_layers(self, lora_scale)
+
+    if not return_dict:
+        return (output,)
+    return Transformer2DModelOutput(sample=output)
+

src/utils/data_utils.py

@@ -0,0 +1,404 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import cv2
+import json
+import torch
+import random
+import base64
+import numpy as np
+from PIL import Image, ImageDraw
+from glob import glob
+from torchvision import transforms as T
+import os
+import gc
+from webdataset.filters import default_collation_fn, pipelinefilter
+import yaml
+
+def get_rank_and_worldsize():
+    try:
+        local_rank = int(os.environ.get("LOCAL_RANK"))
+        global_rank = int(os.environ.get("RANK"))
+        world_size = int(os.getenv('WORLD_SIZE', 1))
+    except:
+        local_rank = 0
+        global_rank = 0
+        world_size = 1
+    return local_rank, global_rank, world_size
+
+def get_train_config(config_path=None):
+    if config_path is None:
+        config_path = os.environ.get("XFL_CONFIG")
+    assert config_path is not None, "Please set the XFL_CONFIG environment variable"
+    with open(config_path, "r") as f:
+        config = yaml.safe_load(f)
+    return config
+
+def calculate_aspect_ratios(resolution):
+    ASPECT_RATIO = {
+        '0.25': [128.0, 512.0], '0.26': [128.0, 496.0], '0.27': [128.0, 480.0], '0.28': [128.0, 464.0],
+        '0.32': [144.0, 448.0], '0.33': [144.0, 432.0], '0.35': [144.0, 416.0], '0.4': [160.0, 400.0],
+        '0.42': [160.0, 384.0], '0.48': [176.0, 368.0], '0.5': [176.0, 352.0], '0.52': [176.0, 336.0],
+        '0.57': [192.0, 336.0], '0.6': [192.0, 320.0], '0.68': [208.0, 304.0], '0.72': [208.0, 288.0],
+        '0.78': [224.0, 288.0], '0.82': [224.0, 272.0], '0.88': [240.0, 272.0], '0.94': [240.0, 256.0],
+        '1.0': [256.0, 256.0], '1.07': [256.0, 240.0], '1.13': [272.0, 240.0], '1.21': [272.0, 224.0],
+        '1.29': [288.0, 224.0], '1.38': [288.0, 208.0], '1.46': [304.0, 208.0], '1.67': [320.0, 192.0],
+        '1.75': [336.0, 192.0], '2.0': [352.0, 176.0], '2.09': [368.0, 176.0], '2.4': [384.0, 160.0],
+        '2.5': [400.0, 160.0], '2.89': [416.0, 144.0], '3.0': [432.0, 144.0], '3.11': [448.0, 144.0],
+        '3.62': [464.0, 128.0], '3.75': [480.0, 128.0], '3.88': [496.0, 128.0], '4.0': [512.0, 128.0]
+    }
+    NEW_ASPECT_RATIO = {}
+    for ratio in ASPECT_RATIO:
+        height, width = ASPECT_RATIO[ratio]
+        width = round(width / 256 * resolution)
+        height = round(height / 256 * resolution)
+        if width % 8 != 0:
+            print(f"skip train resolution {width}, {height}")
+            continue
+        if height % 8 != 0:
+            print(f"skip train resolution {width}, {height}")
+            continue
+        NEW_ASPECT_RATIO[ratio] = [height, width]
+    return NEW_ASPECT_RATIO
+
+ASPECT_RATIO_256 = calculate_aspect_ratios(256)
+ASPECT_RATIO_384 = calculate_aspect_ratios(384)
+ASPECT_RATIO_512 = calculate_aspect_ratios(512)
+ASPECT_RATIO_768 = calculate_aspect_ratios(768)
+ASPECT_RATIO_1024 = calculate_aspect_ratios(1024)
+
+def get_closest_ratio(height: float, width: float, ratios: dict):
+    aspect_ratio = height / width
+    closest_ratio = min(ratios.keys(), key=lambda ratio: abs(float(ratio) - aspect_ratio))
+    return ratios[closest_ratio], closest_ratio
+
+
+def _aspect_ratio_batched(
+    data,
+    batchsize=20,
+    aspect_ratios=ASPECT_RATIO_512,
+    batch_cross=False,
+    collation_fn=default_collation_fn,
+    partial=True,
+):
+    """Create batches of the given size.
+
+    :param data: iterator
+    :param batchsize: target batch size
+    :param tensors: automatically batch lists of ndarrays into ndarrays
+    :param partial: return partial batches
+    :returns: iterator
+
+    """
+    assert collation_fn is not None
+    buckets = {
+        ratio: {"cross": [], "no_cross": []} for ratio in aspect_ratios.keys()
+    }
+
+    def check(buckets):
+        for ratio in buckets:
+            for bucket_name in buckets[ratio]:
+                bucket = buckets[ratio][bucket_name]
+                assert len(bucket) < batchsize
+
+    for sample in data:
+        check(buckets)
+        height, width = sample['original_sizes']
+        (new_height, new_width), closest_ratio = get_closest_ratio(height, width, aspect_ratios)
+
+        bucket_name = "cross" if sample["has_cross"] and batch_cross else "no_cross"
+        bucket = buckets[closest_ratio][bucket_name]
+        bucket.append(sample)
+
+        if len(bucket) >= batchsize:
+            try:
+                batch = collation_fn(bucket)
+                yield batch
+                del batch
+            except Exception as e:
+                print(f"[aspect_ratio_batched] collation_fn batch failed due to error {e}")
+                for sample in bucket:
+                    if "__key__" in sample:
+                        print("error sample key in batch:", sample["__key__"])
+                    if "__url__" in sample:
+                        print("error sample url in batch:", sample["__url__"])
+            buckets[closest_ratio][bucket_name] = []
+            del bucket
+            gc.collect()
+
+    # yield the rest data and reset the buckets
+    for ratio in buckets.keys():
+        for bucket_name in ["cross", "no_cross"]:
+            bucket = buckets[ratio][bucket_name]
+            if len(bucket) > 0:
+                if len(bucket) == batchsize or partial:
+                    batch = collation_fn(bucket)
+                    yield batch
+                    del batch
+                buckets[ratio][bucket_name] = []
+                del bucket
+
+aspect_ratio_batched = pipelinefilter(_aspect_ratio_batched)
+
+def apply_aspect_ratio_batched(dataset, batchsize, aspect_ratios, batch_cross, collation_fn, partial=True):
+    return dataset.compose(
+        aspect_ratio_batched(
+            batchsize, 
+            aspect_ratios=aspect_ratios, 
+            batch_cross=batch_cross,
+            collation_fn=collation_fn, 
+            partial=partial
+        )
+    )
+
+def get_aspect_ratios(enable_aspect_ratio, resolution):
+    if enable_aspect_ratio:
+        # print("[Dataset] Multi Aspect Ratio Training Enabled")
+        if resolution == 256:
+            aspect_ratios = ASPECT_RATIO_256
+        elif resolution == 384:
+            aspect_ratios = ASPECT_RATIO_384
+        elif resolution == 512:
+            aspect_ratios = ASPECT_RATIO_512
+        elif resolution == 768:
+            aspect_ratios = ASPECT_RATIO_768
+        elif resolution == 1024:
+            aspect_ratios = ASPECT_RATIO_1024
+        else:
+            aspect_ratios = calculate_aspect_ratios(resolution)
+    else:
+        # print("[Dataset] Multi Aspect Ratio Training Disabled")
+        aspect_ratios = {
+            '1.0': [resolution, resolution]
+        }
+    return aspect_ratios
+
+def bbox_to_grid(bbox, image_size, output_size=(224, 224)):
+    """
+    Convert bounding box to a grid of points.
+    Args:
+        bbox (list of float): [xmin, ymin, xmax, ymax]
+        output_size (tuple of int): (height, width) of the output grid
+        
+    Returns:
+        torch.Tensor: Grid of points with shape (output_height, output_width, 2)
+    """
+    xmin, ymin, xmax, ymax = bbox
+    
+    # Create a meshgrid for the output grid
+    h, w = output_size
+    yy, xx = torch.meshgrid(
+        torch.linspace(ymin, ymax, h),
+        torch.linspace(xmin, xmax, w)
+    )
+    grid = torch.stack((xx, yy), -1)
+    
+    # Normalize grid to range [-1, 1]
+    H, W = image_size
+    grid[..., 0] = grid[..., 0] / (W - 1) * 2 - 1  # Normalize x to [-1, 1]
+    grid[..., 1] = grid[..., 1] / (H - 1) * 2 - 1  # Normalize y to [-1, 1]
+    
+    return grid
+
+def random_crop_instance(instance, min_crop_ratio):
+    assert 0 < min_crop_ratio <= 1
+    crop_width_ratio = random.uniform(min_crop_ratio, 1)
+    crop_height_ratio = random.uniform(min_crop_ratio, 1)
+    
+    orig_width, orig_height = instance.size
+    
+    crop_width = int(orig_width * crop_width_ratio)
+    crop_height = int(orig_height * crop_height_ratio)
+    
+    crop_left = random.randint(0, orig_width - crop_width)
+    crop_top = random.randint(0, orig_height - crop_height)
+    
+    crop_box = (crop_left, crop_top, crop_left + crop_width, crop_top + crop_height) # (left, upper, right, lower)
+    return instance.crop(crop_box), crop_box
+
+pil2tensor = T.ToTensor()
+tensor2pil = T.ToPILImage()
+
+cv2pil = lambda x: Image.fromarray(cv2.cvtColor(x, cv2.COLOR_BGR2RGB))
+pil2cv2 = lambda x: cv2.cvtColor(np.array(x), cv2.COLOR_RGB2BGR)
+
+def compute_psnr(x, y):
+    y = y.resize(x.size)
+    x = pil2tensor(x) * 255.
+    y = pil2tensor(y) * 255.
+    mse = torch.mean((x - y) ** 2)
+    return 20 * torch.log10(255.0 / torch.sqrt(mse)).item()
+
+def replace_first_occurrence(sentence, word_or_phrase, replace_with):
+    # Escape special characters in word_or_phrase for exact matching
+    escaped_word_or_phrase = re.escape(word_or_phrase)
+    pattern = r'\b' + escaped_word_or_phrase + r'\b'
+    
+    # Finding the first match
+    match = next(re.finditer(pattern, sentence), None)
+    if match:
+        # Perform replacement
+        result = re.sub(pattern, replace_with, sentence, count=1)
+        replaced = True
+        index = match.start()
+    else:
+        # No match found
+        result = sentence
+        replaced = False
+        index = -1
+    
+    return result, replaced, index
+
+
+def decode_base64_to_image(base64_str):
+    # Decode the base64 string to bytes
+    img_bytes = base64.b64decode(base64_str)
+    # Create a BytesIO buffer from the bytes
+    img_buffer = io.BytesIO(img_bytes)
+    # Open the image using Pillow
+    image = Image.open(img_buffer)
+    return image
+
+def jpeg_compression(pil_image, quality):
+    buffer = io.BytesIO()
+    pil_image.save(buffer, format="JPEG", quality=quality)
+    return Image.open(io.BytesIO(buffer.getvalue()))
+
+def pad_to_square(pil_image):
+    new_size = max(pil_image.width, pil_image.height)
+    square_image = Image.new("RGB", (new_size, new_size), "white")
+    left = (new_size - pil_image.width) // 2
+    top = (new_size - pil_image.height) // 2
+    square_image.paste(pil_image, (left, top))
+    return square_image
+
+def pad_to_target(pil_image, target_size):
+    original_width, original_height = pil_image.size
+    target_width, target_height = target_size
+    
+    original_aspect_ratio = original_width / original_height
+    target_aspect_ratio = target_width / target_height
+    
+    # Pad the image to the target aspect ratio
+    if original_aspect_ratio > target_aspect_ratio:
+        new_width = original_width
+        new_height = int(new_width / target_aspect_ratio)
+    else:
+        new_height = original_height
+        new_width = int(new_height * target_aspect_ratio)
+    
+    pad_image = Image.new("RGB", (new_width, new_height), "white")
+    left = (new_width - original_width) // 2
+    top = (new_height - original_height) // 2
+    pad_image.paste(pil_image, (left, top))
+    
+    # Resize the image to the target size
+    resized_image = pad_image.resize(target_size)
+    return resized_image
+
+def image_grid(imgs, rows, cols):
+    # assert len(imgs) == rows * cols
+
+    w, h = imgs[0].size
+    if imgs[0].mode == 'L':
+        grid = Image.new('L', size=(cols * w, rows * h))
+    else:
+        grid = Image.new('RGB', size=(cols * w, rows * h))
+
+    for i, img in enumerate(imgs):
+        grid.paste(img, box=(i % cols * w, i // cols * h))
+    return grid
+
+def split_grid(image):
+    width = image.width // 2
+    height = image.height // 2
+
+    crop_tuples_list = [
+        (0, 0, width, height),
+        (width, 0, width*2, height),
+        (0, height, width, height*2),
+        (width, height, width*2, height*2),
+    ]
+    def crop_image(input_image, crop_tuple=None):
+        if crop_tuple is None:
+            return input_image
+        return input_image.crop((crop_tuple[0], crop_tuple[1], crop_tuple[2], crop_tuple[3]))
+    
+    return [crop_image(image, crop_tuple) for crop_tuple in crop_tuples_list]
+
+def add_border(img, border_color, border_thickness):
+    """
+    Add a colored border to an image without changing its size.
+    
+    Parameters:
+        border_color (tuple): Border color in RGB (e.g., (255, 0, 0) for red).
+        border_thickness (int): Thickness of the border in pixels.
+    """
+    width, height = img.size
+    img = img.copy()
+    draw = ImageDraw.Draw(img)
+    draw.rectangle((0, 0, width, border_thickness), fill=border_color)
+    draw.rectangle((0, height - border_thickness, width, height), fill=border_color)
+    draw.rectangle((0, 0, border_thickness, height), fill=border_color)
+    draw.rectangle((width - border_thickness, 0, width, height), fill=border_color)
+    return img
+
+def merge_bboxes(bboxes):
+    if not bboxes:
+        return None  # Handle empty input
+    
+    # Extract all coordinates
+    x_mins = [b[0] for b in bboxes]
+    y_mins = [b[1] for b in bboxes]
+    x_maxs = [b[2] for b in bboxes]
+    y_maxs = [b[3] for b in bboxes]
+    
+    # Compute the merged box
+    merged_box = (
+        min(x_mins),  # x_min
+        min(y_mins),  # y_min
+        max(x_maxs),  # x_max
+        max(y_maxs)   # y_max
+    )
+    return merged_box
+
+
+def flip_bbox_left_right(bbox, image_width):
+    """
+    Flips the bounding box horizontally on an image.
+    
+    Parameters:
+    bbox (list of float): [x_min, y_min, x_max, y_max]
+    image_width (int): The width of the image
+    
+    Returns:
+    list of float: New bounding box after horizontal flip [x_min', y_min', x_max', y_max']
+    """
+    x_min, y_min, x_max, y_max = bbox
+    new_x_min = image_width - x_max
+    new_x_max = image_width - x_min
+    new_bbox = [new_x_min, y_min, new_x_max, y_max]
+    return new_bbox
+
+def json_load(path, encoding='ascii'):
+    with open(path, 'r', encoding=encoding) as file:
+        return json.load(file)
+
+def json_dump(obj, path, encoding='ascii', indent=4, create_dir=True, verbose=True, **kwargs):
+    if create_dir and os.path.dirname(path) != '':
+        os.makedirs(os.path.dirname(path), exist_ok=True)
+    with open(path, 'w', encoding=encoding) as file:
+        json.dump(obj, file, indent=4, ensure_ascii=False, **kwargs)
+    if verbose:
+        print(type(obj), 'saved to', path)

src/utils/modulation_utils.py

@@ -0,0 +1,55 @@
+# Copyright (c) 2025 Bytedance Ltd. and/or its affiliates
+# Copyright (c) Facebook, Inc. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import torch
+from src.flux.pipeline_tools import tokenize_t5_prompt
+
+def unpad_input_ids(input_ids, attention_mask):
+    return [input_ids[i][attention_mask[i].bool()][:-1] for i in range(input_ids.shape[0])]
+
+def get_word_index(pipe, prompt, input_ids, word, word_count=1, max_length=512, verbose=True, reverse=False):
+    word_inputs = tokenize_t5_prompt(pipe, word, max_length)
+    word_ids = unpad_input_ids(word_inputs.input_ids, word_inputs.attention_mask)[0]
+    if word_ids[0] == 3:
+        word_ids = word_ids[1:] # remove prefix space
+
+    if verbose:
+        print(f"Trying to find {word} {word_ids.tolist()} in {input_ids.tolist()} where")
+        print([(i, pipe.tokenizer_2.decode(input_ids[i])) for i in range(input_ids.shape[0])])
+    
+    count = 0
+    if reverse:
+        for i in range(input_ids.shape[0] - word_ids.shape[0],-1,-1):
+            if torch.equal(input_ids[i:i+word_ids.shape[0]], word_ids):
+                count += 1
+                if count == word_count:
+                    if verbose:
+                        reconstructed_word = pipe.tokenizer_2.decode(input_ids[i:i+word_ids.shape[0]])
+                        assert reconstructed_word == word
+                        print(f"[Reverse] Found index {i} to {i+word_ids.shape[0]} for '{word}' in prompt '{prompt}'")
+                        print("Reconstructed word", reconstructed_word)
+                    return i, i + word_ids.shape[0]
+    else:
+        for i in range(input_ids.shape[0] - word_ids.shape[0] + 1):
+            if torch.equal(input_ids[i:i+word_ids.shape[0]], word_ids):
+                count += 1
+                if count == word_count:
+                    if verbose:
+                        reconstructed_word = pipe.tokenizer_2.decode(input_ids[i:i+word_ids.shape[0]])
+                        assert reconstructed_word == word
+                        print(f"Found index {i} to {i+word_ids.shape[0]} for '{word}' in prompt '{prompt}'")
+                        print("Reconstructed word", reconstructed_word)
+                    return i, i + word_ids.shape[0]
+    print(f"[Error] Could not find '{word}' in prompt '{prompt}' with word_count {word_count}")