'''
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
'''
from ..resnet import Downsample2D, FirDownsample2D, FirUpsample2D, KDownsample2D, KUpsample2D, ResnetBlock2D, ResnetBlockCondNorm2D, Upsample2D
from ..attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0
from ..sapiens_transformers.dual_transformer_2d import DualTransformer2DModel
from ..sapiens_transformers.transformer_2d import Transformer2DModel
from ...utils import deprecate, is_torch_version
from typing import Any, Dict, Optional, Tuple, Union
from ...utils.torch_utils import apply_freeu
from ..activations import get_activation
from ..normalization import AdaGroupNorm
import torch.nn.functional as F
from torch import nn
import numpy as np
import torch
def get_down_block(down_block_type: str, num_layers: int, in_channels: int, out_channels: int, temb_channels: int, add_downsample: bool, resnet_eps: float, resnet_act_fn: str,
transformer_layers_per_block: int=1, num_attention_heads: Optional[int]=None, resnet_groups: Optional[int]=None, cross_attention_dim: Optional[int]=None, downsample_padding: Optional[int]=None,
dual_cross_attention: bool=False, use_linear_projection: bool=False, only_cross_attention: bool=False, upcast_attention: bool=False, resnet_time_scale_shift: str='default', attention_type: str='default',
resnet_skip_time_act: bool=False, resnet_out_scale_factor: float=1.0, cross_attention_norm: Optional[str]=None, attention_head_dim: Optional[int]=None, downsample_type: Optional[str]=None, dropout: float=0.0):
    if attention_head_dim is None: attention_head_dim = num_attention_heads
    down_block_type = down_block_type[7:] if down_block_type.startswith('UNetRes') else down_block_type
    if down_block_type == 'DownBlock2D': return DownBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, dropout=dropout,
    add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, downsample_padding=downsample_padding, resnet_time_scale_shift=resnet_time_scale_shift)
    elif down_block_type == 'ResnetDownsampleBlock2D': return ResnetDownsampleBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels,
    dropout=dropout, add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, resnet_time_scale_shift=resnet_time_scale_shift,
    skip_time_act=resnet_skip_time_act, output_scale_factor=resnet_out_scale_factor)
    elif down_block_type == 'AttnDownBlock2D':
        if add_downsample is False: downsample_type = None
        else: downsample_type = downsample_type or 'conv'
        return AttnDownBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, dropout=dropout, resnet_eps=resnet_eps,
        resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, downsample_padding=downsample_padding, attention_head_dim=attention_head_dim,
        resnet_time_scale_shift=resnet_time_scale_shift, downsample_type=downsample_type)
    elif down_block_type == 'CrossAttnDownBlock2D':
        if cross_attention_dim is None: raise ValueError('cross_attention_dim must be specified for CrossAttnDownBlock2D')
        return CrossAttnDownBlock2D(num_layers=num_layers, transformer_layers_per_block=transformer_layers_per_block, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels,
        dropout=dropout, add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, downsample_padding=downsample_padding, cross_attention_dim=cross_attention_dim,
        num_attention_heads=num_attention_heads, 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, attention_type=attention_type)
    elif down_block_type == 'SimpleCrossAttnDownBlock2D':
        if cross_attention_dim is None: raise ValueError('cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D')
        return SimpleCrossAttnDownBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, dropout=dropout, add_downsample=add_downsample,
        resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, cross_attention_dim=cross_attention_dim, attention_head_dim=attention_head_dim, resnet_time_scale_shift=resnet_time_scale_shift,
        skip_time_act=resnet_skip_time_act, output_scale_factor=resnet_out_scale_factor, only_cross_attention=only_cross_attention, cross_attention_norm=cross_attention_norm)
    elif down_block_type == 'SkipDownBlock2D': return SkipDownBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, dropout=dropout,
    add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, downsample_padding=downsample_padding, resnet_time_scale_shift=resnet_time_scale_shift)
    elif down_block_type == 'AttnSkipDownBlock2D': return AttnSkipDownBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, dropout=dropout,
    add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, attention_head_dim=attention_head_dim, resnet_time_scale_shift=resnet_time_scale_shift)
    elif down_block_type == 'DownEncoderBlock2D': return DownEncoderBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, dropout=dropout, add_downsample=add_downsample,
    resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, downsample_padding=downsample_padding, resnet_time_scale_shift=resnet_time_scale_shift)
    elif down_block_type == 'AttnDownEncoderBlock2D': return AttnDownEncoderBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, dropout=dropout, add_downsample=add_downsample,
    resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, downsample_padding=downsample_padding, attention_head_dim=attention_head_dim, resnet_time_scale_shift=resnet_time_scale_shift)
    elif down_block_type == 'KDownBlock2D': return KDownBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, dropout=dropout,
    add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn)
    elif down_block_type == 'KCrossAttnDownBlock2D': return KCrossAttnDownBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels,
    dropout=dropout, add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, cross_attention_dim=cross_attention_dim,
    attention_head_dim=attention_head_dim, add_self_attention=True if not add_downsample else False)
    raise ValueError(f'{down_block_type} does not exist.')
def get_mid_block(mid_block_type: str, temb_channels: int, in_channels: int, resnet_eps: float, resnet_act_fn: str, resnet_groups: int, output_scale_factor: float=1.0, transformer_layers_per_block: int=1,
num_attention_heads: Optional[int]=None, cross_attention_dim: Optional[int]=None, dual_cross_attention: bool=False, use_linear_projection: bool=False, mid_block_only_cross_attention: bool=False,
upcast_attention: bool=False, resnet_time_scale_shift: str='default', attention_type: str='default', resnet_skip_time_act: bool=False, cross_attention_norm: Optional[str]=None, attention_head_dim: Optional[int]=1, dropout: float=0.0):
    if mid_block_type == 'UNetMidBlock2DCrossAttn': return UNetMidBlock2DCrossAttn(transformer_layers_per_block=transformer_layers_per_block, in_channels=in_channels, temb_channels=temb_channels,
    dropout=dropout, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, output_scale_factor=output_scale_factor, resnet_time_scale_shift=resnet_time_scale_shift, cross_attention_dim=cross_attention_dim,
    num_attention_heads=num_attention_heads, resnet_groups=resnet_groups, dual_cross_attention=dual_cross_attention, use_linear_projection=use_linear_projection, upcast_attention=upcast_attention, attention_type=attention_type)
    elif mid_block_type == 'UNetMidBlock2DSimpleCrossAttn': return UNetMidBlock2DSimpleCrossAttn(in_channels=in_channels, temb_channels=temb_channels, dropout=dropout, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn,
    output_scale_factor=output_scale_factor, cross_attention_dim=cross_attention_dim, attention_head_dim=attention_head_dim, resnet_groups=resnet_groups, resnet_time_scale_shift=resnet_time_scale_shift,
    skip_time_act=resnet_skip_time_act, only_cross_attention=mid_block_only_cross_attention, cross_attention_norm=cross_attention_norm)
    elif mid_block_type == 'UNetMidBlock2D': return UNetMidBlock2D(in_channels=in_channels, temb_channels=temb_channels, dropout=dropout, num_layers=0, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn,
    output_scale_factor=output_scale_factor, resnet_groups=resnet_groups, resnet_time_scale_shift=resnet_time_scale_shift, add_attention=False)
    elif mid_block_type is None: return None
    else: raise ValueError(f'unknown mid_block_type : {mid_block_type}')
def get_up_block(up_block_type: str, num_layers: int, in_channels: int, out_channels: int, prev_output_channel: int, temb_channels: int, add_upsample: bool, resnet_eps: float, resnet_act_fn: str,
resolution_idx: Optional[int]=None, transformer_layers_per_block: int=1, num_attention_heads: Optional[int]=None, resnet_groups: Optional[int]=None, cross_attention_dim: Optional[int]=None,
dual_cross_attention: bool=False, use_linear_projection: bool=False, only_cross_attention: bool=False, upcast_attention: bool=False, resnet_time_scale_shift: str='default', attention_type: str='default',
resnet_skip_time_act: bool=False, resnet_out_scale_factor: float=1.0, cross_attention_norm: Optional[str]=None, attention_head_dim: Optional[int]=None, upsample_type: Optional[str]=None, dropout: float=0.0) -> nn.Module:
    if attention_head_dim is None: attention_head_dim = num_attention_heads
    up_block_type = up_block_type[7:] if up_block_type.startswith('UNetRes') else up_block_type
    if up_block_type == 'UpBlock2D': return UpBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel, temb_channels=temb_channels, resolution_idx=resolution_idx,
    dropout=dropout, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, resnet_time_scale_shift=resnet_time_scale_shift)
    elif up_block_type == 'ResnetUpsampleBlock2D': return ResnetUpsampleBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel, temb_channels=temb_channels,
    resolution_idx=resolution_idx, dropout=dropout, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, resnet_time_scale_shift=resnet_time_scale_shift,
    skip_time_act=resnet_skip_time_act, output_scale_factor=resnet_out_scale_factor)
    elif up_block_type == 'CrossAttnUpBlock2D':
        if cross_attention_dim is None: raise ValueError('cross_attention_dim must be specified for CrossAttnUpBlock2D')
        return CrossAttnUpBlock2D(num_layers=num_layers, transformer_layers_per_block=transformer_layers_per_block, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel,
        temb_channels=temb_channels, resolution_idx=resolution_idx, dropout=dropout, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups,
        cross_attention_dim=cross_attention_dim, num_attention_heads=num_attention_heads, 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, attention_type=attention_type)
    elif up_block_type == 'SimpleCrossAttnUpBlock2D':
        if cross_attention_dim is None: raise ValueError('cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D')
        return SimpleCrossAttnUpBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel, temb_channels=temb_channels,
        resolution_idx=resolution_idx, dropout=dropout, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups,
        cross_attention_dim=cross_attention_dim, attention_head_dim=attention_head_dim, resnet_time_scale_shift=resnet_time_scale_shift,
        skip_time_act=resnet_skip_time_act, output_scale_factor=resnet_out_scale_factor, only_cross_attention=only_cross_attention, cross_attention_norm=cross_attention_norm)
    elif up_block_type == 'AttnUpBlock2D':
        if add_upsample is False: upsample_type = None
        else: upsample_type = upsample_type or 'conv'
        return AttnUpBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel, temb_channels=temb_channels, resolution_idx=resolution_idx,
        dropout=dropout, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, attention_head_dim=attention_head_dim, resnet_time_scale_shift=resnet_time_scale_shift, upsample_type=upsample_type)
    elif up_block_type == 'SkipUpBlock2D': return SkipUpBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel, temb_channels=temb_channels,
    resolution_idx=resolution_idx, dropout=dropout, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_time_scale_shift=resnet_time_scale_shift)
    elif up_block_type == 'AttnSkipUpBlock2D': return AttnSkipUpBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel,
    temb_channels=temb_channels, resolution_idx=resolution_idx, dropout=dropout, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, attention_head_dim=attention_head_dim,
    resnet_time_scale_shift=resnet_time_scale_shift)
    elif up_block_type == 'UpDecoderBlock2D': return UpDecoderBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, resolution_idx=resolution_idx, dropout=dropout,
    add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, resnet_time_scale_shift=resnet_time_scale_shift, temb_channels=temb_channels)
    elif up_block_type == 'AttnUpDecoderBlock2D': return AttnUpDecoderBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, resolution_idx=resolution_idx, dropout=dropout,
    add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, attention_head_dim=attention_head_dim, resnet_time_scale_shift=resnet_time_scale_shift, temb_channels=temb_channels)
    elif up_block_type == 'KUpBlock2D': return KUpBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, resolution_idx=resolution_idx, dropout=dropout,
    add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn)
    elif up_block_type == 'KCrossAttnUpBlock2D': return KCrossAttnUpBlock2D(num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels,
    resolution_idx=resolution_idx, dropout=dropout, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, cross_attention_dim=cross_attention_dim, attention_head_dim=attention_head_dim)
    raise ValueError(f'{up_block_type} does not exist.')
class AutoencoderTinyBlock(nn.Module):
    """Returns:"""
    def __init__(self, in_channels: int, out_channels: int, act_fn: str):
        super().__init__()
        act_fn = get_activation(act_fn)
        self.conv = nn.Sequential(nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1), act_fn, nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1), act_fn, nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1))
        self.skip = nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False) if in_channels != out_channels else nn.Identity()
        self.fuse = nn.ReLU()
    def forward(self, x: torch.Tensor) -> torch.Tensor: return self.fuse(self.conv(x) + self.skip(x))
class UNetMidBlock2D(nn.Module):
    """Returns:"""
    def __init__(self, in_channels: int, temb_channels: int, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish', resnet_groups: int=32,
    attn_groups: Optional[int]=None, resnet_pre_norm: bool=True, add_attention: bool=True, attention_head_dim: int=1, output_scale_factor: float=1.0):
        super().__init__()
        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
        self.add_attention = add_attention
        if attn_groups is None: attn_groups = resnet_groups if resnet_time_scale_shift == 'default' else None
        if resnet_time_scale_shift == 'spatial': resnets = [ResnetBlockCondNorm2D(in_channels=in_channels, out_channels=in_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
        time_embedding_norm='spatial', non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor)]
        else: resnets = [ResnetBlock2D(in_channels=in_channels, out_channels=in_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift,
        non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm)]
        attentions = []
        if attention_head_dim is None: attention_head_dim = in_channels
        for _ in range(num_layers):
            if self.add_attention: attentions.append(Attention(in_channels, heads=in_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps,
            norm_num_groups=attn_groups, spatial_norm_dim=temb_channels if resnet_time_scale_shift == 'spatial' else None, residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True))
            else: attentions.append(None)
            if resnet_time_scale_shift == 'spatial': resnets.append(ResnetBlockCondNorm2D(in_channels=in_channels, out_channels=in_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm='spatial', non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor))
            else: resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=in_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift,
            non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
        self.gradient_checkpointing = False
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None) -> torch.Tensor:
        hidden_states = self.resnets[0](hidden_states, temb)
        for attn, resnet in zip(self.attentions, self.resnets[1:]):
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None: return module(*inputs, return_dict=return_dict)
                        else: return module(*inputs)
                    return custom_forward
                ckpt_kwargs: Dict[str, Any] = {'use_reentrant': False} if is_torch_version('>=', '1.11.0') else {}
                if attn is not None: hidden_states = attn(hidden_states, temb=temb)
                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs)
            else:
                if attn is not None: hidden_states = attn(hidden_states, temb=temb)
                hidden_states = resnet(hidden_states, temb)
        return hidden_states
class UNetMidBlock2DCrossAttn(nn.Module):
    def __init__(self, in_channels: int, temb_channels: int, out_channels: Optional[int]=None, dropout: float=0.0, num_layers: int=1, transformer_layers_per_block: Union[int, Tuple[int]]=1, resnet_eps: float=1e-06,
    resnet_time_scale_shift: str='default', resnet_act_fn: str='swish', resnet_groups: int=32, resnet_groups_out: Optional[int]=None, resnet_pre_norm: bool=True, num_attention_heads: int=1, output_scale_factor: float=1.0,
    cross_attention_dim: int=1280, dual_cross_attention: bool=False, use_linear_projection: bool=False, upcast_attention: bool=False, attention_type: str='default'):
        super().__init__()
        out_channels = out_channels or in_channels
        self.in_channels = in_channels
        self.out_channels = out_channels
        self.has_cross_attention = True
        self.num_attention_heads = num_attention_heads
        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
        if isinstance(transformer_layers_per_block, int): transformer_layers_per_block = [transformer_layers_per_block] * num_layers
        resnet_groups_out = resnet_groups_out or resnet_groups
        resnets = [ResnetBlock2D(in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, groups_out=resnet_groups_out, dropout=dropout,
        time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm)]
        attentions = []
        for i in range(num_layers):
            if not dual_cross_attention: attentions.append(Transformer2DModel(num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=transformer_layers_per_block[i],
            cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups_out, use_linear_projection=use_linear_projection, upcast_attention=upcast_attention, attention_type=attention_type))
            else: attentions.append(DualTransformer2DModel(num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups))
            resnets.append(ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups_out, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
        self.gradient_checkpointing = False
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None,
    cross_attention_kwargs: Optional[Dict[str, Any]]=None, encoder_attention_mask: Optional[torch.Tensor]=None) -> torch.Tensor:
        hidden_states = self.resnets[0](hidden_states, temb)
        for attn, resnet in zip(self.attentions, self.resnets[1:]):
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None: return module(*inputs, return_dict=return_dict)
                        else: return module(*inputs)
                    return custom_forward
                ckpt_kwargs: Dict[str, Any] = {'use_reentrant': False} if is_torch_version('>=', '1.11.0') else {}
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs,
                attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False)[0]
                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs)
            else:
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs,
                attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False)[0]
                hidden_states = resnet(hidden_states, temb)
        return hidden_states
class UNetMidBlock2DSimpleCrossAttn(nn.Module):
    def __init__(self, in_channels: int, temb_channels: int, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish',
    resnet_groups: int=32, resnet_pre_norm: bool=True, attention_head_dim: int=1, output_scale_factor: float=1.0, cross_attention_dim: int=1280,
    skip_time_act: bool=False, only_cross_attention: bool=False, cross_attention_norm: Optional[str]=None):
        super().__init__()
        self.has_cross_attention = True
        self.attention_head_dim = attention_head_dim
        resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
        self.num_heads = in_channels // self.attention_head_dim
        resnets = [ResnetBlock2D(in_channels=in_channels, out_channels=in_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift,
        non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=skip_time_act)]
        attentions = []
        for _ in range(num_layers):
            processor = AttnAddedKVProcessor2_0() if hasattr(F, 'scaled_dot_product_attention') else AttnAddedKVProcessor()
            attentions.append(Attention(query_dim=in_channels, cross_attention_dim=in_channels, heads=self.num_heads, dim_head=self.attention_head_dim, added_kv_proj_dim=cross_attention_dim,
            norm_num_groups=resnet_groups, bias=True, upcast_softmax=True, only_cross_attention=only_cross_attention, cross_attention_norm=cross_attention_norm, processor=processor))
            resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=in_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=skip_time_act))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None,
    cross_attention_kwargs: Optional[Dict[str, Any]]=None, encoder_attention_mask: Optional[torch.Tensor]=None) -> torch.Tensor:
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
        if attention_mask is None: mask = None if encoder_hidden_states is None else encoder_attention_mask
        else: mask = attention_mask
        hidden_states = self.resnets[0](hidden_states, temb)
        for attn, resnet in zip(self.attentions, self.resnets[1:]):
            hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=mask, **cross_attention_kwargs)
            hidden_states = resnet(hidden_states, temb)
        return hidden_states
class AttnDownBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default',
    resnet_act_fn: str='swish', resnet_groups: int=32, resnet_pre_norm: bool=True, attention_head_dim: int=1, output_scale_factor: float=1.0, downsample_padding: int=1, downsample_type: str='conv'):
        super().__init__()
        resnets = []
        attentions = []
        self.downsample_type = downsample_type
        if attention_head_dim is None: attention_head_dim = out_channels
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
            attentions.append(Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor,
            eps=resnet_eps, norm_num_groups=resnet_groups, residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
        if downsample_type == 'conv': self.downsamplers = nn.ModuleList([Downsample2D(out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name='op')])
        elif downsample_type == 'resnet': self.downsamplers = nn.ModuleList([ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps,
        groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, down=True)])
        else: self.downsamplers = None
        self.gradient_checkpointing = False
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, upsample_size: Optional[int]=None,
    cross_attention_kwargs: Optional[Dict[str, Any]]=None) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
        output_states = ()
        for resnet, attn in zip(self.resnets, self.attentions):
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None: return module(*inputs, return_dict=return_dict)
                        else: return module(*inputs)
                    return custom_forward
                ckpt_kwargs: Dict[str, Any] = {'use_reentrant': False} if is_torch_version('>=', '1.11.0') else {}
                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs)
                hidden_states = attn(hidden_states, **cross_attention_kwargs)
                output_states = output_states + (hidden_states,)
            else:
                hidden_states = resnet(hidden_states, temb)
                hidden_states = attn(hidden_states, **cross_attention_kwargs)
                output_states = output_states + (hidden_states,)
        if self.downsamplers is not None:
            for downsampler in self.downsamplers:
                if self.downsample_type == 'resnet': hidden_states = downsampler(hidden_states, temb=temb)
                else: hidden_states = downsampler(hidden_states)
            output_states += (hidden_states,)
        return (hidden_states, output_states)
class CrossAttnDownBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, dropout: float=0.0, num_layers: int=1, transformer_layers_per_block: Union[int, Tuple[int]]=1,
    resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish', resnet_groups: int=32, resnet_pre_norm: bool=True, num_attention_heads: int=1,
    cross_attention_dim: int=1280, output_scale_factor: float=1.0, downsample_padding: int=1, add_downsample: bool=True, dual_cross_attention: bool=False, use_linear_projection: bool=False,
    only_cross_attention: bool=False, upcast_attention: bool=False, attention_type: str='default'):
        super().__init__()
        resnets = []
        attentions = []
        self.has_cross_attention = True
        self.num_attention_heads = num_attention_heads
        if isinstance(transformer_layers_per_block, int): transformer_layers_per_block = [transformer_layers_per_block] * num_layers
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
            if not dual_cross_attention: attentions.append(Transformer2DModel(num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=transformer_layers_per_block[i],
            cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, attention_type=attention_type))
            else: attentions.append(DualTransformer2DModel(num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
        if add_downsample: self.downsamplers = nn.ModuleList([Downsample2D(out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name='op')])
        else: self.downsamplers = None
        self.gradient_checkpointing = False
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None,
    cross_attention_kwargs: Optional[Dict[str, Any]]=None, encoder_attention_mask: Optional[torch.Tensor]=None, additional_residuals: Optional[torch.Tensor]=None) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
        output_states = ()
        blocks = list(zip(self.resnets, self.attentions))
        for i, (resnet, attn) in enumerate(blocks):
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None: return module(*inputs, return_dict=return_dict)
                        else: return module(*inputs)
                    return custom_forward
                ckpt_kwargs: Dict[str, Any] = {'use_reentrant': False} if is_torch_version('>=', '1.11.0') else {}
                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False)[0]
            else:
                hidden_states = resnet(hidden_states, temb)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False)[0]
            if i == len(blocks) - 1 and additional_residuals is not None: hidden_states = hidden_states + additional_residuals
            output_states = output_states + (hidden_states,)
        if self.downsamplers is not None:
            for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states)
            output_states = output_states + (hidden_states,)
        return (hidden_states, output_states)
class DownBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish',
    resnet_groups: int=32, resnet_pre_norm: bool=True, output_scale_factor: float=1.0, add_downsample: bool=True, downsample_padding: int=1):
        super().__init__()
        resnets = []
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
        self.resnets = nn.ModuleList(resnets)
        if add_downsample: self.downsamplers = nn.ModuleList([Downsample2D(out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name='op')])
        else: self.downsamplers = None
        self.gradient_checkpointing = False
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, *args, **kwargs) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        output_states = ()
        for resnet in self.resnets:
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module):
                    def custom_forward(*inputs): return module(*inputs)
                    return custom_forward
                if is_torch_version('>=', '1.11.0'): hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, use_reentrant=False)
                else: hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
            else: hidden_states = resnet(hidden_states, temb)
            output_states = output_states + (hidden_states,)
        if self.downsamplers is not None:
            for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states)
            output_states = output_states + (hidden_states,)
        return (hidden_states, output_states)
class DownEncoderBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish',
    resnet_groups: int=32, resnet_pre_norm: bool=True, output_scale_factor: float=1.0, add_downsample: bool=True, downsample_padding: int=1):
        super().__init__()
        resnets = []
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            if resnet_time_scale_shift == 'spatial': resnets.append(ResnetBlockCondNorm2D(in_channels=in_channels, out_channels=out_channels, temb_channels=None, eps=resnet_eps,
            groups=resnet_groups, dropout=dropout, time_embedding_norm='spatial', non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor))
            else: resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=out_channels, temb_channels=None, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
        self.resnets = nn.ModuleList(resnets)
        if add_downsample: self.downsamplers = nn.ModuleList([Downsample2D(out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name='op')])
        else: self.downsamplers = None
    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        for resnet in self.resnets: hidden_states = resnet(hidden_states, temb=None)
        if self.downsamplers is not None:
            for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states)
        return hidden_states
class AttnDownEncoderBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish',
    resnet_groups: int=32, resnet_pre_norm: bool=True, attention_head_dim: int=1, output_scale_factor: float=1.0, add_downsample: bool=True, downsample_padding: int=1):
        super().__init__()
        resnets = []
        attentions = []
        if attention_head_dim is None: attention_head_dim = out_channels
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            if resnet_time_scale_shift == 'spatial': resnets.append(ResnetBlockCondNorm2D(in_channels=in_channels, out_channels=out_channels, temb_channels=None, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm='spatial', non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor))
            else: resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=out_channels, temb_channels=None, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift,
            non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
            attentions.append(Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps, norm_num_groups=resnet_groups,
            residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
        if add_downsample: self.downsamplers = nn.ModuleList([Downsample2D(out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name='op')])
        else: self.downsamplers = None
    def forward(self, hidden_states: torch.Tensor, *args, **kwargs) -> torch.Tensor:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        for resnet, attn in zip(self.resnets, self.attentions):
            hidden_states = resnet(hidden_states, temb=None)
            hidden_states = attn(hidden_states)
        if self.downsamplers is not None:
            for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states)
        return hidden_states
class AttnSkipDownBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish',
    resnet_pre_norm: bool=True, attention_head_dim: int=1, output_scale_factor: float=np.sqrt(2.0), add_downsample: bool=True):
        super().__init__()
        self.attentions = nn.ModuleList([])
        self.resnets = nn.ModuleList([])
        if attention_head_dim is None: attention_head_dim = out_channels
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            self.resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=min(in_channels // 4, 32),
            groups_out=min(out_channels // 4, 32), dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
            self.attentions.append(Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor,
            eps=resnet_eps, norm_num_groups=32, residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True))
        if add_downsample:
            self.resnet_down = ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=min(out_channels // 4, 32),
            dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_in_shortcut=True, down=True, kernel='fir')
            self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])
            self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1))
        else:
            self.resnet_down = None
            self.downsamplers = None
            self.skip_conv = None
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, skip_sample: Optional[torch.Tensor]=None, *args, **kwargs) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...], torch.Tensor]:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        output_states = ()
        for resnet, attn in zip(self.resnets, self.attentions):
            hidden_states = resnet(hidden_states, temb)
            hidden_states = attn(hidden_states)
            output_states += (hidden_states,)
        if self.downsamplers is not None:
            hidden_states = self.resnet_down(hidden_states, temb)
            for downsampler in self.downsamplers: skip_sample = downsampler(skip_sample)
            hidden_states = self.skip_conv(skip_sample) + hidden_states
            output_states += (hidden_states,)
        return (hidden_states, output_states, skip_sample)
class SkipDownBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish',
    resnet_pre_norm: bool=True, output_scale_factor: float=np.sqrt(2.0), add_downsample: bool=True, downsample_padding: int=1):
        super().__init__()
        self.resnets = nn.ModuleList([])
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            self.resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=min(in_channels // 4, 32),
            groups_out=min(out_channels // 4, 32), dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
        if add_downsample:
            self.resnet_down = ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=min(out_channels // 4, 32),
            dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_in_shortcut=True, down=True, kernel='fir')
            self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])
            self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1))
        else:
            self.resnet_down = None
            self.downsamplers = None
            self.skip_conv = None
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, skip_sample: Optional[torch.Tensor]=None, *args, **kwargs) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...], torch.Tensor]:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        output_states = ()
        for resnet in self.resnets:
            hidden_states = resnet(hidden_states, temb)
            output_states += (hidden_states,)
        if self.downsamplers is not None:
            hidden_states = self.resnet_down(hidden_states, temb)
            for downsampler in self.downsamplers: skip_sample = downsampler(skip_sample)
            hidden_states = self.skip_conv(skip_sample) + hidden_states
            output_states += (hidden_states,)
        return (hidden_states, output_states, skip_sample)
class ResnetDownsampleBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish',
    resnet_groups: int=32, resnet_pre_norm: bool=True, output_scale_factor: float=1.0, add_downsample: bool=True, skip_time_act: bool=False):
        super().__init__()
        resnets = []
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=skip_time_act))
        self.resnets = nn.ModuleList(resnets)
        if add_downsample: self.downsamplers = nn.ModuleList([ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps,
        groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor,
        pre_norm=resnet_pre_norm, skip_time_act=skip_time_act, down=True)])
        else: self.downsamplers = None
        self.gradient_checkpointing = False
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, *args, **kwargs) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        output_states = ()
        for resnet in self.resnets:
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module):
                    def custom_forward(*inputs): return module(*inputs)
                    return custom_forward
                if is_torch_version('>=', '1.11.0'): hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, use_reentrant=False)
                else: hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
            else: hidden_states = resnet(hidden_states, temb)
            output_states = output_states + (hidden_states,)
        if self.downsamplers is not None:
            for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states, temb)
            output_states = output_states + (hidden_states,)
        return (hidden_states, output_states)
class SimpleCrossAttnDownBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish',
    resnet_groups: int=32, resnet_pre_norm: bool=True, attention_head_dim: int=1, cross_attention_dim: int=1280, output_scale_factor: float=1.0, add_downsample: bool=True, skip_time_act: bool=False,
    only_cross_attention: bool=False, cross_attention_norm: Optional[str]=None):
        super().__init__()
        self.has_cross_attention = True
        resnets = []
        attentions = []
        self.attention_head_dim = attention_head_dim
        self.num_heads = out_channels // self.attention_head_dim
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            resnets.append(ResnetBlock2D(in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift,
            non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=skip_time_act))
            processor = AttnAddedKVProcessor2_0() if hasattr(F, 'scaled_dot_product_attention') else AttnAddedKVProcessor()
            attentions.append(Attention(query_dim=out_channels, cross_attention_dim=out_channels, heads=self.num_heads, dim_head=attention_head_dim, added_kv_proj_dim=cross_attention_dim, norm_num_groups=resnet_groups,
            bias=True, upcast_softmax=True, only_cross_attention=only_cross_attention, cross_attention_norm=cross_attention_norm, processor=processor))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
        if add_downsample: self.downsamplers = nn.ModuleList([ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
        time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=skip_time_act, down=True)])
        else: self.downsamplers = None
        self.gradient_checkpointing = False
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None,
    cross_attention_kwargs: Optional[Dict[str, Any]]=None, encoder_attention_mask: Optional[torch.Tensor]=None) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
        output_states = ()
        if attention_mask is None: mask = None if encoder_hidden_states is None else encoder_attention_mask
        else: mask = attention_mask
        for resnet, attn in zip(self.resnets, self.attentions):
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None: return module(*inputs, return_dict=return_dict)
                        else: return module(*inputs)
                    return custom_forward
                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=mask, **cross_attention_kwargs)
            else:
                hidden_states = resnet(hidden_states, temb)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=mask, **cross_attention_kwargs)
            output_states = output_states + (hidden_states,)
        if self.downsamplers is not None:
            for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states, temb)
            output_states = output_states + (hidden_states,)
        return (hidden_states, output_states)
class KDownBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, dropout: float=0.0, num_layers: int=4, resnet_eps: float=1e-05, resnet_act_fn: str='gelu', resnet_group_size: int=32, add_downsample: bool=False):
        super().__init__()
        resnets = []
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            groups = in_channels // resnet_group_size
            groups_out = out_channels // resnet_group_size
            resnets.append(ResnetBlockCondNorm2D(in_channels=in_channels, out_channels=out_channels, dropout=dropout, temb_channels=temb_channels, groups=groups, groups_out=groups_out,
            eps=resnet_eps, non_linearity=resnet_act_fn, time_embedding_norm='ada_group', conv_shortcut_bias=False))
        self.resnets = nn.ModuleList(resnets)
        if add_downsample: self.downsamplers = nn.ModuleList([KDownsample2D()])
        else: self.downsamplers = None
        self.gradient_checkpointing = False
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, *args, **kwargs) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        output_states = ()
        for resnet in self.resnets:
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module):
                    def custom_forward(*inputs): return module(*inputs)
                    return custom_forward
                if is_torch_version('>=', '1.11.0'): hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, use_reentrant=False)
                else: hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
            else: hidden_states = resnet(hidden_states, temb)
            output_states += (hidden_states,)
        if self.downsamplers is not None:
            for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states)
        return (hidden_states, output_states)
class KCrossAttnDownBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, cross_attention_dim: int, dropout: float=0.0, num_layers: int=4, resnet_group_size: int=32, add_downsample: bool=True,
    attention_head_dim: int=64, add_self_attention: bool=False, resnet_eps: float=1e-05, resnet_act_fn: str='gelu'):
        super().__init__()
        resnets = []
        attentions = []
        self.has_cross_attention = True
        for i in range(num_layers):
            in_channels = in_channels if i == 0 else out_channels
            groups = in_channels // resnet_group_size
            groups_out = out_channels // resnet_group_size
            resnets.append(ResnetBlockCondNorm2D(in_channels=in_channels, out_channels=out_channels, dropout=dropout, temb_channels=temb_channels, groups=groups, groups_out=groups_out,
            eps=resnet_eps, non_linearity=resnet_act_fn, time_embedding_norm='ada_group', conv_shortcut_bias=False))
            attentions.append(KAttentionBlock(out_channels, out_channels // attention_head_dim, attention_head_dim, cross_attention_dim=cross_attention_dim, temb_channels=temb_channels,
            attention_bias=True, add_self_attention=add_self_attention, cross_attention_norm='layer_norm', group_size=resnet_group_size))
        self.resnets = nn.ModuleList(resnets)
        self.attentions = nn.ModuleList(attentions)
        if add_downsample: self.downsamplers = nn.ModuleList([KDownsample2D()])
        else: self.downsamplers = None
        self.gradient_checkpointing = False
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None,
    cross_attention_kwargs: Optional[Dict[str, Any]]=None, encoder_attention_mask: Optional[torch.Tensor]=None) -> Tuple[torch.Tensor, Tuple[torch.Tensor, ...]]:
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
        output_states = ()
        for resnet, attn in zip(self.resnets, self.attentions):
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None: return module(*inputs, return_dict=return_dict)
                        else: return module(*inputs)
                    return custom_forward
                ckpt_kwargs: Dict[str, Any] = {'use_reentrant': False} if is_torch_version('>=', '1.11.0') else {}
                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, emb=temb, attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs, encoder_attention_mask=encoder_attention_mask)
            else:
                hidden_states = resnet(hidden_states, temb)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, emb=temb, attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs, encoder_attention_mask=encoder_attention_mask)
            if self.downsamplers is None: output_states += (None,)
            else: output_states += (hidden_states,)
        if self.downsamplers is not None:
            for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states)
        return (hidden_states, output_states)
class AttnUpBlock2D(nn.Module):
    def __init__(self, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, resolution_idx: int=None, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06,
    resnet_time_scale_shift: str='default', resnet_act_fn: str='swish', resnet_groups: int=32, resnet_pre_norm: bool=True, attention_head_dim: int=1, output_scale_factor: float=1.0, upsample_type: str='conv'):
        super().__init__()
        resnets = []
        attentions = []
        self.upsample_type = upsample_type
        if attention_head_dim is None: attention_head_dim = out_channels
        for i in range(num_layers):
            res_skip_channels = in_channels if i == num_layers - 1 else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels
            resnets.append(ResnetBlock2D(in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
            attentions.append(Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps,
            norm_num_groups=resnet_groups, residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
        if upsample_type == 'conv': self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
        elif upsample_type == 'resnet': self.upsamplers = nn.ModuleList([ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups,
        dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, up=True)])
        else: self.upsamplers = None
        self.gradient_checkpointing = False
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, res_hidden_states_tuple: Tuple[torch.Tensor, ...], temb: Optional[torch.Tensor]=None, upsample_size: Optional[int]=None, *args, **kwargs) -> torch.Tensor:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        for resnet, attn in zip(self.resnets, self.attentions):
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None: return module(*inputs, return_dict=return_dict)
                        else: return module(*inputs)
                    return custom_forward
                ckpt_kwargs: Dict[str, Any] = {'use_reentrant': False} if is_torch_version('>=', '1.11.0') else {}
                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs)
                hidden_states = attn(hidden_states)
            else:
                hidden_states = resnet(hidden_states, temb)
                hidden_states = attn(hidden_states)
        if self.upsamplers is not None:
            for upsampler in self.upsamplers:
                if self.upsample_type == 'resnet': hidden_states = upsampler(hidden_states, temb=temb)
                else: hidden_states = upsampler(hidden_states)
        return hidden_states
class CrossAttnUpBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, prev_output_channel: int, temb_channels: int, resolution_idx: Optional[int]=None, dropout: float=0.0, num_layers: int=1, transformer_layers_per_block: Union[int, Tuple[int]]=1, resnet_eps: float=1e-06,
    resnet_time_scale_shift: str='default', resnet_act_fn: str='swish', resnet_groups: int=32, resnet_pre_norm: bool=True, num_attention_heads: int=1, cross_attention_dim: int=1280, output_scale_factor: float=1.0, add_upsample: bool=True,
    dual_cross_attention: bool=False, use_linear_projection: bool=False, only_cross_attention: bool=False, upcast_attention: bool=False, attention_type: str='default'):
        super().__init__()
        resnets = []
        attentions = []
        self.has_cross_attention = True
        self.num_attention_heads = num_attention_heads
        if isinstance(transformer_layers_per_block, int): transformer_layers_per_block = [transformer_layers_per_block] * num_layers
        for i in range(num_layers):
            res_skip_channels = in_channels if i == num_layers - 1 else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels
            resnets.append(ResnetBlock2D(in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift,
            non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
            if not dual_cross_attention: attentions.append(Transformer2DModel(num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=transformer_layers_per_block[i], cross_attention_dim=cross_attention_dim,
            norm_num_groups=resnet_groups, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, attention_type=attention_type))
            else: attentions.append(DualTransformer2DModel(num_attention_heads, out_channels // num_attention_heads, in_channels=out_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
        if add_upsample: self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
        else: self.upsamplers = None
        self.gradient_checkpointing = False
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, res_hidden_states_tuple: Tuple[torch.Tensor, ...], temb: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.Tensor]=None, cross_attention_kwargs: Optional[Dict[str, Any]]=None,
    upsample_size: Optional[int]=None, attention_mask: Optional[torch.Tensor]=None, encoder_attention_mask: Optional[torch.Tensor]=None) -> torch.Tensor:
        is_freeu_enabled = getattr(self, 's1', None) and getattr(self, 's2', None) and getattr(self, 'b1', None) and getattr(self, 'b2', None)
        for resnet, attn in zip(self.resnets, self.attentions):
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            if is_freeu_enabled: hidden_states, res_hidden_states = apply_freeu(self.resolution_idx, hidden_states, res_hidden_states, s1=self.s1, s2=self.s2, b1=self.b1, b2=self.b2)
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None: return module(*inputs, return_dict=return_dict)
                        else: return module(*inputs)
                    return custom_forward
                ckpt_kwargs: Dict[str, Any] = {'use_reentrant': False} if is_torch_version('>=', '1.11.0') else {}
                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False)[0]
            else:
                hidden_states = resnet(hidden_states, temb)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, attention_mask=attention_mask, encoder_attention_mask=encoder_attention_mask, return_dict=False)[0]
        if self.upsamplers is not None:
            for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states, upsample_size)
        return hidden_states
class UpBlock2D(nn.Module):
    def __init__(self, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, resolution_idx: Optional[int]=None, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default',
    resnet_act_fn: str='swish', resnet_groups: int=32, resnet_pre_norm: bool=True, output_scale_factor: float=1.0, add_upsample: bool=True):
        super().__init__()
        resnets = []
        for i in range(num_layers):
            res_skip_channels = in_channels if i == num_layers - 1 else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels
            resnets.append(ResnetBlock2D(in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift,
            non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
        self.resnets = nn.ModuleList(resnets)
        if add_upsample: self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
        else: self.upsamplers = None
        self.gradient_checkpointing = False
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, res_hidden_states_tuple: Tuple[torch.Tensor, ...], temb: Optional[torch.Tensor]=None, upsample_size: Optional[int]=None, *args, **kwargs) -> torch.Tensor:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        is_freeu_enabled = getattr(self, 's1', None) and getattr(self, 's2', None) and getattr(self, 'b1', None) and getattr(self, 'b2', None)
        for resnet in self.resnets:
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            if is_freeu_enabled: hidden_states, res_hidden_states = apply_freeu(self.resolution_idx, hidden_states, res_hidden_states, s1=self.s1, s2=self.s2, b1=self.b1, b2=self.b2)
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module):
                    def custom_forward(*inputs): return module(*inputs)
                    return custom_forward
                if is_torch_version('>=', '1.11.0'): hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, use_reentrant=False)
                else: hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
            else: hidden_states = resnet(hidden_states, temb)
        if self.upsamplers is not None:
            for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states, upsample_size)
        return hidden_states
class UpDecoderBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, resolution_idx: Optional[int]=None, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish', resnet_groups: int=32,
    resnet_pre_norm: bool=True, output_scale_factor: float=1.0, add_upsample: bool=True, temb_channels: Optional[int]=None):
        super().__init__()
        resnets = []
        for i in range(num_layers):
            input_channels = in_channels if i == 0 else out_channels
            if resnet_time_scale_shift == 'spatial': resnets.append(ResnetBlockCondNorm2D(in_channels=input_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm='spatial', non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor))
            else: resnets.append(ResnetBlock2D(in_channels=input_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift,
            non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
        self.resnets = nn.ModuleList(resnets)
        if add_upsample: self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
        else: self.upsamplers = None
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None) -> torch.Tensor:
        for resnet in self.resnets: hidden_states = resnet(hidden_states, temb=temb)
        if self.upsamplers is not None:
            for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states)
        return hidden_states
class AttnUpDecoderBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, resolution_idx: Optional[int]=None, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06, resnet_time_scale_shift: str='default', resnet_act_fn: str='swish',
    resnet_groups: int=32, resnet_pre_norm: bool=True, attention_head_dim: int=1, output_scale_factor: float=1.0, add_upsample: bool=True, temb_channels: Optional[int]=None):
        super().__init__()
        resnets = []
        attentions = []
        if attention_head_dim is None: attention_head_dim = out_channels
        for i in range(num_layers):
            input_channels = in_channels if i == 0 else out_channels
            if resnet_time_scale_shift == 'spatial': resnets.append(ResnetBlockCondNorm2D(in_channels=input_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm='spatial', non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor))
            else: resnets.append(ResnetBlock2D(in_channels=input_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
            attentions.append(Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps,
            norm_num_groups=resnet_groups if resnet_time_scale_shift != 'spatial' else None, spatial_norm_dim=temb_channels if resnet_time_scale_shift == 'spatial' else None,
            residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
        if add_upsample: self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
        else: self.upsamplers = None
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, temb: Optional[torch.Tensor]=None) -> torch.Tensor:
        for resnet, attn in zip(self.resnets, self.attentions):
            hidden_states = resnet(hidden_states, temb=temb)
            hidden_states = attn(hidden_states, temb=temb)
        if self.upsamplers is not None:
            for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states)
        return hidden_states
class AttnSkipUpBlock2D(nn.Module):
    def __init__(self, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, resolution_idx: Optional[int]=None, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06,
    resnet_time_scale_shift: str='default', resnet_act_fn: str='swish', resnet_pre_norm: bool=True, attention_head_dim: int=1, output_scale_factor: float=np.sqrt(2.0), add_upsample: bool=True):
        super().__init__()
        self.attentions = nn.ModuleList([])
        self.resnets = nn.ModuleList([])
        for i in range(num_layers):
            res_skip_channels = in_channels if i == num_layers - 1 else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels
            self.resnets.append(ResnetBlock2D(in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=min(resnet_in_channels + res_skip_channels // 4, 32),
            groups_out=min(out_channels // 4, 32), dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
        if attention_head_dim is None: attention_head_dim = out_channels
        self.attentions.append(Attention(out_channels, heads=out_channels // attention_head_dim, dim_head=attention_head_dim, rescale_output_factor=output_scale_factor, eps=resnet_eps, norm_num_groups=32,
        residual_connection=True, bias=True, upcast_softmax=True, _from_deprecated_attn_block=True))
        self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels)
        if add_upsample:
            self.resnet_up = ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=min(out_channels // 4, 32),
            groups_out=min(out_channels // 4, 32), dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_in_shortcut=True, up=True, kernel='fir')
            self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            self.skip_norm = torch.nn.GroupNorm(num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True)
            self.act = nn.SiLU()
        else:
            self.resnet_up = None
            self.skip_conv = None
            self.skip_norm = None
            self.act = None
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, res_hidden_states_tuple: Tuple[torch.Tensor, ...], temb: Optional[torch.Tensor]=None, skip_sample=None, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        for resnet in self.resnets:
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
            hidden_states = resnet(hidden_states, temb)
        hidden_states = self.attentions[0](hidden_states)
        if skip_sample is not None: skip_sample = self.upsampler(skip_sample)
        else: skip_sample = 0
        if self.resnet_up is not None:
            skip_sample_states = self.skip_norm(hidden_states)
            skip_sample_states = self.act(skip_sample_states)
            skip_sample_states = self.skip_conv(skip_sample_states)
            skip_sample = skip_sample + skip_sample_states
            hidden_states = self.resnet_up(hidden_states, temb)
        return (hidden_states, skip_sample)
class SkipUpBlock2D(nn.Module):
    def __init__(self, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, resolution_idx: Optional[int]=None, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06,
    resnet_time_scale_shift: str='default', resnet_act_fn: str='swish', resnet_pre_norm: bool=True, output_scale_factor: float=np.sqrt(2.0), add_upsample: bool=True, upsample_padding: int=1):
        super().__init__()
        self.resnets = nn.ModuleList([])
        for i in range(num_layers):
            res_skip_channels = in_channels if i == num_layers - 1 else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels
            self.resnets.append(ResnetBlock2D(in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=min((resnet_in_channels + res_skip_channels) // 4, 32),
            groups_out=min(out_channels // 4, 32), dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm))
        self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels)
        if add_upsample:
            self.resnet_up = ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=min(out_channels // 4, 32), groups_out=min(out_channels // 4, 32), dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, use_in_shortcut=True, up=True, kernel='fir')
            self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
            self.skip_norm = torch.nn.GroupNorm(num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True)
            self.act = nn.SiLU()
        else:
            self.resnet_up = None
            self.skip_conv = None
            self.skip_norm = None
            self.act = None
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, res_hidden_states_tuple: Tuple[torch.Tensor, ...], temb: Optional[torch.Tensor]=None, skip_sample=None, *args, **kwargs) -> Tuple[torch.Tensor, torch.Tensor]:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        for resnet in self.resnets:
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
            hidden_states = resnet(hidden_states, temb)
        if skip_sample is not None: skip_sample = self.upsampler(skip_sample)
        else: skip_sample = 0
        if self.resnet_up is not None:
            skip_sample_states = self.skip_norm(hidden_states)
            skip_sample_states = self.act(skip_sample_states)
            skip_sample_states = self.skip_conv(skip_sample_states)
            skip_sample = skip_sample + skip_sample_states
            hidden_states = self.resnet_up(hidden_states, temb)
        return (hidden_states, skip_sample)
class ResnetUpsampleBlock2D(nn.Module):
    def __init__(self, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, resolution_idx: Optional[int]=None, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06,
    resnet_time_scale_shift: str='default', resnet_act_fn: str='swish', resnet_groups: int=32, resnet_pre_norm: bool=True, output_scale_factor: float=1.0, add_upsample: bool=True, skip_time_act: bool=False):
        super().__init__()
        resnets = []
        for i in range(num_layers):
            res_skip_channels = in_channels if i == num_layers - 1 else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels
            resnets.append(ResnetBlock2D(in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=skip_time_act))
        self.resnets = nn.ModuleList(resnets)
        if add_upsample: self.upsamplers = nn.ModuleList([ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
        time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=skip_time_act, up=True)])
        else: self.upsamplers = None
        self.gradient_checkpointing = False
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, res_hidden_states_tuple: Tuple[torch.Tensor, ...], temb: Optional[torch.Tensor]=None, upsample_size: Optional[int]=None, *args, **kwargs) -> torch.Tensor:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        for resnet in self.resnets:
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module):
                    def custom_forward(*inputs): return module(*inputs)
                    return custom_forward
                if is_torch_version('>=', '1.11.0'): hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, use_reentrant=False)
                else: hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
            else: hidden_states = resnet(hidden_states, temb)
        if self.upsamplers is not None:
            for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states, temb)
        return hidden_states
class SimpleCrossAttnUpBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, prev_output_channel: int, temb_channels: int, resolution_idx: Optional[int]=None, dropout: float=0.0, num_layers: int=1, resnet_eps: float=1e-06,
    resnet_time_scale_shift: str='default', resnet_act_fn: str='swish', resnet_groups: int=32, resnet_pre_norm: bool=True, attention_head_dim: int=1, cross_attention_dim: int=1280, output_scale_factor: float=1.0,
    add_upsample: bool=True, skip_time_act: bool=False, only_cross_attention: bool=False, cross_attention_norm: Optional[str]=None):
        super().__init__()
        resnets = []
        attentions = []
        self.has_cross_attention = True
        self.attention_head_dim = attention_head_dim
        self.num_heads = out_channels // self.attention_head_dim
        for i in range(num_layers):
            res_skip_channels = in_channels if i == num_layers - 1 else out_channels
            resnet_in_channels = prev_output_channel if i == 0 else out_channels
            resnets.append(ResnetBlock2D(in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
            time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=skip_time_act))
            processor = AttnAddedKVProcessor2_0() if hasattr(F, 'scaled_dot_product_attention') else AttnAddedKVProcessor()
            attentions.append(Attention(query_dim=out_channels, cross_attention_dim=out_channels, heads=self.num_heads, dim_head=self.attention_head_dim, added_kv_proj_dim=cross_attention_dim, norm_num_groups=resnet_groups,
            bias=True, upcast_softmax=True, only_cross_attention=only_cross_attention, cross_attention_norm=cross_attention_norm, processor=processor))
        self.attentions = nn.ModuleList(attentions)
        self.resnets = nn.ModuleList(resnets)
        if add_upsample: self.upsamplers = nn.ModuleList([ResnetBlock2D(in_channels=out_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout,
        time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=skip_time_act, up=True)])
        else: self.upsamplers = None
        self.gradient_checkpointing = False
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, res_hidden_states_tuple: Tuple[torch.Tensor, ...], temb: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.Tensor]=None, upsample_size: Optional[int]=None,
    attention_mask: Optional[torch.Tensor]=None, cross_attention_kwargs: Optional[Dict[str, Any]]=None, encoder_attention_mask: Optional[torch.Tensor]=None) -> torch.Tensor:
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
        if attention_mask is None: mask = None if encoder_hidden_states is None else encoder_attention_mask
        else: mask = attention_mask
        for resnet, attn in zip(self.resnets, self.attentions):
            res_hidden_states = res_hidden_states_tuple[-1]
            res_hidden_states_tuple = res_hidden_states_tuple[:-1]
            hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None: return module(*inputs, return_dict=return_dict)
                        else: return module(*inputs)
                    return custom_forward
                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=mask, **cross_attention_kwargs)
            else:
                hidden_states = resnet(hidden_states, temb)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=mask, **cross_attention_kwargs)
        if self.upsamplers is not None:
            for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states, temb)
        return hidden_states
class KUpBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, resolution_idx: int, dropout: float=0.0, num_layers: int=5, resnet_eps: float=1e-05,
    resnet_act_fn: str='gelu', resnet_group_size: Optional[int]=32, add_upsample: bool=True):
        super().__init__()
        resnets = []
        k_in_channels = 2 * out_channels
        k_out_channels = in_channels
        num_layers = num_layers - 1
        for i in range(num_layers):
            in_channels = k_in_channels if i == 0 else out_channels
            groups = in_channels // resnet_group_size
            groups_out = out_channels // resnet_group_size
            resnets.append(ResnetBlockCondNorm2D(in_channels=in_channels, out_channels=k_out_channels if i == num_layers - 1 else out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=groups,
            groups_out=groups_out, dropout=dropout, non_linearity=resnet_act_fn, time_embedding_norm='ada_group', conv_shortcut_bias=False))
        self.resnets = nn.ModuleList(resnets)
        if add_upsample: self.upsamplers = nn.ModuleList([KUpsample2D()])
        else: self.upsamplers = None
        self.gradient_checkpointing = False
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, res_hidden_states_tuple: Tuple[torch.Tensor, ...], temb: Optional[torch.Tensor]=None, upsample_size: Optional[int]=None, *args, **kwargs) -> torch.Tensor:
        if len(args) > 0 or kwargs.get('scale', None) is not None:
            deprecation_message = 'The `scale` argument is deprecated and will be ignored. Please remove it, as passing it will raise an error in the future. `scale` should directly be passed while calling the underlying pipeline component i.e., via `cross_attention_kwargs`.'
            deprecate('scale', '1.0.0', deprecation_message)
        res_hidden_states_tuple = res_hidden_states_tuple[-1]
        if res_hidden_states_tuple is not None: hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1)
        for resnet in self.resnets:
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module):
                    def custom_forward(*inputs): return module(*inputs)
                    return custom_forward
                if is_torch_version('>=', '1.11.0'): hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, use_reentrant=False)
                else: hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
            else: hidden_states = resnet(hidden_states, temb)
        if self.upsamplers is not None:
            for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states)
        return hidden_states
class KCrossAttnUpBlock2D(nn.Module):
    def __init__(self, in_channels: int, out_channels: int, temb_channels: int, resolution_idx: int, dropout: float=0.0, num_layers: int=4, resnet_eps: float=1e-05, resnet_act_fn: str='gelu', resnet_group_size: int=32,
    attention_head_dim: int=1, cross_attention_dim: int=768, add_upsample: bool=True, upcast_attention: bool=False):
        super().__init__()
        resnets = []
        attentions = []
        is_first_block = in_channels == out_channels == temb_channels
        is_middle_block = in_channels != out_channels
        add_self_attention = True if is_first_block else False
        self.has_cross_attention = True
        self.attention_head_dim = attention_head_dim
        k_in_channels = out_channels if is_first_block else 2 * out_channels
        k_out_channels = in_channels
        num_layers = num_layers - 1
        for i in range(num_layers):
            in_channels = k_in_channels if i == 0 else out_channels
            groups = in_channels // resnet_group_size
            groups_out = out_channels // resnet_group_size
            if is_middle_block and i == num_layers - 1: conv_2d_out_channels = k_out_channels
            else: conv_2d_out_channels = None
            resnets.append(ResnetBlockCondNorm2D(in_channels=in_channels, out_channels=out_channels, conv_2d_out_channels=conv_2d_out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=groups,
            groups_out=groups_out, dropout=dropout, non_linearity=resnet_act_fn, time_embedding_norm='ada_group', conv_shortcut_bias=False))
            attentions.append(KAttentionBlock(k_out_channels if i == num_layers - 1 else out_channels, k_out_channels // attention_head_dim if i == num_layers - 1 else out_channels // attention_head_dim,
            attention_head_dim, cross_attention_dim=cross_attention_dim, temb_channels=temb_channels, attention_bias=True, add_self_attention=add_self_attention, cross_attention_norm='layer_norm', upcast_attention=upcast_attention))
        self.resnets = nn.ModuleList(resnets)
        self.attentions = nn.ModuleList(attentions)
        if add_upsample: self.upsamplers = nn.ModuleList([KUpsample2D()])
        else: self.upsamplers = None
        self.gradient_checkpointing = False
        self.resolution_idx = resolution_idx
    def forward(self, hidden_states: torch.Tensor, res_hidden_states_tuple: Tuple[torch.Tensor, ...], temb: Optional[torch.Tensor]=None, encoder_hidden_states: Optional[torch.Tensor]=None,
    cross_attention_kwargs: Optional[Dict[str, Any]]=None, upsample_size: Optional[int]=None, attention_mask: Optional[torch.Tensor]=None, encoder_attention_mask: Optional[torch.Tensor]=None) -> torch.Tensor:
        res_hidden_states_tuple = res_hidden_states_tuple[-1]
        if res_hidden_states_tuple is not None: hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1)
        for resnet, attn in zip(self.resnets, self.attentions):
            if torch.is_grad_enabled() and self.gradient_checkpointing:
                def create_custom_forward(module, return_dict=None):
                    def custom_forward(*inputs):
                        if return_dict is not None: return module(*inputs, return_dict=return_dict)
                        else: return module(*inputs)
                    return custom_forward
                ckpt_kwargs: Dict[str, Any] = {'use_reentrant': False} if is_torch_version('>=', '1.11.0') else {}
                hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, emb=temb, attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs, encoder_attention_mask=encoder_attention_mask)
            else:
                hidden_states = resnet(hidden_states, temb)
                hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states, emb=temb, attention_mask=attention_mask, cross_attention_kwargs=cross_attention_kwargs, encoder_attention_mask=encoder_attention_mask)
        if self.upsamplers is not None:
            for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states)
        return hidden_states
class KAttentionBlock(nn.Module):
    def __init__(self, dim: int, num_attention_heads: int, attention_head_dim: int, dropout: float=0.0, cross_attention_dim: Optional[int]=None, attention_bias: bool=False, upcast_attention: bool=False, temb_channels: int=768,
    add_self_attention: bool=False, cross_attention_norm: Optional[str]=None, group_size: int=32):
        super().__init__()
        self.add_self_attention = add_self_attention
        if add_self_attention:
            self.norm1 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size))
            self.attn1 = Attention(query_dim=dim, heads=num_attention_heads, dim_head=attention_head_dim, dropout=dropout, bias=attention_bias, cross_attention_dim=None, cross_attention_norm=None)
        self.norm2 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size))
        self.attn2 = Attention(query_dim=dim, cross_attention_dim=cross_attention_dim, heads=num_attention_heads, dim_head=attention_head_dim, dropout=dropout,
        bias=attention_bias, upcast_attention=upcast_attention, cross_attention_norm=cross_attention_norm)
    def _to_3d(self, hidden_states: torch.Tensor, height: int, weight: int) -> torch.Tensor: return hidden_states.permute(0, 2, 3, 1).reshape(hidden_states.shape[0], height * weight, -1)
    def _to_4d(self, hidden_states: torch.Tensor, height: int, weight: int) -> torch.Tensor: return hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, height, weight)
    def forward(self, hidden_states: torch.Tensor, encoder_hidden_states: Optional[torch.Tensor]=None, emb: Optional[torch.Tensor]=None, attention_mask: Optional[torch.Tensor]=None,
    cross_attention_kwargs: Optional[Dict[str, Any]]=None, encoder_attention_mask: Optional[torch.Tensor]=None) -> torch.Tensor:
        cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}
        if self.add_self_attention:
            norm_hidden_states = self.norm1(hidden_states, emb)
            height, weight = norm_hidden_states.shape[2:]
            norm_hidden_states = self._to_3d(norm_hidden_states, height, weight)
            attn_output = self.attn1(norm_hidden_states, encoder_hidden_states=None, attention_mask=attention_mask, **cross_attention_kwargs)
            attn_output = self._to_4d(attn_output, height, weight)
            hidden_states = attn_output + hidden_states
        norm_hidden_states = self.norm2(hidden_states, emb)
        height, weight = norm_hidden_states.shape[2:]
        norm_hidden_states = self._to_3d(norm_hidden_states, height, weight)
        attn_output = self.attn2(norm_hidden_states, encoder_hidden_states=encoder_hidden_states, attention_mask=attention_mask if encoder_hidden_states is None else encoder_attention_mask, **cross_attention_kwargs)
        attn_output = self._to_4d(attn_output, height, weight)
        hidden_states = attn_output + hidden_states
        return hidden_states
'''
    #################################################################################################################################################################################################
    # This code library is an adaptation of the original Transformers and was designed, developed and programmed by Sapiens Technology®.                                                            #
    # Any alteration and/or disclosure of this code without prior authorization is strictly prohibited and is subject to legal action that will be forwarded by the Sapiens Technology® legal team. #
    # This set of algorithms aims to download, train, fine-tune and/or infer large language models from various sources and slopes.                                                                 #
    #################################################################################################################################################################################################
'''
