Delete src

src/adapters/mod_adapters.py

@@ -1,243 +0,0 @@
-# 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/environmentdata.py

@@ -1,55 +0,0 @@
-
-from huggingface_hub import snapshot_download
-import os
-
-# FLUX.1-dev
-snapshot_download(
-    repo_id="black-forest-labs/FLUX.1-dev",
-    local_dir="./FLUX.1-dev",
-    local_dir_use_symlinks=False
-)
-
-# Florence-2-large
-snapshot_download(
-    repo_id="microsoft/Florence-2-large",
-    local_dir="./Florence-2-large",
-    local_dir_use_symlinks=False
-)
-
-# CLIP ViT Large
-snapshot_download(
-    repo_id="openai/clip-vit-large-patch14",
-    local_dir="./clip-vit-large-patch14",
-    local_dir_use_symlinks=False
-)
-
-# DINO ViT-s16
-snapshot_download(
-    repo_id="facebook/dino-vits16",
-    local_dir="./dino-vits16",
-    local_dir_use_symlinks=False
-)
-
-# mPLUG Visual Question Answering
-snapshot_download(
-    repo_id="xingjianleng/mplug_visual-question-answering_coco_large_en",
-    local_dir="./mplug_visual-question-answering_coco_large_en",
-    local_dir_use_symlinks=False
-)
-
-# XVerse
-snapshot_download(
-    repo_id="ByteDance/XVerse",
-    local_dir="./XVerse",
-    local_dir_use_symlinks=False
-)
-
-
-os.environ["FLORENCE2_MODEL_PATH"]    = "./checkpoints/Florence-2-large"
-os.environ["SAM2_MODEL_PATH"]         = "./checkpoints/sam2.1_hiera_large.pt"
-os.environ["FACE_ID_MODEL_PATH"]      = "./checkpoints/model_ir_se50.pth"
-os.environ["CLIP_MODEL_PATH"]         = "./checkpoints/clip-vit-large-patch14"
-os.environ["FLUX_MODEL_PATH"]         = "./checkpoints/FLUX.1-dev"
-os.environ["DPG_VQA_MODEL_PATH"]      = "./checkpoints/mplug_visual-question-answering_coco_large_en"
-os.environ["DINO_MODEL_PATH"]         = "./checkpoints/dino-vits16"
-

src/flux/block.py

@@ -1,814 +0,0 @@
-# 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

@@ -1,133 +0,0 @@
-# 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
-
-import src.flux.pipeline_tools
-
-# 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

@@ -1,838 +0,0 @@
-# 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
-import src.flux.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

@@ -1,99 +0,0 @@
-# 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

@@ -1,685 +0,0 @@
-# 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

@@ -1,363 +0,0 @@
-# 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,
-    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 is_torch_version(spec: str) -> bool:
-    # e.g. spec = ">=1.12.0"
-    return version.parse(torch.__version__) in version.SpecifierSet(spec)
-    
-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

@@ -1,404 +0,0 @@
-# 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

@@ -1,55 +0,0 @@
-# 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}")