models.py

import torch
import torch.nn as nn
import torch.nn.functional as F

from typing import List, Tuple, Union
from dataclasses import dataclass
from diffusers.utils.outputs import BaseOutput
from diffusers.configuration_utils import ConfigMixin, register_to_config
from diffusers.models.modeling_utils import ModelMixin
from diffusers.models.unet_2d_blocks import get_down_block as get_down_block_default
from diffusers.models.resnet import Mish, Upsample2D, Downsample2D, upsample_2d, downsample_2d, partial
from diffusers.models.cross_attention import CrossAttention, LoRALinearLayer # , LoRACrossAttnProcessor


def get_down_block(
    down_block_type,
    num_layers,
    in_channels,
    out_channels,
    temb_channels,
    add_downsample,
    resnet_eps,
    resnet_act_fn,
    attn_num_head_channels,
    resnet_groups=None,
    cross_attention_dim=None,
    downsample_padding=None,
    dual_cross_attention=False,
    use_linear_projection=False,
    only_cross_attention=False,
    upcast_attention=False,
    resnet_time_scale_shift="default",
    resnet_kernel_size=3,
):
    down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
    if down_block_type == "SimpleDownEncoderBlock2D":
        return SimpleDownEncoderBlock2D(
            num_layers=num_layers,
            in_channels=in_channels,
            out_channels=out_channels,
            add_downsample=add_downsample,
            convnet_eps=resnet_eps,
            convnet_act_fn=resnet_act_fn,
            convnet_groups=resnet_groups,
            downsample_padding=downsample_padding,
            convnet_time_scale_shift=resnet_time_scale_shift,
            convnet_kernel_size=resnet_kernel_size
        )
    else:
        return get_down_block_default(
            down_block_type,
            num_layers,
            in_channels,
            out_channels,
            temb_channels,
            add_downsample,
            resnet_eps,
            resnet_act_fn,
            attn_num_head_channels,
            resnet_groups=resnet_groups,
            cross_attention_dim=cross_attention_dim,
            downsample_padding=downsample_padding,
            dual_cross_attention=dual_cross_attention,
            use_linear_projection=use_linear_projection,
            only_cross_attention=only_cross_attention,
            upcast_attention=upcast_attention,
            resnet_time_scale_shift=resnet_time_scale_shift,
            # resnet_kernel_size=resnet_kernel_size
        )


class LoRACrossAttnProcessor(nn.Module):
    def __init__(
            self, 
            hidden_size, 
            cross_attention_dim=None, 
            rank=4, 
            post_add=False,
            key_states_skipped=False,
            value_states_skipped=False,
            output_states_skipped=False):
        super().__init__()

        self.hidden_size = hidden_size
        self.cross_attention_dim = cross_attention_dim
        self.rank = rank
        self.post_add = post_add

        self.to_q_lora = LoRALinearLayer(hidden_size, hidden_size, rank)
        if not key_states_skipped:
            self.to_k_lora = LoRALinearLayer(
                hidden_size if post_add else (cross_attention_dim or hidden_size), hidden_size, rank)
        if not value_states_skipped:
            self.to_v_lora = LoRALinearLayer(
                hidden_size if post_add else (cross_attention_dim or hidden_size), hidden_size, rank)
        if not output_states_skipped:
            self.to_out_lora = LoRALinearLayer(hidden_size, hidden_size, rank)

        self.key_states_skipped: bool = key_states_skipped
        self.value_states_skipped: bool = value_states_skipped
        self.output_states_skipped: bool = output_states_skipped

    def skip_key_states(self, is_skipped: bool = True):
        if is_skipped == False:
            assert hasattr(self, 'to_k_lora')
        self.key_states_skipped = is_skipped

    def skip_value_states(self, is_skipped: bool = True):
        if is_skipped == False:
            assert hasattr(self, 'to_q_lora')
        self.value_states_skipped = is_skipped

    def skip_output_states(self, is_skipped: bool = True):
        if is_skipped == False:
            assert hasattr(self, 'to_out_lora')
        self.output_states_skipped = is_skipped

    def __call__(
        self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None, scale=1.0
    ):
        batch_size, sequence_length, _ = hidden_states.shape
        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length, batch_size)

        query = attn.to_q(hidden_states) 
        query = query + scale * self.to_q_lora(query if self.post_add else hidden_states)
        query = attn.head_to_batch_dim(query)

        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states

        key = attn.to_k(encoder_hidden_states) 
        if not self.key_states_skipped:
            key = key + scale * self.to_k_lora(key if self.post_add else encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)
        if not self.value_states_skipped:
            value = value + scale * self.to_v_lora(value if self.post_add else encoder_hidden_states)

        key = attn.head_to_batch_dim(key)
        value = attn.head_to_batch_dim(value)

        attention_probs = attn.get_attention_scores(query, key, attention_mask)
        hidden_states = torch.bmm(attention_probs, value)
        hidden_states = attn.batch_to_head_dim(hidden_states)

        # linear proj
        out = attn.to_out[0](hidden_states)
        if not self.output_states_skipped:
            out = out + scale * self.to_out_lora(out if self.post_add else hidden_states)
        hidden_states = out
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        return hidden_states


class ControlLoRACrossAttnProcessor(LoRACrossAttnProcessor):
    def __init__(
            self, 
            hidden_size, 
            cross_attention_dim=None, 
            rank=4, 
            control_rank=None, 
            post_add=False, 
            concat_hidden=False,
            control_channels=None,
            control_self_add=True,
            key_states_skipped=False,
            value_states_skipped=False,
            output_states_skipped=False,
            **kwargs):
        super().__init__(
            hidden_size, 
            cross_attention_dim, 
            rank, 
            post_add=post_add,
            key_states_skipped=key_states_skipped,
            value_states_skipped=value_states_skipped,
            output_states_skipped=output_states_skipped)

        control_rank = rank if control_rank is None else control_rank
        control_channels = hidden_size if control_channels is None else control_channels
        self.concat_hidden = concat_hidden
        self.control_self_add = control_self_add if control_channels is None else False
        self.control_states: torch.Tensor = None

        self.to_control = LoRALinearLayer(
            control_channels + (hidden_size if concat_hidden else 0), 
            hidden_size, 
            control_rank)
        self.pre_loras: List[LoRACrossAttnProcessor] = []
        self.post_loras: List[LoRACrossAttnProcessor] = []

    def inject_pre_lora(self, lora_layer):
        self.pre_loras.append(lora_layer)
    
    def inject_post_lora(self, lora_layer):
        self.post_loras.append(lora_layer)

    def inject_control_states(self, control_states):
        self.control_states = control_states

    def process_control_states(self, hidden_states, scale=1.0):
        control_states = self.control_states.to(hidden_states.dtype)
        if hidden_states.ndim == 3 and control_states.ndim == 4:
            batch, _, height, width = control_states.shape
            control_states = control_states.permute(0, 2, 3, 1).reshape(batch, height * width, -1)
            self.control_states = control_states
        _control_states = control_states
        if self.concat_hidden:
            b1, b2 = control_states.shape[0], hidden_states.shape[0]
            if b1 != b2:
                control_states = control_states[:,None].repeat(1, b2//b1, *([1]*(len(control_states.shape)-1)))
                control_states = control_states.view(-1, *control_states.shape[2:])
            _control_states = torch.cat([hidden_states, control_states], -1)
        _control_states = scale * self.to_control(_control_states)
        if self.control_self_add:
            control_states = control_states + _control_states
        else:
            control_states = _control_states

        return control_states

    def __call__(
        self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None, scale=1.0
    ):
        pre_lora: LoRACrossAttnProcessor
        post_lora: LoRACrossAttnProcessor
        assert self.control_states is not None

        batch_size, sequence_length, _ = hidden_states.shape
        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)
        query = attn.to_q(hidden_states)
        for pre_lora in self.pre_loras:
            lora_in = query if pre_lora.post_add else hidden_states
            if isinstance(pre_lora, ControlLoRACrossAttnProcessor):
                lora_in = lora_in + pre_lora.process_control_states(hidden_states, scale)
            query = query + scale * pre_lora.to_q_lora(lora_in)
        query = query + scale * self.to_q_lora((
            query if self.post_add else hidden_states) + self.process_control_states(hidden_states, scale))
        for post_lora in self.post_loras:
            lora_in = query if post_lora.post_add else hidden_states
            if isinstance(post_lora, ControlLoRACrossAttnProcessor):
                lora_in = lora_in + post_lora.process_control_states(hidden_states, scale)
            query = query + scale * post_lora.to_q_lora(lora_in)
        query = attn.head_to_batch_dim(query)

        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states

        key = attn.to_k(encoder_hidden_states)
        for pre_lora in self.pre_loras:
            if not pre_lora.key_states_skipped:
                key = key + scale * pre_lora.to_k_lora(key if pre_lora.post_add else encoder_hidden_states)
        if not self.key_states_skipped:
            key = key + scale * self.to_k_lora(key if self.post_add else encoder_hidden_states)
        for post_lora in self.post_loras:
            if not post_lora.key_states_skipped:
                key = key + scale * post_lora.to_k_lora(key if post_lora.post_add else encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)
        for pre_lora in self.pre_loras:
            if not pre_lora.value_states_skipped:
                value = value + pre_lora.to_v_lora(value if pre_lora.post_add else encoder_hidden_states)
        if not self.value_states_skipped:
            value = value + scale * self.to_v_lora(value if self.post_add else encoder_hidden_states)
        for post_lora in self.post_loras:
            if not post_lora.value_states_skipped:
                value = value + post_lora.to_v_lora(value if post_lora.post_add else encoder_hidden_states)

        key = attn.head_to_batch_dim(key)
        value = attn.head_to_batch_dim(value)

        attention_probs = attn.get_attention_scores(query, key, attention_mask)
        hidden_states = torch.bmm(attention_probs, value)
        hidden_states = attn.batch_to_head_dim(hidden_states)

        # linear proj
        out = attn.to_out[0](hidden_states)
        for pre_lora in self.pre_loras:
            if not pre_lora.output_states_skipped:
                out = out + scale * pre_lora.to_out_lora(out if pre_lora.post_add else hidden_states)
        out = out + scale * self.to_out_lora(out if self.post_add else hidden_states)
        for post_lora in self.post_loras:
            if not post_lora.output_states_skipped:
                out = out + scale * post_lora.to_out_lora(out if post_lora.post_add else hidden_states)
        hidden_states = out
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        return hidden_states




class ControlLoRACrossAttnProcessorV2(LoRACrossAttnProcessor):
    def __init__(
            self, 
            hidden_size, 
            cross_attention_dim=None, 
            rank=4, 
            control_rank=None, 
            control_channels=None,
            **kwargs):
        super().__init__(
            hidden_size, 
            cross_attention_dim, 
            rank, 
            post_add=False,
            key_states_skipped=True,
            value_states_skipped=True,
            output_states_skipped=False)

        control_rank = rank if control_rank is None else control_rank
        control_channels = hidden_size if control_channels is None else control_channels
        self.concat_hidden = True
        self.control_self_add = False
        self.control_states: torch.Tensor = None

        self.to_control = LoRALinearLayer(
            hidden_size + control_channels, 
            hidden_size, 
            control_rank)
        self.to_control_out = LoRALinearLayer(
            hidden_size + control_channels, 
            hidden_size, 
            control_rank)
        self.pre_loras: List[LoRACrossAttnProcessor] = []
        self.post_loras: List[LoRACrossAttnProcessor] = []

    def inject_pre_lora(self, lora_layer):
        self.pre_loras.append(lora_layer)
    
    def inject_post_lora(self, lora_layer):
        self.post_loras.append(lora_layer)

    def inject_control_states(self, control_states):
        self.control_states = control_states

    def process_control_states(self, hidden_states, scale=1.0, is_out=False):
        control_states = self.control_states.to(hidden_states.dtype)
        if hidden_states.ndim == 3 and control_states.ndim == 4:
            batch, _, height, width = control_states.shape
            control_states = control_states.permute(0, 2, 3, 1).reshape(batch, height * width, -1)
            self.control_states = control_states
        _control_states = control_states
        if self.concat_hidden:
            b1, b2 = control_states.shape[0], hidden_states.shape[0]
            if b1 != b2:
                control_states = control_states[:,None].repeat(1, b2//b1, *([1]*(len(control_states.shape)-1)))
                control_states = control_states.view(-1, *control_states.shape[2:])
            _control_states = torch.cat([hidden_states, control_states], -1)
        _control_states = scale * (self.to_control_out if is_out else self.to_control)(_control_states)
        if self.control_self_add:
            control_states = control_states + _control_states
        else:
            control_states = _control_states

        return control_states

    def __call__(
        self, attn: CrossAttention, hidden_states, encoder_hidden_states=None, attention_mask=None, scale=1.0
    ):
        pre_lora: LoRACrossAttnProcessor
        post_lora: LoRACrossAttnProcessor
        assert self.control_states is not None

        batch_size, sequence_length, _ = hidden_states.shape
        attention_mask = attn.prepare_attention_mask(attention_mask, sequence_length)
        for pre_lora in self.pre_loras:
            if isinstance(pre_lora, ControlLoRACrossAttnProcessorV2):
                hidden_states = hidden_states + pre_lora.process_control_states(hidden_states, scale)
        hidden_states = hidden_states + self.process_control_states(hidden_states, scale)
        for post_lora in self.post_loras:
            if isinstance(post_lora, ControlLoRACrossAttnProcessorV2):
                hidden_states = hidden_states + post_lora.process_control_states(hidden_states, scale)
        query = attn.to_q(hidden_states)
        for pre_lora in self.pre_loras:
            lora_in = query if pre_lora.post_add else hidden_states
            query = query + scale * pre_lora.to_q_lora(lora_in)
        query = query + scale * self.to_q_lora(query if self.post_add else hidden_states)
        for post_lora in self.post_loras:
            lora_in = query if post_lora.post_add else hidden_states
            query = query + scale * post_lora.to_q_lora(lora_in)
        query = attn.head_to_batch_dim(query)

        encoder_hidden_states = encoder_hidden_states if encoder_hidden_states is not None else hidden_states

        key = attn.to_k(encoder_hidden_states)
        for pre_lora in self.pre_loras:
            if not pre_lora.key_states_skipped:
                key = key + scale * pre_lora.to_k_lora(key if pre_lora.post_add else encoder_hidden_states)
        if not self.key_states_skipped:
            key = key + scale * self.to_k_lora(key if self.post_add else encoder_hidden_states)
        for post_lora in self.post_loras:
            if not post_lora.key_states_skipped:
                key = key + scale * post_lora.to_k_lora(key if post_lora.post_add else encoder_hidden_states)
        value = attn.to_v(encoder_hidden_states)
        for pre_lora in self.pre_loras:
            if not pre_lora.value_states_skipped:
                value = value + pre_lora.to_v_lora(value if pre_lora.post_add else encoder_hidden_states)
        if not self.value_states_skipped:
            value = value + scale * self.to_v_lora(value if self.post_add else encoder_hidden_states)
        for post_lora in self.post_loras:
            if not post_lora.value_states_skipped:
                value = value + post_lora.to_v_lora(value if post_lora.post_add else encoder_hidden_states)

        key = attn.head_to_batch_dim(key)
        value = attn.head_to_batch_dim(value)

        attention_probs = attn.get_attention_scores(query, key, attention_mask)
        hidden_states = torch.bmm(attention_probs, value)
        hidden_states = attn.batch_to_head_dim(hidden_states)

        # linear proj
        for pre_lora in self.pre_loras:
            if isinstance(pre_lora, ControlLoRACrossAttnProcessorV2):
                hidden_states = hidden_states + pre_lora.process_control_states(hidden_states, scale, is_out=True)
        hidden_states = hidden_states + self.process_control_states(hidden_states, scale, is_out=True)
        for post_lora in self.post_loras:
            if isinstance(post_lora, ControlLoRACrossAttnProcessorV2):
                hidden_states = hidden_states + post_lora.process_control_states(hidden_states, scale, is_out=True)
        out = attn.to_out[0](hidden_states)
        for pre_lora in self.pre_loras:
            if not pre_lora.output_states_skipped:
                out = out + scale * pre_lora.to_out_lora(out if pre_lora.post_add else hidden_states)
        out = out + scale * self.to_out_lora(out if self.post_add else hidden_states)
        for post_lora in self.post_loras:
            if not post_lora.output_states_skipped:
                out = out + scale * post_lora.to_out_lora(out if post_lora.post_add else hidden_states)
        hidden_states = out
        # dropout
        hidden_states = attn.to_out[1](hidden_states)

        return hidden_states


class ConvBlock2D(nn.Module):
    def __init__(
        self,
        *,
        in_channels,
        out_channels=None,
        conv_kernel_size=3,
        dropout=0.0,
        temb_channels=512,
        groups=32,
        groups_out=None,
        pre_norm=True,
        eps=1e-6,
        non_linearity="swish",
        time_embedding_norm="default",
        kernel=None,
        output_scale_factor=1.0,
        up=False,
        down=False,
    ):
        super().__init__()
        self.pre_norm = pre_norm
        self.pre_norm = True
        self.in_channels = in_channels
        out_channels = in_channels if out_channels is None else out_channels
        self.out_channels = out_channels
        self.time_embedding_norm = time_embedding_norm
        self.up = up
        self.down = down
        self.output_scale_factor = output_scale_factor

        if groups_out is None:
            groups_out = groups

        self.norm1 = torch.nn.GroupNorm(num_groups=groups, num_channels=in_channels, eps=eps, affine=True)

        self.conv1 = torch.nn.Conv2d(in_channels, out_channels, kernel_size=conv_kernel_size, stride=1, padding=conv_kernel_size//2)

        if temb_channels is not None:
            if self.time_embedding_norm == "default":
                time_emb_proj_out_channels = out_channels
            elif self.time_embedding_norm == "scale_shift":
                time_emb_proj_out_channels = out_channels * 2
            else:
                raise ValueError(f"unknown time_embedding_norm : {self.time_embedding_norm} ")

            self.time_emb_proj = torch.nn.Linear(temb_channels, time_emb_proj_out_channels)
        else:
            self.time_emb_proj = None

        self.norm2 = torch.nn.GroupNorm(num_groups=groups_out, num_channels=out_channels, eps=eps, affine=True)
        self.dropout = torch.nn.Dropout(dropout)

        if non_linearity == "swish":
            self.nonlinearity = lambda x: F.silu(x)
        elif non_linearity == "mish":
            self.nonlinearity = Mish()
        elif non_linearity == "silu":
            self.nonlinearity = nn.SiLU()

        self.upsample = self.downsample = None
        if self.up:
            if kernel == "fir":
                fir_kernel = (1, 3, 3, 1)
                self.upsample = lambda x: upsample_2d(x, kernel=fir_kernel)
            elif kernel == "sde_vp":
                self.upsample = partial(F.interpolate, scale_factor=2.0, mode="nearest")
            else:
                self.upsample = Upsample2D(in_channels, use_conv=False)
        elif self.down:
            if kernel == "fir":
                fir_kernel = (1, 3, 3, 1)
                self.downsample = lambda x: downsample_2d(x, kernel=fir_kernel)
            elif kernel == "sde_vp":
                self.downsample = partial(F.avg_pool2d, kernel_size=2, stride=2)
            else:
                self.downsample = Downsample2D(in_channels, use_conv=False, padding=1, name="op")

    def forward(self, input_tensor, temb):
        hidden_states = input_tensor

        hidden_states = self.norm1(hidden_states)
        hidden_states = self.nonlinearity(hidden_states)

        if self.upsample is not None:
            # upsample_nearest_nhwc fails with large batch sizes. see https://github.com/huggingface/diffusers/issues/984
            if hidden_states.shape[0] >= 64:
                input_tensor = input_tensor.contiguous()
                hidden_states = hidden_states.contiguous()
            input_tensor = self.upsample(input_tensor)
            hidden_states = self.upsample(hidden_states)
        elif self.downsample is not None:
            input_tensor = self.downsample(input_tensor)
            hidden_states = self.downsample(hidden_states)

        hidden_states = self.conv1(hidden_states)

        if temb is not None:
            temb = self.time_emb_proj(self.nonlinearity(temb))[:, :, None, None]

        if temb is not None and self.time_embedding_norm == "default":
            hidden_states = hidden_states + temb

        hidden_states = self.norm2(hidden_states)

        if temb is not None and self.time_embedding_norm == "scale_shift":
            scale, shift = torch.chunk(temb, 2, dim=1)
            hidden_states = hidden_states * (1 + scale) + shift

        hidden_states = self.nonlinearity(hidden_states)

        output_tensor = self.dropout(hidden_states)

        return output_tensor


class SimpleDownEncoderBlock2D(nn.Module):
    def __init__(
        self,
        in_channels: int,
        out_channels: int,
        dropout: float = 0.0,
        num_layers: int = 1,
        convnet_eps: float = 1e-6,
        convnet_time_scale_shift: str = "default",
        convnet_act_fn: str = "swish",
        convnet_groups: int = 32,
        convnet_pre_norm: bool = True,
        convnet_kernel_size: int = 3,
        output_scale_factor=1.0,
        add_downsample=True,
        downsample_padding=1,
    ):
        super().__init__()
        convnets = []

        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            convnets.append(
                ConvBlock2D(
                    in_channels=in_channels,
                    out_channels=out_channels,
                    temb_channels=None,
                    eps=convnet_eps,
                    groups=convnet_groups,
                    dropout=dropout,
                    time_embedding_norm=convnet_time_scale_shift,
                    non_linearity=convnet_act_fn,
                    output_scale_factor=output_scale_factor,
                    pre_norm=convnet_pre_norm,
                    conv_kernel_size=convnet_kernel_size,
                )
            )
        in_channels = in_channels if num_layers == 0 else out_channels

        self.convnets = nn.ModuleList(convnets)

        if add_downsample:
            self.downsamplers = nn.ModuleList(
                [
                    Downsample2D(
                        in_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
                    )
                ]
            )
        else:
            self.downsamplers = None

    def forward(self, hidden_states):
        for convnet in self.convnets:
            hidden_states = convnet(hidden_states, temb=None)

        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
                hidden_states = downsampler(hidden_states)

        return hidden_states


@dataclass
class ControlLoRAOutput(BaseOutput):
    control_states: Tuple[torch.FloatTensor]


class ControlLoRA(ModelMixin, ConfigMixin):
    @register_to_config
    def __init__(
        self,
        in_channels: int = 3,
        down_block_types: Tuple[str] = (
            "SimpleDownEncoderBlock2D",
            "SimpleDownEncoderBlock2D",
            "SimpleDownEncoderBlock2D",
            "SimpleDownEncoderBlock2D",
        ),
        block_out_channels: Tuple[int] = (32, 64, 128, 256),
        layers_per_block: int = 1,
        act_fn: str = "silu",
        norm_num_groups: int = 32,
        lora_pre_down_block_types: Tuple[str] = (
            None,
            "SimpleDownEncoderBlock2D",
            "SimpleDownEncoderBlock2D",
            "SimpleDownEncoderBlock2D",
        ),
        lora_pre_down_layers_per_block: int = 1,
        lora_pre_conv_skipped: bool = False,
        lora_pre_conv_types: Tuple[str] = (
            "SimpleDownEncoderBlock2D",
            "SimpleDownEncoderBlock2D",
            "SimpleDownEncoderBlock2D",
            "SimpleDownEncoderBlock2D",
        ),
        lora_pre_conv_layers_per_block: int = 1,
        lora_pre_conv_layers_kernel_size: int = 1,
        lora_block_in_channels: Tuple[int] = (256, 256, 256, 256),
        lora_block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
        lora_cross_attention_dims: Tuple[List[int]] = (
            [None, 768, None, 768, None, 768, None, 768, None, 768], 
            [None, 768, None, 768, None, 768, None, 768, None, 768], 
            [None, 768, None, 768, None, 768, None, 768, None, 768], 
            [None, 768]
        ),
        lora_rank: int = 4,
        lora_control_rank: int = None,
        lora_post_add: bool = False,
        lora_concat_hidden: bool = False,
        lora_control_channels: Tuple[int] = (None, None, None, None),
        lora_control_self_add: bool = True,
        lora_key_states_skipped: bool = False,
        lora_value_states_skipped: bool = False,
        lora_output_states_skipped: bool = False,
        lora_control_version: int = 1
    ):
        super().__init__()

        lora_control_cls = ControlLoRACrossAttnProcessor
        if lora_control_version == 2:
            lora_control_cls = ControlLoRACrossAttnProcessorV2

        assert lora_block_in_channels[0] == block_out_channels[-1]
        
        if lora_pre_conv_skipped:
            lora_control_channels = lora_block_in_channels
            lora_control_self_add = False

        self.layers_per_block = layers_per_block
        self.lora_pre_down_layers_per_block = lora_pre_down_layers_per_block
        self.lora_pre_conv_layers_per_block = lora_pre_conv_layers_per_block

        self.conv_in = torch.nn.Conv2d(in_channels, block_out_channels[0], kernel_size=3, stride=1, padding=1)

        self.down_blocks = nn.ModuleList([])
        self.pre_lora_layers = nn.ModuleList([])
        self.lora_layers = nn.ModuleList([])

        # pre_down
        pre_down_blocks = []
        output_channel = block_out_channels[0]
        for i, down_block_type in enumerate(down_block_types):
            input_channel = output_channel
            output_channel = block_out_channels[i]
            is_final_block = i == len(block_out_channels) - 1

            pre_down_block = get_down_block(
                down_block_type,
                num_layers=self.layers_per_block,
                in_channels=input_channel,
                out_channels=output_channel,
                add_downsample=not is_final_block,
                resnet_eps=1e-6,
                downsample_padding=0,
                resnet_act_fn=act_fn,
                resnet_groups=norm_num_groups,
                attn_num_head_channels=None,
                temb_channels=None,
            )
            pre_down_blocks.append(pre_down_block)
        self.down_blocks.append(nn.Sequential(*pre_down_blocks))
        self.pre_lora_layers.append(
            get_down_block(
                lora_pre_conv_types[0],
                num_layers=self.lora_pre_conv_layers_per_block,
                in_channels=lora_block_in_channels[0],
                out_channels=(
                    lora_block_out_channels[0] 
                    if lora_control_channels[0] is None 
                    else lora_control_channels[0]),
                add_downsample=False,
                resnet_eps=1e-6,
                downsample_padding=0,
                resnet_act_fn=act_fn,
                resnet_groups=norm_num_groups,
                attn_num_head_channels=None,
                temb_channels=None,
                resnet_kernel_size=lora_pre_conv_layers_kernel_size,
            ) if not lora_pre_conv_skipped else nn.Identity()
        )
        self.lora_layers.append(
            nn.ModuleList([
                lora_control_cls(
                    lora_block_out_channels[0], 
                    cross_attention_dim=cross_attention_dim, 
                    rank=lora_rank, 
                    control_rank=lora_control_rank,
                    post_add=lora_post_add,
                    concat_hidden=lora_concat_hidden,
                    control_channels=lora_control_channels[0],
                    control_self_add=lora_control_self_add,
                    key_states_skipped=lora_key_states_skipped,
                    value_states_skipped=lora_value_states_skipped,
                    output_states_skipped=lora_output_states_skipped)
                for cross_attention_dim in lora_cross_attention_dims[0]
            ])
        )
        
        # down
        output_channel = lora_block_in_channels[0]
        for i, down_block_type in enumerate(lora_pre_down_block_types):
            if i == 0:
                continue
            input_channel = output_channel
            output_channel = lora_block_in_channels[i]

            down_block = get_down_block(
                down_block_type,
                num_layers=self.lora_pre_down_layers_per_block,
                in_channels=input_channel,
                out_channels=output_channel,
                add_downsample=True,
                resnet_eps=1e-6,
                downsample_padding=0,
                resnet_act_fn=act_fn,
                resnet_groups=norm_num_groups,
                attn_num_head_channels=None,
                temb_channels=None,
            )
            self.down_blocks.append(down_block)

            self.pre_lora_layers.append(
                get_down_block(
                    lora_pre_conv_types[i],
                    num_layers=self.lora_pre_conv_layers_per_block,
                    in_channels=output_channel,
                    out_channels=(
                        lora_block_out_channels[i] 
                        if lora_control_channels[i] is None 
                        else lora_control_channels[i]),
                    add_downsample=False,
                    resnet_eps=1e-6,
                    downsample_padding=0,
                    resnet_act_fn=act_fn,
                    resnet_groups=norm_num_groups,
                    attn_num_head_channels=None,
                    temb_channels=None,
                    resnet_kernel_size=lora_pre_conv_layers_kernel_size,
                ) if not lora_pre_conv_skipped else nn.Identity()
            )
            self.lora_layers.append(
                nn.ModuleList([
                    lora_control_cls(
                        lora_block_out_channels[i], 
                        cross_attention_dim=cross_attention_dim, 
                        rank=lora_rank, 
                        control_rank=lora_control_rank,
                        post_add=lora_post_add,
                        concat_hidden=lora_concat_hidden,
                        control_channels=lora_control_channels[i],
                        control_self_add=lora_control_self_add,
                        key_states_skipped=lora_key_states_skipped,
                        value_states_skipped=lora_value_states_skipped,
                        output_states_skipped=lora_output_states_skipped)
                    for cross_attention_dim in lora_cross_attention_dims[i]
                ])
            )

    def forward(self, x: torch.FloatTensor, return_dict: bool = True) -> Union[ControlLoRAOutput, Tuple]:
        lora_layer: ControlLoRACrossAttnProcessor
        
        orig_dtype = x.dtype
        dtype = self.conv_in.weight.dtype

        h = x.to(dtype)
        h = self.conv_in(h)
        control_states_list = []

        # down
        for down_block, pre_lora_layer, lora_layer_list in zip(
            self.down_blocks, self.pre_lora_layers, self.lora_layers):
            h = down_block(h)
            control_states = pre_lora_layer(h)
            if isinstance(control_states, tuple):
                control_states = control_states[0]
            control_states = control_states.to(orig_dtype)
            for lora_layer in lora_layer_list:
                lora_layer.inject_control_states(control_states)
            control_states_list.append(control_states)

        if not return_dict:
            return tuple(control_states_list)

        return ControlLoRAOutput(control_states=tuple(control_states_list))