content: # Copyright 2023 The HuggingFace Team. All rights reserved.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  from dataclasses import dataclass
  from typing import Any, Dict, List, Optional, Tuple, Union
  import os
  
  import torch
  import torch.nn as nn
  import torch.utils.checkpoint
  
  from diffusers.configuration_utils import ConfigMixin, register_to_config
  from diffusers.loaders import UNet2DConditionLoadersMixin
  from diffusers.utils import BaseOutput, logging
  from diffusers.models.activations import get_activation
  from diffusers.models.attention_processor import AttentionProcessor, AttnProcessor
  from diffusers.models.embeddings import (
      GaussianFourierProjection,
      ImageHintTimeEmbedding,
      ImageProjection,
      ImageTimeEmbedding,
      TextImageProjection,
      TextImageTimeEmbedding,
      TextTimeEmbedding,
      TimestepEmbedding,
      Timesteps,
  )
  from diffusers.models.modeling_utils import ModelMixin
  from diffusers.models.model_loading_utils import load_state_dict, _load_state_dict_into_model
  from diffusers.models.unets.unet_2d_blocks import (
      CrossAttnDownBlock2D,
      CrossAttnUpBlock2D,
      DownBlock2D,
      UNetMidBlock2DCrossAttn,
      UNetMidBlock2DSimpleCrossAttn,
      UpBlock2D,
  )
  from huggingface_hub.constants import HF_HUB_OFFLINE, HUGGINGFACE_HUB_CACHE
  from diffusers.utils import (
      CONFIG_NAME,
      FLAX_WEIGHTS_NAME,
      SAFETENSORS_WEIGHTS_NAME,
      WEIGHTS_NAME,
      _add_variant,
      _get_model_file,
      deprecate,
      is_accelerate_available,
      is_torch_version,
      logging,
  )
  from diffusers import __version__
  from .unet_mv2d_blocks import (
      CrossAttnDownBlockMV2D,
      CrossAttnUpBlockMV2D,
      UNetMidBlockMV2DCrossAttn,
      get_down_block,
      get_up_block,
  )
  
  
  logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  @dataclass
  class UNetMV2DConditionOutput(BaseOutput):
      """
      The output of [`UNet2DConditionModel`].
  
      Args:
          sample (`torch.FloatTensor` of shape `(batch_size, num_channels, height, width)`):
              The hidden states output conditioned on `encoder_hidden_states` input. Output of last layer of model.
      """
  
      sample: torch.FloatTensor = None
  
  
  class UNetMV2DConditionModel(ModelMixin, ConfigMixin, UNet2DConditionLoadersMixin):
      r"""
      A conditional 2D UNet model that takes a noisy sample, conditional state, and a timestep and returns a sample
      shaped output.
  
      This model inherits from [`ModelMixin`]. Check the superclass documentation for it's generic methods implemented
      for all models (such as downloading or saving).
  
      Parameters:
          sample_size (`int` or `Tuple[int, int]`, *optional*, defaults to `None`):
              Height and width of input/output sample.
          in_channels (`int`, *optional*, defaults to 4): Number of channels in the input sample.
          out_channels (`int`, *optional*, defaults to 4): Number of channels in the output.
          center_input_sample (`bool`, *optional*, defaults to `False`): Whether to center the input sample.
          flip_sin_to_cos (`bool`, *optional*, defaults to `False`):
              Whether to flip the sin to cos in the time embedding.
          freq_shift (`int`, *optional*, defaults to 0): The frequency shift to apply to the time embedding.
          down_block_types (`Tuple[str]`, *optional*, defaults to `("CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "CrossAttnDownBlock2D", "DownBlock2D")`):
              The tuple of downsample blocks to use.
          mid_block_type (`str`, *optional*, defaults to `"UNetMidBlock2DCrossAttn"`):
              Block type for middle of UNet, it can be either `UNetMidBlock2DCrossAttn` or
              `UNetMidBlock2DSimpleCrossAttn`. If `None`, the mid block layer is skipped.
          up_block_types (`Tuple[str]`, *optional*, defaults to `("UpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D", "CrossAttnUpBlock2D")`):
              The tuple of upsample blocks to use.
          only_cross_attention(`bool` or `Tuple[bool]`, *optional*, default to `False`):
              Whether to include self-attention in the basic transformer blocks, see
              [`~models.attention.BasicTransformerBlock`].
          block_out_channels (`Tuple[int]`, *optional*, defaults to `(320, 640, 1280, 1280)`):
              The tuple of output channels for each block.
          layers_per_block (`int`, *optional*, defaults to 2): The number of layers per block.
          downsample_padding (`int`, *optional*, defaults to 1): The padding to use for the downsampling convolution.
          mid_block_scale_factor (`float`, *optional*, defaults to 1.0): The scale factor to use for the mid block.
          act_fn (`str`, *optional*, defaults to `"silu"`): The activation function to use.
          norm_num_groups (`int`, *optional*, defaults to 32): The number of groups to use for the normalization.
              If `None`, normalization and activation layers is skipped in post-processing.
          norm_eps (`float`, *optional*, defaults to 1e-5): The epsilon to use for the normalization.
          cross_attention_dim (`int` or `Tuple[int]`, *optional*, defaults to 1280):
              The dimension of the cross attention features.
          transformer_layers_per_block (`int` or `Tuple[int]`, *optional*, defaults to 1):
              The number of transformer blocks of type [`~models.attention.BasicTransformerBlock`]. Only relevant for
              [`~models.unet_2d_blocks.CrossAttnDownBlock2D`], [`~models.unet_2d_blocks.CrossAttnUpBlock2D`],
              [`~models.unet_2d_blocks.UNetMidBlock2DCrossAttn`].
          encoder_hid_dim (`int`, *optional*, defaults to None):
              If `encoder_hid_dim_type` is defined, `encoder_hidden_states` will be projected from `encoder_hid_dim`
              dimension to `cross_attention_dim`.
          encoder_hid_dim_type (`str`, *optional*, defaults to `None`):
              If given, the `encoder_hidden_states` and potentially other embeddings are down-projected to text
              embeddings of dimension `cross_attention` according to `encoder_hid_dim_type`.
          attention_head_dim (`int`, *optional*, defaults to 8): The dimension of the attention heads.
          num_attention_heads (`int`, *optional*):
              The number of attention heads. If not defined, defaults to `attention_head_dim`
          resnet_time_scale_shift (`str`, *optional*, defaults to `"default"`): Time scale shift config
              for ResNet blocks (see [`~models.resnet.ResnetBlock2D`]). Choose from `default` or `scale_shift`.
          class_embed_type (`str`, *optional*, defaults to `None`):
              The type of class embedding to use which is ultimately summed with the time embeddings. Choose from `None`,
              `"timestep"`, `"identity"`, `"projection"`, or `"simple_projection"`.
          addition_embed_type (`str`, *optional*, defaults to `None`):
              Configures an optional embedding which will be summed with the time embeddings. Choose from `None` or
              "text". "text" will use the `TextTimeEmbedding` layer.
          addition_time_embed_dim: (`int`, *optional*, defaults to `None`):
              Dimension for the timestep embeddings.
          num_class_embeds (`int`, *optional*, defaults to `None`):
              Input dimension of the learnable embedding matrix to be projected to `time_embed_dim`, when performing
              class conditioning with `class_embed_type` equal to `None`.
          time_embedding_type (`str`, *optional*, defaults to `positional`):
              The type of position embedding to use for timesteps. Choose from `positional` or `fourier`.
          time_embedding_dim (`int`, *optional*, defaults to `None`):
              An optional override for the dimension of the projected time embedding.
          time_embedding_act_fn (`str`, *optional*, defaults to `None`):
              Optional activation function to use only once on the time embeddings before they are passed to the rest of
              the UNet. Choose from `silu`, `mish`, `gelu`, and `swish`.
          timestep_post_act (`str`, *optional*, defaults to `None`):
              The second activation function to use in timestep embedding. Choose from `silu`, `mish` and `gelu`.
          time_cond_proj_dim (`int`, *optional*, defaults to `None`):
              The dimension of `cond_proj` layer in the timestep embedding.
          conv_in_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_in` layer.
          conv_out_kernel (`int`, *optional*, default to `3`): The kernel size of `conv_out` layer.
          projection_class_embeddings_input_dim (`int`, *optional*): The dimension of the `class_labels` input when
              `class_embed_type="projection"`. Required when `class_embed_type="projection"`.
          class_embeddings_concat (`bool`, *optional*, defaults to `False`): Whether to concatenate the time
              embeddings with the class embeddings.
          mid_block_only_cross_attention (`bool`, *optional*, defaults to `None`):
              Whether to use cross attention with the mid block when using the `UNetMidBlock2DSimpleCrossAttn`. If
              `only_cross_attention` is given as a single boolean and `mid_block_only_cross_attention` is `None`, the
              `only_cross_attention` value is used as the value for `mid_block_only_cross_attention`. Default to `False`
              otherwise.
      """
  
      _supports_gradient_checkpointing = True
  
      @register_to_config
      def __init__(
          self,
          sample_size: Optional[int] = None,
          in_channels: int = 4,
          out_channels: int = 4,
          center_input_sample: bool = False,
          flip_sin_to_cos: bool = True,
          freq_shift: int = 0,
          down_block_types: Tuple[str] = (
              "CrossAttnDownBlockMV2D",
              "CrossAttnDownBlockMV2D",
              "CrossAttnDownBlockMV2D",
              "DownBlock2D",
          ),
          mid_block_type: Optional[str] = "UNetMidBlockMV2DCrossAttn",
          up_block_types: Tuple[str] = ("UpBlock2D", "CrossAttnUpBlockMV2D", "CrossAttnUpBlockMV2D", "CrossAttnUpBlockMV2D"),
          only_cross_attention: Union[bool, Tuple[bool]] = False,
          block_out_channels: Tuple[int] = (320, 640, 1280, 1280),
          layers_per_block: Union[int, Tuple[int]] = 2,
          downsample_padding: int = 1,
          mid_block_scale_factor: float = 1,
          act_fn: str = "silu",
          norm_num_groups: Optional[int] = 32,
          norm_eps: float = 1e-5,
          cross_attention_dim: Union[int, Tuple[int]] = 1280,
          transformer_layers_per_block: Union[int, Tuple[int]] = 1,
          encoder_hid_dim: Optional[int] = None,
          encoder_hid_dim_type: Optional[str] = None,
          attention_head_dim: Union[int, Tuple[int]] = 8,
          num_attention_heads: Optional[Union[int, Tuple[int]]] = None,
          dual_cross_attention: bool = False,
          use_linear_projection: bool = False,
          class_embed_type: Optional[str] = None,
          addition_embed_type: Optional[str] = None,
          addition_time_embed_dim: Optional[int] = None,
          num_class_embeds: Optional[int] = None,
          upcast_attention: bool = False,
          resnet_time_scale_shift: str = "default",
          resnet_skip_time_act: bool = False,
          resnet_out_scale_factor: int = 1.0,
          time_embedding_type: str = "positional",
          time_embedding_dim: Optional[int] = None,
          time_embedding_act_fn: Optional[str] = None,
          timestep_post_act: Optional[str] = None,
          time_cond_proj_dim: Optional[int] = None,
          conv_in_kernel: int = 3,
          conv_out_kernel: int = 3,
          projection_class_embeddings_input_dim: Optional[int] = None,
          class_embeddings_concat: bool = False,
          mid_block_only_cross_attention: Optional[bool] = None,
          cross_attention_norm: Optional[str] = None,
          addition_embed_type_num_heads=64,
          num_views: int = 1,
          cd_attention_last: bool = False,
          cd_attention_mid: bool = False,
          multiview_attention: bool = True,
          sparse_mv_attention: bool = False,
          mvcd_attention: bool = False
      ):
          super().__init__()
  
          self.sample_size = sample_size
  
          if num_attention_heads is not None:
              raise ValueError(
                  "At the moment it is not possible to define the number of attention heads via `num_attention_heads` because of a naming issue as described in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131. Passing `num_attention_heads` will only be supported in diffusers v0.19."
              )
  
          # If `num_attention_heads` is not defined (which is the case for most models)
          # it will default to `attention_head_dim`. This looks weird upon first reading it and it is.
          # The reason for this behavior is to correct for incorrectly named variables that were introduced
          # when this library was created. The incorrect naming was only discovered much later in https://github.com/huggingface/diffusers/issues/2011#issuecomment-1547958131
          # Changing `attention_head_dim` to `num_attention_heads` for 40,000+ configurations is too backwards breaking
          # which is why we correct for the naming here.
          num_attention_heads = num_attention_heads or attention_head_dim
  
          # Check inputs
          if len(down_block_types) != len(up_block_types):
              raise ValueError(
                  f"Must provide the same number of `down_block_types` as `up_block_types`. `down_block_types`: {down_block_types}. `up_block_types`: {up_block_types}."
              )
  
          if len(block_out_channels) != len(down_block_types):
              raise ValueError(
                  f"Must provide the same number of `block_out_channels` as `down_block_types`. `block_out_channels`: {block_out_channels}. `down_block_types`: {down_block_types}."
              )
  
          if not isinstance(only_cross_attention, bool) and len(only_cross_attention) != len(down_block_types):
              raise ValueError(
                  f"Must provide the same number of `only_cross_attention` as `down_block_types`. `only_cross_attention`: {only_cross_attention}. `down_block_types`: {down_block_types}."
              )
  
          if not isinstance(num_attention_heads, int) and len(num_attention_heads) != len(down_block_types):
              raise ValueError(
                  f"Must provide the same number of `num_attention_heads` as `down_block_types`. `num_attention_heads`: {num_attention_heads}. `down_block_types`: {down_block_types}."
              )
  
          if not isinstance(attention_head_dim, int) and len(attention_head_dim) != len(down_block_types):
              raise ValueError(
                  f"Must provide the same number of `attention_head_dim` as `down_block_types`. `attention_head_dim`: {attention_head_dim}. `down_block_types`: {down_block_types}."
              )
  
          if isinstance(cross_attention_dim, list) and len(cross_attention_dim) != len(down_block_types):
              raise ValueError(
                  f"Must provide the same number of `cross_attention_dim` as `down_block_types`. `cross_attention_dim`: {cross_attention_dim}. `down_block_types`: {down_block_types}."
              )
  
          if not isinstance(layers_per_block, int) and len(layers_per_block) != len(down_block_types):
              raise ValueError(
                  f"Must provide the same number of `layers_per_block` as `down_block_types`. `layers_per_block`: {layers_per_block}. `down_block_types`: {down_block_types}."
              )
  
          # input
          conv_in_padding = (conv_in_kernel - 1) // 2
          self.conv_in = nn.Conv2d(
              in_channels, block_out_channels[0], kernel_size=conv_in_kernel, padding=conv_in_padding
          )
  
          # time
          if time_embedding_type == "fourier":
              time_embed_dim = time_embedding_dim or block_out_channels[0] * 2
              if time_embed_dim % 2 != 0:
                  raise ValueError(f"`time_embed_dim` should be divisible by 2, but is {time_embed_dim}.")
              self.time_proj = GaussianFourierProjection(
                  time_embed_dim // 2, set_W_to_weight=False, log=False, flip_sin_to_cos=flip_sin_to_cos
              )
              timestep_input_dim = time_embed_dim
          elif time_embedding_type == "positional":
              time_embed_dim = time_embedding_dim or block_out_channels[0] * 4
  
              self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)
              timestep_input_dim = block_out_channels[0]
          else:
              raise ValueError(
                  f"{time_embedding_type} does not exist. Please make sure to use one of `fourier` or `positional`."
              )
  
          self.time_embedding = TimestepEmbedding(
              timestep_input_dim,
              time_embed_dim,
              act_fn=act_fn,
              post_act_fn=timestep_post_act,
              cond_proj_dim=time_cond_proj_dim,
          )
  
          if encoder_hid_dim_type is None and encoder_hid_dim is not None:
              encoder_hid_dim_type = "text_proj"
              self.register_to_config(encoder_hid_dim_type=encoder_hid_dim_type)
              logger.info("encoder_hid_dim_type defaults to 'text_proj' as `encoder_hid_dim` is defined.")
  
          if encoder_hid_dim is None and encoder_hid_dim_type is not None:
              raise ValueError(
                  f"`encoder_hid_dim` has to be defined when `encoder_hid_dim_type` is set to {encoder_hid_dim_type}."
              )
  
          if encoder_hid_dim_type == "text_proj":
              self.encoder_hid_proj = nn.Linear(encoder_hid_dim, cross_attention_dim)
          elif encoder_hid_dim_type == "text_image_proj":
              # image_embed_dim DOESN'T have to be `cross_attention_dim`. To not clutter the __init__ too much
              # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
              # case when `addition_embed_type == "text_image_proj"` (Kadinsky 2.1)`
              self.encoder_hid_proj = TextImageProjection(
                  text_embed_dim=encoder_hid_dim,
                  image_embed_dim=cross_attention_dim,
                  cross_attention_dim=cross_attention_dim,
              )
          elif encoder_hid_dim_type == "image_proj":
              # Kandinsky 2.2
              self.encoder_hid_proj = ImageProjection(
                  image_embed_dim=encoder_hid_dim,
                  cross_attention_dim=cross_attention_dim,
              )
          elif encoder_hid_dim_type is not None:
              raise ValueError(
                  f"encoder_hid_dim_type: {encoder_hid_dim_type} must be None, 'text_proj' or 'text_image_proj'."
              )
          else:
              self.encoder_hid_proj = None
  
          # class embedding
          if class_embed_type is None and num_class_embeds is not None:
              self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)
          elif class_embed_type == "timestep":
              self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim, act_fn=act_fn)
          elif class_embed_type == "identity":
              self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)
          elif class_embed_type == "projection":
              if projection_class_embeddings_input_dim is None:
                  raise ValueError(
                      "`class_embed_type`: 'projection' requires `projection_class_embeddings_input_dim` be set"
                  )
              # The projection `class_embed_type` is the same as the timestep `class_embed_type` except
              # 1. the `class_labels` inputs are not first converted to sinusoidal embeddings
              # 2. it projects from an arbitrary input dimension.
              #
              # Note that `TimestepEmbedding` is quite general, being mainly linear layers and activations.
              # When used for embedding actual timesteps, the timesteps are first converted to sinusoidal embeddings.
              # As a result, `TimestepEmbedding` can be passed arbitrary vectors.
              self.class_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
          elif class_embed_type == "simple_projection":
              if projection_class_embeddings_input_dim is None:
                  raise ValueError(
                      "`class_embed_type`: 'simple_projection' requires `projection_class_embeddings_input_dim` be set"
                  )
              self.class_embedding = nn.Linear(projection_class_embeddings_input_dim, time_embed_dim)
          else:
              self.class_embedding = None
  
          if addition_embed_type == "text":
              if encoder_hid_dim is not None:
                  text_time_embedding_from_dim = encoder_hid_dim
              else:
                  text_time_embedding_from_dim = cross_attention_dim
  
              self.add_embedding = TextTimeEmbedding(
                  text_time_embedding_from_dim, time_embed_dim, num_heads=addition_embed_type_num_heads
              )
          elif addition_embed_type == "text_image":
              # text_embed_dim and image_embed_dim DON'T have to be `cross_attention_dim`. To not clutter the __init__ too much
              # they are set to `cross_attention_dim` here as this is exactly the required dimension for the currently only use
              # case when `addition_embed_type == "text_image"` (Kadinsky 2.1)`
              self.add_embedding = TextImageTimeEmbedding(
                  text_embed_dim=cross_attention_dim, image_embed_dim=cross_attention_dim, time_embed_dim=time_embed_dim
              )
          elif addition_embed_type == "text_time":
              self.add_time_proj = Timesteps(addition_time_embed_dim, flip_sin_to_cos, freq_shift)
              self.add_embedding = TimestepEmbedding(projection_class_embeddings_input_dim, time_embed_dim)
          elif addition_embed_type == "image":
              # Kandinsky 2.2
              self.add_embedding = ImageTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
          elif addition_embed_type == "image_hint":
              # Kandinsky 2.2 ControlNet
              self.add_embedding = ImageHintTimeEmbedding(image_embed_dim=encoder_hid_dim, time_embed_dim=time_embed_dim)
          elif addition_embed_type is not None:
              raise ValueError(f"addition_embed_type: {addition_embed_type} must be None, 'text' or 'text_image'.")
  
          if time_embedding_act_fn is None:
              self.time_embed_act = None
          else:
              self.time_embed_act = get_activation(time_embedding_act_fn)
  
          self.down_blocks = nn.ModuleList([])
          self.up_blocks = nn.ModuleList([])
  
          if isinstance(only_cross_attention, bool):
              if mid_block_only_cross_attention is None:
                  mid_block_only_cross_attention = only_cross_attention
  
              only_cross_attention = [only_cross_attention] * len(down_block_types)
  
          if mid_block_only_cross_attention is None:
              mid_block_only_cross_attention = False
  
          if isinstance(num_attention_heads, int):
              num_attention_heads = (num_attention_heads,) * len(down_block_types)
  
          if isinstance(attention_head_dim, int):
              attention_head_dim = (attention_head_dim,) * len(down_block_types)
  
          if isinstance(cross_attention_dim, int):
              cross_attention_dim = (cross_attention_dim,) * len(down_block_types)
  
          if isinstance(layers_per_block, int):
              layers_per_block = [layers_per_block] * len(down_block_types)
  
          if isinstance(transformer_layers_per_block, int):
              transformer_layers_per_block = [transformer_layers_per_block] * len(down_block_types)
  
          if class_embeddings_concat:
              # The time embeddings are concatenated with the class embeddings. The dimension of the
              # time embeddings passed to the down, middle, and up blocks is twice the dimension of the
              # regular time embeddings
              blocks_time_embed_dim = time_embed_dim * 2
          else:
              blocks_time_embed_dim = time_embed_dim
  
          # down
          output_channel = block_out_channels[0]
          for i, down_block_type in enumerate(down_block_types):
              input_channel = output_channel
              output_channel = block_out_channels[i]
              is_final_block = i == len(block_out_channels) - 1
  
              down_block = get_down_block(
                  down_block_type,
                  num_layers=layers_per_block[i],
                  transformer_layers_per_block=transformer_layers_per_block[i],
                  in_channels=input_channel,
                  out_channels=output_channel,
                  temb_channels=blocks_time_embed_dim,
                  add_downsample=not is_final_block,
                  resnet_eps=norm_eps,
                  resnet_act_fn=act_fn,
                  resnet_groups=norm_num_groups,
                  cross_attention_dim=cross_attention_dim[i],
                  num_attention_heads=num_attention_heads[i],
                  downsample_padding=downsample_padding,
                  dual_cross_attention=dual_cross_attention,
                  use_linear_projection=use_linear_projection,
                  only_cross_attention=only_cross_attention[i],
                  upcast_attention=upcast_attention,
                  resnet_time_scale_shift=resnet_time_scale_shift,
                  resnet_skip_time_act=resnet_skip_time_act,
                  resnet_out_scale_factor=resnet_out_scale_factor,
                  cross_attention_norm=cross_attention_norm,
                  attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
                  num_views=num_views,
                  cd_attention_last=cd_attention_last,
                  cd_attention_mid=cd_attention_mid,
                  multiview_attention=multiview_attention,
                  sparse_mv_attention=sparse_mv_attention,
                  mvcd_attention=mvcd_attention
              )
              self.down_blocks.append(down_block)
  
          # mid
          if mid_block_type == "UNetMidBlock2DCrossAttn":
              self.mid_block = UNetMidBlock2DCrossAttn(
                  transformer_layers_per_block=transformer_layers_per_block[-1],
                  in_channels=block_out_channels[-1],
                  temb_channels=blocks_time_embed_dim,
                  resnet_eps=norm_eps,
                  resnet_act_fn=act_fn,
                  output_scale_factor=mid_block_scale_factor,
                  resnet_time_scale_shift=resnet_time_scale_shift,
                  cross_attention_dim=cross_attention_dim[-1],
                  num_attention_heads=num_attention_heads[-1],
                  resnet_groups=norm_num_groups,
                  dual_cross_attention=dual_cross_attention,
                  use_linear_projection=use_linear_projection,
                  upcast_attention=upcast_attention,
              )
          # custom MV2D attention block  
          elif mid_block_type == "UNetMidBlockMV2DCrossAttn":
              self.mid_block = UNetMidBlockMV2DCrossAttn(
                  transformer_layers_per_block=transformer_layers_per_block[-1],
                  in_channels=block_out_channels[-1],
                  temb_channels=blocks_time_embed_dim,
                  resnet_eps=norm_eps,
                  resnet_act_fn=act_fn,
                  output_scale_factor=mid_block_scale_factor,
                  resnet_time_scale_shift=resnet_time_scale_shift,
                  cross_attention_dim=cross_attention_dim[-1],
                  num_attention_heads=num_attention_heads[-1],
                  resnet_groups=norm_num_groups,
                  dual_cross_attention=dual_cross_attention,
                  use_linear_projection=use_linear_projection,
                  upcast_attention=upcast_attention,
                  num_views=num_views,
                  cd_attention_last=cd_attention_last,
                  cd_attention_mid=cd_attention_mid,
                  multiview_attention=multiview_attention,
                  sparse_mv_attention=sparse_mv_attention,
                  mvcd_attention=mvcd_attention
              )
          elif mid_block_type == "UNetMidBlock2DSimpleCrossAttn":
              self.mid_block = UNetMidBlock2DSimpleCrossAttn(
                  in_channels=block_out_channels[-1],
                  temb_channels=blocks_time_embed_dim,
                  resnet_eps=norm_eps,
                  resnet_act_fn=act_fn,
                  output_scale_factor=mid_block_scale_factor,
                  cross_attention_dim=cross_attention_dim[-1],
                  attention_head_dim=attention_head_dim[-1],
                  resnet_groups=norm_num_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 is None:
              self.mid_block = None
          else:
              raise ValueError(f"unknown mid_block_type : {mid_block_type}")
  
          # count how many layers upsample the images
          self.num_upsamplers = 0
  
          # up
          reversed_block_out_channels = list(reversed(block_out_channels))
          reversed_num_attention_heads = list(reversed(num_attention_heads))
          reversed_layers_per_block = list(reversed(layers_per_block))
          reversed_cross_attention_dim = list(reversed(cross_attention_dim))
          reversed_transformer_layers_per_block = list(reversed(transformer_layers_per_block))
          only_cross_attention = list(reversed(only_cross_attention))
  
          output_channel = reversed_block_out_channels[0]
          for i, up_block_type in enumerate(up_block_types):
              is_final_block = i == len(block_out_channels) - 1
  
              prev_output_channel = output_channel
              output_channel = reversed_block_out_channels[i]
              input_channel = reversed_block_out_channels[min(i + 1, len(block_out_channels) - 1)]
  
              # add upsample block for all BUT final layer
              if not is_final_block:
                  add_upsample = True
                  self.num_upsamplers += 1
              else:
                  add_upsample = False
  
              up_block = get_up_block(
                  up_block_type,
                  num_layers=reversed_layers_per_block[i] + 1,
                  transformer_layers_per_block=reversed_transformer_layers_per_block[i],
                  in_channels=input_channel,
                  out_channels=output_channel,
                  prev_output_channel=prev_output_channel,
                  temb_channels=blocks_time_embed_dim,
                  add_upsample=add_upsample,
                  resnet_eps=norm_eps,
                  resnet_act_fn=act_fn,
                  resnet_groups=norm_num_groups,
                  cross_attention_dim=reversed_cross_attention_dim[i],
                  num_attention_heads=reversed_num_attention_heads[i],
                  dual_cross_attention=dual_cross_attention,
                  use_linear_projection=use_linear_projection,
                  only_cross_attention=only_cross_attention[i],
                  upcast_attention=upcast_attention,
                  resnet_time_scale_shift=resnet_time_scale_shift,
                  resnet_skip_time_act=resnet_skip_time_act,
                  resnet_out_scale_factor=resnet_out_scale_factor,
                  cross_attention_norm=cross_attention_norm,
                  attention_head_dim=attention_head_dim[i] if attention_head_dim[i] is not None else output_channel,
                  num_views=num_views,
                  cd_attention_last=cd_attention_last,
                  cd_attention_mid=cd_attention_mid,
                  multiview_attention=multiview_attention,
                  sparse_mv_attention=sparse_mv_attention,
                  mvcd_attention=mvcd_attention
              )
              self.up_blocks.append(up_block)
              prev_output_channel = output_channel
  
          # out
          if norm_num_groups is not None:
              self.conv_norm_out = nn.GroupNorm(
                  num_channels=block_out_channels[0], num_groups=norm_num_groups, eps=norm_eps
              )
  
              self.conv_act = get_activation(act_fn)
  
          else:
              self.conv_norm_out = None
              self.conv_act = None
  
          conv_out_padding = (conv_out_kernel - 1) // 2
          self.conv_out = nn.Conv2d(
  