ai-toolkit/toolkit/models/vd_adapter.py

import sys

import torch
import torch.nn as nn
import torch.nn.functional as F
import weakref
from typing import Union, TYPE_CHECKING, Optional
from collections import OrderedDict

from diffusers import Transformer2DModel, FluxTransformer2DModel
from transformers import T5EncoderModel, CLIPTextModel, CLIPTokenizer, T5Tokenizer, CLIPVisionModelWithProjection
from toolkit.models.pixtral_vision import PixtralVisionEncoder, PixtralVisionImagePreprocessor, VisionLanguageAdapter
from transformers import SiglipImageProcessor, SiglipVisionModel
import traceback
from toolkit.config_modules import AdapterConfig
from toolkit.paths import REPOS_ROOT
sys.path.append(REPOS_ROOT)


if TYPE_CHECKING:
    from toolkit.stable_diffusion_model import StableDiffusion
    from toolkit.custom_adapter import CustomAdapter
    

# matches distribution of randn
class Norm(nn.Module):
    def __init__(self, target_mean=0.0, target_std=1.0, eps=1e-6):
        super(Norm, self).__init__()
        self.target_mean = target_mean
        self.target_std = target_std
        self.eps = eps

    def forward(self, x):
        dims = tuple(range(1, x.dim()))
        mean = x.mean(dim=dims, keepdim=True)
        std = x.std(dim=dims, keepdim=True)
        
        # Normalize
        return self.target_std * (x - mean) / (std + self.eps) + self.target_mean


norm_layer = Norm()
    
class SparseAutoencoder(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim):
        super(SparseAutoencoder, self).__init__()
        self.encoder = nn.Sequential(
            nn.Linear(input_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, output_dim),
        )
        self.norm = Norm()
        self.decoder = nn.Sequential(
            nn.Linear(output_dim, hidden_dim),
            nn.GELU(),
            nn.Linear(hidden_dim, input_dim),
        )
        self.last_run = None

    def forward(self, x):
        self.last_run = {
            "input": x
        }
        x = self.encoder(x)
        x = self.norm(x)
        self.last_run["sparse"] = x
        x = self.decoder(x)
        x = self.norm(x)
        self.last_run["output"] = x
        return x


class MLPR(nn.Module):  # MLP with reshaping
    def __init__(
            self,
            in_dim,
            in_channels,
            out_dim,
            out_channels,
            use_residual=True
    ):
        super().__init__()
        if use_residual:
            assert in_dim == out_dim
        # dont normalize if using conv
        self.layer_norm = nn.LayerNorm(in_dim)

        self.fc1 = nn.Linear(in_dim, out_dim)
        self.act_fn = nn.GELU()
        self.conv1 = nn.Conv1d(in_channels, out_channels, 1)

    def forward(self, x):
        residual = x
        x = self.layer_norm(x)
        x = self.fc1(x)
        x = self.act_fn(x)
        x = self.conv1(x)
        return x

class AttnProcessor2_0(torch.nn.Module):
    r"""
    Processor for implementing scaled dot-product attention (enabled by default if you're using PyTorch 2.0).
    """

    def __init__(
        self,
        hidden_size=None,
        cross_attention_dim=None,
    ):
        super().__init__()
        if not hasattr(F, "scaled_dot_product_attention"):
            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")

    def __call__(
        self,
        attn,
        hidden_states,
        encoder_hidden_states=None,
        attention_mask=None,
        temb=None,
    ):
        residual = hidden_states

        if attn.spatial_norm is not None:
            hidden_states = attn.spatial_norm(hidden_states, temb)

        input_ndim = hidden_states.ndim

        if input_ndim == 4:
            batch_size, channel, height, width = hidden_states.shape
            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)

        batch_size, sequence_length, _ = (
            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
        )

        if attention_mask is not None:
            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
            # scaled_dot_product_attention expects attention_mask shape to be
            # (batch, heads, source_length, target_length)
            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])

        if attn.group_norm is not None:
            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)

        query = attn.to_q(hidden_states)

        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        elif attn.norm_cross:
            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)

        key = attn.to_k(encoder_hidden_states)
        value = attn.to_v(encoder_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)

        # the output of sdp = (batch, num_heads, seq_len, head_dim)
        # TODO: add support for attn.scale when we move to Torch 2.1
        hidden_states = F.scaled_dot_product_attention(
            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
        )

        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
        hidden_states = hidden_states.to(query.dtype)

        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        if input_ndim == 4:
            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)

        if attn.residual_connection:
            hidden_states = hidden_states + residual

        hidden_states = hidden_states / attn.rescale_output_factor

        return hidden_states

class VisionDirectAdapterAttnProcessor(nn.Module):
    r"""
    Attention processor for Custom TE for PyTorch 2.0.
    Args:
        hidden_size (`int`):
            The hidden size of the attention layer.
        cross_attention_dim (`int`):
            The number of channels in the `encoder_hidden_states`.
        scale (`float`, defaults to 1.0):
            the weight scale of image prompt.
        adapter
    """

    def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, adapter=None,
                 adapter_hidden_size=None, has_bias=False, **kwargs):
        super().__init__()

        if not hasattr(F, "scaled_dot_product_attention"):
            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")

        self.adapter_ref: weakref.ref = weakref.ref(adapter)

        self.hidden_size = hidden_size
        self.adapter_hidden_size = adapter_hidden_size
        self.cross_attention_dim = cross_attention_dim
        self.scale = scale

        self.to_k_adapter = nn.Linear(adapter_hidden_size, hidden_size, bias=has_bias)
        self.to_v_adapter = nn.Linear(adapter_hidden_size, hidden_size, bias=has_bias)

    @property
    def is_active(self):
        return self.adapter_ref().is_active
        # return False

    @property
    def unconditional_embeds(self):
        return self.adapter_ref().adapter_ref().unconditional_embeds

    @property
    def conditional_embeds(self):
        return self.adapter_ref().adapter_ref().conditional_embeds

    def __call__(
            self,
            attn,
            hidden_states,
            encoder_hidden_states=None,
            attention_mask=None,
            temb=None,
    ):
        is_active = self.adapter_ref().is_active
        residual = hidden_states

        if attn.spatial_norm is not None:
            hidden_states = attn.spatial_norm(hidden_states, temb)

        input_ndim = hidden_states.ndim

        if input_ndim == 4:
            batch_size, channel, height, width = hidden_states.shape
            hidden_states = hidden_states.view(batch_size, channel, height * width).transpose(1, 2)

        batch_size, sequence_length, _ = (
            hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape
        )

        if attention_mask is not None:
            attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)
            # scaled_dot_product_attention expects attention_mask shape to be
            # (batch, heads, source_length, target_length)
            attention_mask = attention_mask.view(batch_size, attn.heads, -1, attention_mask.shape[-1])

        if attn.group_norm is not None:
            hidden_states = attn.group_norm(hidden_states.transpose(1, 2)).transpose(1, 2)

        query = attn.to_q(hidden_states)

        # will be none if disabled
        if encoder_hidden_states is None:
            encoder_hidden_states = hidden_states
        elif attn.norm_cross:
            encoder_hidden_states = attn.norm_encoder_hidden_states(encoder_hidden_states)

        key = attn.to_k(encoder_hidden_states)
        value = attn.to_v(encoder_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)

        # the output of sdp = (batch, num_heads, seq_len, head_dim)
        # TODO: add support for attn.scale when we move to Torch 2.1
        hidden_states = F.scaled_dot_product_attention(
            query, key, value, attn_mask=attention_mask, dropout_p=0.0, is_causal=False
        )

        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
        hidden_states = hidden_states.to(query.dtype)

        # only use one TE or the other. If our adapter is active only use ours
        if self.is_active and self.conditional_embeds is not None:
            try:

                adapter_hidden_states = self.conditional_embeds
                if adapter_hidden_states.shape[0] == batch_size // 2:
                    adapter_hidden_states = torch.cat([
                        self.unconditional_embeds,
                        adapter_hidden_states
                    ], dim=0)
                    # if it is image embeds, we need to add a 1 dim at inx 1
                if len(adapter_hidden_states.shape) == 2:
                    adapter_hidden_states = adapter_hidden_states.unsqueeze(1)
                # conditional_batch_size = adapter_hidden_states.shape[0]
                # conditional_query = query

                # for ip-adapter
                vd_key = self.to_k_adapter(adapter_hidden_states)
                vd_value = self.to_v_adapter(adapter_hidden_states)

                vd_key = vd_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
                vd_value = vd_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

                # the output of sdp = (batch, num_heads, seq_len, head_dim)
                # TODO: add support for attn.scale when we move to Torch 2.1
                vd_hidden_states = F.scaled_dot_product_attention(
                    query, vd_key, vd_value, attn_mask=None, dropout_p=0.0, is_causal=False
                )

                vd_hidden_states = vd_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
                vd_hidden_states = vd_hidden_states.to(query.dtype)

                hidden_states = hidden_states + self.scale * vd_hidden_states
            except Exception as e:
                print("Error in VisionDirectAdapterAttnProcessor")
                # print shapes of all tensors
                print(f"hidden_states: {hidden_states.shape}")
                print(f"adapter_hidden_states: {adapter_hidden_states.shape}")
                print(f"vd_key: {vd_key.shape}")
                print(f"vd_value: {vd_value.shape}")
                print(f"vd_hidden_states: {vd_hidden_states.shape}")
                print(f"query: {query.shape}")
                print(f"key: {key.shape}")
                print(f"value: {value.shape}")
                print(f"inner_dim: {inner_dim}")
                print(f"head_dim: {head_dim}")
                print(f"batch_size: {batch_size}")
                traceback.print_exc()


        # linear proj
        hidden_states = attn.to_out[0](hidden_states)
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        if input_ndim == 4:
            hidden_states = hidden_states.transpose(-1, -2).reshape(batch_size, channel, height, width)

        if attn.residual_connection:
            hidden_states = hidden_states + residual

        hidden_states = hidden_states / attn.rescale_output_factor

        return hidden_states


class CustomFluxVDAttnProcessor2_0(torch.nn.Module):
    """Attention processor used typically in processing the SD3-like self-attention projections."""

    def __init__(self, hidden_size, cross_attention_dim=None, scale=1.0, adapter=None,
                 adapter_hidden_size=None, has_bias=False, block_idx=0, **kwargs):
        super().__init__()

        if not hasattr(F, "scaled_dot_product_attention"):
            raise ImportError("AttnProcessor2_0 requires PyTorch 2.0, to use it, please upgrade PyTorch to 2.0.")

        self.adapter_ref: weakref.ref = weakref.ref(adapter)

        self.hidden_size = hidden_size
        self.adapter_hidden_size = adapter_hidden_size
        self.cross_attention_dim = cross_attention_dim
        self.scale = scale
        self.block_idx = block_idx

        self.to_k_adapter = nn.Linear(adapter_hidden_size, hidden_size, bias=has_bias)
        self.to_v_adapter = nn.Linear(adapter_hidden_size, hidden_size, bias=has_bias)

    @property
    def is_active(self):
        return self.adapter_ref().is_active
        # return False

    @property
    def unconditional_embeds(self):
        return self.adapter_ref().adapter_ref().unconditional_embeds

    @property
    def conditional_embeds(self):
        return self.adapter_ref().adapter_ref().conditional_embeds

    def __call__(
        self,
        attn,
        hidden_states: torch.FloatTensor,
        encoder_hidden_states: torch.FloatTensor = None,
        attention_mask: Optional[torch.FloatTensor] = None,
        image_rotary_emb: Optional[torch.Tensor] = None,
    ) -> torch.FloatTensor:
        batch_size, _, _ = hidden_states.shape if encoder_hidden_states is None else encoder_hidden_states.shape

        # `sample` projections.
        query = attn.to_q(hidden_states)
        key = attn.to_k(hidden_states)
        value = attn.to_v(hidden_states)

        inner_dim = key.shape[-1]
        head_dim = inner_dim // attn.heads

        query = query.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        key = key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
        value = value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

        if attn.norm_q is not None:
            query = attn.norm_q(query)
        if attn.norm_k is not None:
            key = attn.norm_k(key)

        # the attention in FluxSingleTransformerBlock does not use `encoder_hidden_states`
        if encoder_hidden_states is not None:
            # `context` projections.
            encoder_hidden_states_query_proj = attn.add_q_proj(encoder_hidden_states)
            encoder_hidden_states_key_proj = attn.add_k_proj(encoder_hidden_states)
            encoder_hidden_states_value_proj = attn.add_v_proj(encoder_hidden_states)

            encoder_hidden_states_query_proj = encoder_hidden_states_query_proj.view(
                batch_size, -1, attn.heads, head_dim
            ).transpose(1, 2)
            encoder_hidden_states_key_proj = encoder_hidden_states_key_proj.view(
                batch_size, -1, attn.heads, head_dim
            ).transpose(1, 2)
            encoder_hidden_states_value_proj = encoder_hidden_states_value_proj.view(
                batch_size, -1, attn.heads, head_dim
            ).transpose(1, 2)

            if attn.norm_added_q is not None:
                encoder_hidden_states_query_proj = attn.norm_added_q(encoder_hidden_states_query_proj)
            if attn.norm_added_k is not None:
                encoder_hidden_states_key_proj = attn.norm_added_k(encoder_hidden_states_key_proj)

            # attention
            query = torch.cat([encoder_hidden_states_query_proj, query], dim=2)
            key = torch.cat([encoder_hidden_states_key_proj, key], dim=2)
            value = torch.cat([encoder_hidden_states_value_proj, value], dim=2)

        if image_rotary_emb is not None:
            from diffusers.models.embeddings import apply_rotary_emb

            query = apply_rotary_emb(query, image_rotary_emb)
            key = apply_rotary_emb(key, image_rotary_emb)

        hidden_states = F.scaled_dot_product_attention(query, key, value, dropout_p=0.0, is_causal=False)
        hidden_states = hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
        hidden_states = hidden_states.to(query.dtype)

        # begin ip adapter
        if self.is_active and self.conditional_embeds is not None:
            adapter_hidden_states = self.conditional_embeds
            block_scaler = self.adapter_ref().block_scaler
            if block_scaler is not None:
                # add 1 to block scaler so we can decay its weight to 1.0
                block_scaler = block_scaler[self.block_idx] + 1.0

            if adapter_hidden_states.shape[0] < batch_size:
                adapter_hidden_states = torch.cat([
                    self.unconditional_embeds,
                    adapter_hidden_states
                ], dim=0)
                # if it is image embeds, we need to add a 1 dim at inx 1
            if len(adapter_hidden_states.shape) == 2:
                adapter_hidden_states = adapter_hidden_states.unsqueeze(1)
            # conditional_batch_size = adapter_hidden_states.shape[0]
            # conditional_query = query

            # for ip-adapter
            vd_key = self.to_k_adapter(adapter_hidden_states)
            vd_value = self.to_v_adapter(adapter_hidden_states)

            vd_key = vd_key.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)
            vd_value = vd_value.view(batch_size, -1, attn.heads, head_dim).transpose(1, 2)

            vd_hidden_states = F.scaled_dot_product_attention(
                query, vd_key, vd_value, attn_mask=None, dropout_p=0.0, is_causal=False
            )

            vd_hidden_states = vd_hidden_states.transpose(1, 2).reshape(batch_size, -1, attn.heads * head_dim)
            vd_hidden_states = vd_hidden_states.to(query.dtype)

            # scale to block scaler
            if block_scaler is not None:
                orig_dtype = vd_hidden_states.dtype
                if block_scaler.dtype != vd_hidden_states.dtype:
                    vd_hidden_states = vd_hidden_states.to(block_scaler.dtype)
                vd_hidden_states = vd_hidden_states * block_scaler
                if block_scaler.dtype != orig_dtype:
                    vd_hidden_states = vd_hidden_states.to(orig_dtype)

            hidden_states = hidden_states + self.scale * vd_hidden_states

        if encoder_hidden_states is not None:
            encoder_hidden_states, hidden_states = (
                hidden_states[:, : encoder_hidden_states.shape[1]],
                hidden_states[:, encoder_hidden_states.shape[1] :],
            )

            # linear proj
            hidden_states = attn.to_out[0](hidden_states)
            # dropout
            hidden_states = attn.to_out[1](hidden_states)
            encoder_hidden_states = attn.to_add_out(encoder_hidden_states)

            return hidden_states, encoder_hidden_states
        else:
            return hidden_states

class VisionDirectAdapter(torch.nn.Module):
    def __init__(
            self,
            adapter: 'CustomAdapter',
            sd: 'StableDiffusion',
            vision_model: Union[CLIPVisionModelWithProjection],
    ):
        super(VisionDirectAdapter, self).__init__()
        is_pixart = sd.is_pixart
        is_flux = sd.is_flux
        self.adapter_ref: weakref.ref = weakref.ref(adapter)
        self.sd_ref: weakref.ref = weakref.ref(sd)
        self.config: AdapterConfig = adapter.config
        self.vision_model_ref: weakref.ref = weakref.ref(vision_model)
        self.resampler = None
        is_pixtral = self.config.image_encoder_arch == "pixtral"

        if adapter.config.clip_layer == "image_embeds":
            if isinstance(vision_model, SiglipVisionModel):
                self.token_size = vision_model.config.hidden_size
            else:
                self.token_size = vision_model.config.projection_dim
        else:
            self.token_size = vision_model.config.hidden_size
            
        self.mid_size = self.token_size
        
        if self.config.conv_pooling and self.config.conv_pooling_stacks > 1:
            self.mid_size = self.mid_size * self.config.conv_pooling_stacks
        
        # if pixtral, use cross attn dim for more sparse representation if only doing double transformers
        if is_pixtral and self.config.flux_only_double:
            if is_flux:
                hidden_size = 3072
            else:
                hidden_size = sd.unet.config['cross_attention_dim']
            self.mid_size = hidden_size

        # init adapter modules
        attn_procs = {}
        unet_sd = sd.unet.state_dict()

        attn_processor_keys = []
        if is_pixart:
            transformer: Transformer2DModel = sd.unet
            for i, module in transformer.transformer_blocks.named_children():

                attn_processor_keys.append(f"transformer_blocks.{i}.attn1")

                # cross attention
                attn_processor_keys.append(f"transformer_blocks.{i}.attn2")

        elif is_flux:
            transformer: FluxTransformer2DModel = sd.unet
            for i, module in transformer.transformer_blocks.named_children():
                attn_processor_keys.append(f"transformer_blocks.{i}.attn")

            if not self.config.flux_only_double:
                # single transformer blocks do not have cross attn, but we will do them anyway
                for i, module in transformer.single_transformer_blocks.named_children():
                    attn_processor_keys.append(f"single_transformer_blocks.{i}.attn")
        else:
            attn_processor_keys = list(sd.unet.attn_processors.keys())

        current_idx = 0

        for name in attn_processor_keys:
            if is_flux:
                cross_attention_dim = None
            else:
                cross_attention_dim = None if name.endswith("attn1.processor") or name.endswith("attn.1") else sd.unet.config['cross_attention_dim']
            if name.startswith("mid_block"):
                hidden_size = sd.unet.config['block_out_channels'][-1]
            elif name.startswith("up_blocks"):
                block_id = int(name[len("up_blocks.")])
                hidden_size = list(reversed(sd.unet.config['block_out_channels']))[block_id]
            elif name.startswith("down_blocks"):
                block_id = int(name[len("down_blocks.")])
                hidden_size = sd.unet.config['block_out_channels'][block_id]
            elif name.startswith("transformer") or name.startswith("single_transformer"):
                if is_flux:
                    hidden_size = 3072
                else:
                    hidden_size = sd.unet.config['cross_attention_dim']
            else:
                # they didnt have this, but would lead to undefined below
                raise ValueError(f"unknown attn processor name: {name}")
            if cross_attention_dim is None and not is_flux:
                attn_procs[name] = AttnProcessor2_0()
            else:
                layer_name = name.split(".processor")[0]
                if f"{layer_name}.to_k.weight._data" in unet_sd and is_flux:
                    # is quantized

                    to_k_adapter = torch.randn(hidden_size, hidden_size) * 0.01
                    to_v_adapter = torch.randn(hidden_size, hidden_size) * 0.01
                    to_k_adapter = to_k_adapter.to(self.sd_ref().torch_dtype)
                    to_v_adapter = to_v_adapter.to(self.sd_ref().torch_dtype)
                else:
                    to_k_adapter = unet_sd[layer_name + ".to_k.weight"]
                    to_v_adapter = unet_sd[layer_name + ".to_v.weight"]

                # add zero padding to the adapter
                if to_k_adapter.shape[1] < self.mid_size:
                    to_k_adapter = torch.cat([
                        to_k_adapter,
                        torch.randn(to_k_adapter.shape[0], self.mid_size - to_k_adapter.shape[1]).to(
                            to_k_adapter.device, dtype=to_k_adapter.dtype) * 0.01
                    ],
                        dim=1
                    )
                    to_v_adapter = torch.cat([
                        to_v_adapter,
                        torch.randn(to_v_adapter.shape[0], self.mid_size - to_v_adapter.shape[1]).to(
                            to_k_adapter.device, dtype=to_k_adapter.dtype) * 0.01
                    ],
                        dim=1
                    )
                elif to_k_adapter.shape[1] > self.mid_size:
                    to_k_adapter = to_k_adapter[:, :self.mid_size]
                    to_v_adapter = to_v_adapter[:, :self.mid_size]
                    # if is_pixart:
                    #     to_k_bias = to_k_bias[:self.mid_size]
                    #     to_v_bias = to_v_bias[:self.mid_size]
                else:
                    to_k_adapter = to_k_adapter
                    to_v_adapter = to_v_adapter
                    # if is_pixart:
                    #     to_k_bias = to_k_bias
                    #     to_v_bias = to_v_bias

                weights = {
                    "to_k_adapter.weight": to_k_adapter * 0.01,
                    "to_v_adapter.weight": to_v_adapter * 0.01,
                }
                # if is_pixart:
                #     weights["to_k_adapter.bias"] = to_k_bias
                #     weights["to_v_adapter.bias"] = to_v_bias\

                if is_flux:
                    attn_procs[name] = CustomFluxVDAttnProcessor2_0(
                        hidden_size=hidden_size,
                        cross_attention_dim=cross_attention_dim,
                        scale=1.0,
                        adapter=self,
                        adapter_hidden_size=self.mid_size,
                        has_bias=False,
                        block_idx=current_idx
                    )
                else:
                    attn_procs[name] = VisionDirectAdapterAttnProcessor(
                        hidden_size=hidden_size,
                        cross_attention_dim=cross_attention_dim,
                        scale=1.0,
                        adapter=self,
                        adapter_hidden_size=self.mid_size,
                        has_bias=False,
                    )
                current_idx += 1
                attn_procs[name].load_state_dict(weights)

        if self.sd_ref().is_pixart:
            # we have to set them ourselves
            transformer: Transformer2DModel = sd.unet
            for i, module in transformer.transformer_blocks.named_children():
                module.attn1.processor = attn_procs[f"transformer_blocks.{i}.attn1"]
                module.attn2.processor = attn_procs[f"transformer_blocks.{i}.attn2"]
            self.adapter_modules = torch.nn.ModuleList([
                transformer.transformer_blocks[i].attn1.processor for i in range(len(transformer.transformer_blocks))
            ] + [
                transformer.transformer_blocks[i].attn2.processor for i in range(len(transformer.transformer_blocks))
            ])
        elif self.sd_ref().is_flux:
            # we have to set them ourselves
            transformer: FluxTransformer2DModel = sd.unet
            for i, module in transformer.transformer_blocks.named_children():
                module.attn.processor = attn_procs[f"transformer_blocks.{i}.attn"]

            if not self.config.flux_only_double:
                # do single blocks too even though they dont have cross attn
                for i, module in transformer.single_transformer_blocks.named_children():
                    module.attn.processor = attn_procs[f"single_transformer_blocks.{i}.attn"]

            if not self.config.flux_only_double:
                self.adapter_modules = torch.nn.ModuleList(
                    [
                        transformer.transformer_blocks[i].attn.processor for i in
                        range(len(transformer.transformer_blocks))
                    ] + [
                        transformer.single_transformer_blocks[i].attn.processor for i in
                        range(len(transformer.single_transformer_blocks))
                    ]
                )
            else:
                self.adapter_modules = torch.nn.ModuleList(
                    [
                        transformer.transformer_blocks[i].attn.processor for i in
                        range(len(transformer.transformer_blocks))
                    ]
                )
        else:
            sd.unet.set_attn_processor(attn_procs)
            self.adapter_modules = torch.nn.ModuleList(sd.unet.attn_processors.values())

        num_modules = len(self.adapter_modules)
        if self.config.train_scaler:
            self.block_scaler = torch.nn.Parameter(torch.tensor([0.0] * num_modules).to(
                dtype=torch.float32,
                device=self.sd_ref().device_torch
            ))
            self.block_scaler.data = self.block_scaler.data.to(torch.float32)
            self.block_scaler.requires_grad = True
        else:
            self.block_scaler = None
        
        self.pool = None

        if self.config.num_tokens is not None:
            # image_encoder_state_dict = self.adapter_ref().vision_encoder.state_dict()
            # max_seq_len = CLIP tokens + CLS token
            # max_seq_len = 257
            # if "vision_model.embeddings.position_embedding.weight" in image_encoder_state_dict:
            #     # clip
            #     max_seq_len = int(
            #         image_encoder_state_dict["vision_model.embeddings.position_embedding.weight"].shape[0])
            # self.resampler = MLPR(
            #     in_dim=self.token_size,
            #     in_channels=max_seq_len,
            #     out_dim=self.mid_size,
            #     out_channels=self.config.num_tokens,
            # )
            vision_config = self.adapter_ref().vision_encoder.config
            # sequence_length = int((vision_config.image_size / vision_config.patch_size) ** 2 + 1)
            # siglip doesnt add 1
            sequence_length = int((vision_config.image_size / vision_config.patch_size) ** 2)
            self.pool = nn.Sequential(
                nn.Conv1d(sequence_length, self.config.num_tokens, 1, bias=False),
                Norm(),
            )
        
        elif self.config.image_encoder_arch == "pixtral":
            self.resampler = VisionLanguageAdapter(
                in_dim=self.token_size,
                out_dim=self.mid_size,
            )
        
        self.sparse_autoencoder = None
        if self.config.conv_pooling:
            vision_config = self.adapter_ref().vision_encoder.config
            # sequence_length = int((vision_config.image_size / vision_config.patch_size) ** 2 + 1)
            # siglip doesnt add 1
            sequence_length = int((vision_config.image_size / vision_config.patch_size) ** 2)
            self.pool = nn.Sequential(
                nn.Conv1d(sequence_length, self.config.conv_pooling_stacks, 1, bias=False),
                Norm(),
            )
        if self.config.sparse_autoencoder_dim is not None:
            hidden_dim  = self.token_size * 2
            if hidden_dim > self.config.sparse_autoencoder_dim:
                hidden_dim = self.config.sparse_autoencoder_dim
            self.sparse_autoencoder = SparseAutoencoder(
                input_dim=self.token_size,
                hidden_dim=hidden_dim,
                output_dim=self.config.sparse_autoencoder_dim
            )
        
        if self.config.clip_layer == "image_embeds":
            self.proj = nn.Linear(self.token_size, self.token_size)

    def state_dict(self, destination=None, prefix='', keep_vars=False):
        if self.config.train_scaler:
            # only return the block scaler
            if destination is None:
                destination = OrderedDict()
            destination[prefix + 'block_scaler'] = self.block_scaler
            return destination
        return super().state_dict(destination, prefix, keep_vars)

    # make a getter to see if is active
    @property
    def is_active(self):
        return self.adapter_ref().is_active

    def forward(self, input):
        # block scaler keeps moving dtypes. make sure it is float32 here
        # todo remove this when we have a real solution
        
        if self.block_scaler is not None and self.block_scaler.dtype != torch.float32:
            self.block_scaler.data = self.block_scaler.data.to(torch.float32)
        # if doing image_embeds, normalize here
        if self.config.clip_layer == "image_embeds":
            input = norm_layer(input)
            input = self.proj(input)
        if self.resampler is not None:
            input = self.resampler(input)
        if self.pool is not None:
            input = self.pool(input)
            if self.config.conv_pooling_stacks > 1:
                input = torch.cat(torch.chunk(input, self.config.conv_pooling_stacks, dim=1), dim=2)
        if self.sparse_autoencoder is not None:
            input = self.sparse_autoencoder(input)
        return input

    def to(self, *args, **kwargs):
        super().to(*args, **kwargs)
        if self.block_scaler is not None:
            if self.block_scaler.dtype != torch.float32:
                self.block_scaler.data = self.block_scaler.data.to(torch.float32)
        return self

    def post_weight_update(self):
        # force block scaler to be mean of 1
        pass