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 typing import Any, Dict, Optional, Tuple
  
  import numpy as np
  import torch
  import torch.nn.functional as F
  from torch import nn
  
  from diffusers.utils import is_torch_version, logging
  from diffusers.models.normalization import AdaGroupNorm
  from diffusers.models.attention_processor import Attention, AttnAddedKVProcessor, AttnAddedKVProcessor2_0
  from diffusers.models.transformers.dual_transformer_2d import DualTransformer2DModel
  from diffusers.models.resnet import Downsample2D, FirDownsample2D, FirUpsample2D, KDownsample2D, KUpsample2D, ResnetBlock2D, Upsample2D
  
  from diffusers.models.unets.unet_2d_blocks import DownBlock2D, ResnetDownsampleBlock2D, AttnDownBlock2D, CrossAttnDownBlock2D, SimpleCrossAttnDownBlock2D, SkipDownBlock2D, AttnSkipDownBlock2D, DownEncoderBlock2D, AttnDownEncoderBlock2D, KDownBlock2D, KCrossAttnDownBlock2D
  from diffusers.models.unets.unet_2d_blocks import UpBlock2D, ResnetUpsampleBlock2D, CrossAttnUpBlock2D, SimpleCrossAttnUpBlock2D, AttnUpBlock2D, SkipUpBlock2D, AttnSkipUpBlock2D, UpDecoderBlock2D, AttnUpDecoderBlock2D, KUpBlock2D, KCrossAttnUpBlock2D
  
  
  logger = logging.get_logger(__name__)  # pylint: disable=invalid-name
  
  
  def get_down_block(
      down_block_type,
      num_layers,
      in_channels,
      out_channels,
      temb_channels,
      add_downsample,
      resnet_eps,
      resnet_act_fn,
      transformer_layers_per_block=1,
      num_attention_heads=None,
      resnet_groups=None,
      cross_attention_dim=None,
      downsample_padding=None,
      dual_cross_attention=False,
      use_linear_projection=False,
      only_cross_attention=False,
      upcast_attention=False,
      resnet_time_scale_shift="default",
      resnet_skip_time_act=False,
      resnet_out_scale_factor=1.0,
      cross_attention_norm=None,
      attention_head_dim=None,
      downsample_type=None,
      num_views=1,
      cd_attention_last: bool = False,
      cd_attention_mid: bool = False,
      multiview_attention: bool = True,
      sparse_mv_attention: bool = False,
      selfattn_block: str = "custom",
      mvcd_attention: bool=False,
      use_dino: bool = False
  ):
      # If attn head dim is not defined, we default it to the number of heads
      if attention_head_dim is None:
          logger.warn(
              f"It is recommended to provide `attention_head_dim` when calling `get_down_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
          )
          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,
              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,
              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"  # default to 'conv'
          return AttnDownBlock2D(
              num_layers=num_layers,
              in_channels=in_channels,
              out_channels=out_channels,
              temb_channels=temb_channels,
              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,
              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,
          )
      # custom MV2D attention block
      elif down_block_type == "CrossAttnDownBlockMV2D":
          if cross_attention_dim is None:
              raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlockMV2D")
          return CrossAttnDownBlockMV2D(
              num_layers=num_layers,
              transformer_layers_per_block=transformer_layers_per_block,
              in_channels=in_channels,
              out_channels=out_channels,
              temb_channels=temb_channels,
              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,
              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,
              selfattn_block=selfattn_block,
              mvcd_attention=mvcd_attention,
              use_dino=use_dino
          )
      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,
              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,
              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,
              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,
              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,
              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,
              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,
              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_up_block(
      up_block_type,
      num_layers,
      in_channels,
      out_channels,
      prev_output_channel,
      temb_channels,
      add_upsample,
      resnet_eps,
      resnet_act_fn,
      transformer_layers_per_block=1,
      num_attention_heads=None,
      resnet_groups=None,
      cross_attention_dim=None,
      dual_cross_attention=False,
      use_linear_projection=False,
      only_cross_attention=False,
      upcast_attention=False,
      resnet_time_scale_shift="default",
      resnet_skip_time_act=False,
      resnet_out_scale_factor=1.0,
      cross_attention_norm=None,
      attention_head_dim=None,
      upsample_type=None,
      num_views=1,
      cd_attention_last: bool = False,
      cd_attention_mid: bool = False,
      multiview_attention: bool = True,
      sparse_mv_attention: bool = False,
      selfattn_block: str = "custom",
      mvcd_attention: bool=False,
      use_dino: bool = False  
  ):
      # If attn head dim is not defined, we default it to the number of heads
      if attention_head_dim is None:
          logger.warn(
              f"It is recommended to provide `attention_head_dim` when calling `get_up_block`. Defaulting `attention_head_dim` to {num_attention_heads}."
          )
          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,
              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,
              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,
              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,
          )
      # custom MV2D attention block
      elif up_block_type == "CrossAttnUpBlockMV2D":
          if cross_attention_dim is None:
              raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlockMV2D")
          return CrossAttnUpBlockMV2D(
              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,
              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,
              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,
              selfattn_block=selfattn_block,
              mvcd_attention=mvcd_attention,
              use_dino=use_dino
          )    
      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,
              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"  # default to '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,
              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,
              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,
              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,
              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,
              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,
              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,
              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 UNetMidBlockMV2DCrossAttn(nn.Module):
      def __init__(
          self,
          in_channels: int,
          temb_channels: int,
          dropout: float = 0.0,
          num_layers: int = 1,
          transformer_layers_per_block: int = 1,
          resnet_eps: float = 1e-6,
          resnet_time_scale_shift: str = "default",
          resnet_act_fn: str = "swish",
          resnet_groups: int = 32,
          resnet_pre_norm: bool = True,
          num_attention_heads=1,
          output_scale_factor=1.0,
          cross_attention_dim=1280,
          dual_cross_attention=False,
          use_linear_projection=False,
          upcast_attention=False,
          num_views: int = 1,
          cd_attention_last: bool = False,
          cd_attention_mid: bool = False,
          multiview_attention: bool = True,
          sparse_mv_attention: bool = False,
          selfattn_block: str = "custom", 
          mvcd_attention: bool=False,
          use_dino: bool = False
      ):
          super().__init__()
  
          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 selfattn_block == "custom":
              from .transformer_mv2d import TransformerMV2DModel
          elif selfattn_block == "rowwise":
              from .transformer_mv2d_rowwise import TransformerMV2DModel
          elif selfattn_block == "self_rowwise":
              from .transformer_mv2d_self_rowwise import TransformerMV2DModel
          else:
              raise NotImplementedError
          
          # there is always at least one resnet
          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 = []
  
          for _ in range(num_layers):
              if not dual_cross_attention:
                  attentions.append(
                      TransformerMV2DModel(
                          num_attention_heads,
                          in_channels // num_attention_heads,
                          in_channels=in_channels,
                          num_layers=transformer_layers_per_block,
                          cross_attention_dim=cross_attention_dim,
                          norm_num_groups=resnet_groups,
                          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,
                          use_dino=use_dino
                      )
                  )
              else:
                  raise NotImplementedError
              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)
  
      def forward(
          self,
          hidden_states: torch.FloatTensor,
          temb: Optional[torch.FloatTensor] = None,
          encoder_hidden_states: Optional[torch.FloatTensor] = None,
          attention_mask: Optional[torch.FloatTensor] = None,
          cross_attention_kwargs: Optional[Dict[str, Any]] = None,
          encoder_attention_mask: Optional[torch.FloatTensor] = None,
          dino_feature: Optional[torch.FloatTensor] = None
      ) -> torch.FloatTensor:
          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,
                  cross_attention_kwargs=cross_attention_kwargs,
                  attention_mask=attention_mask,
                  encoder_attention_mask=encoder_attention_mask,
                  dino_feature=dino_feature,
                  return_dict=False,
              )[0]
              hidden_states = resnet(hidden_states, temb)
  
          return hidden_states
  
  
  class CrossAttnUpBlockMV2D(nn.Module):
      def __init__(
          self,
          in_channels: int,
          out_channels: int,
          prev_output_channel: int,
          temb_channels: int,
          dropout: float = 0.0,
          num_layers: int = 1,
          transformer_layers_per_block: int = 1,
          resnet_eps: float = 1e-6,
          resnet_time_scale_shift: str = "default",
          resnet_act_fn: str = "swish",
          resnet_groups: int = 32,
          resnet_pre_norm: bool = True,
          num_attention_heads=1,
          cross_attention_dim=1280,
          output_scale_factor=1.0,
          add_upsample=True,
          dual_cross_attention=False,
          use_linear_projection=False,
          only_cross_attention=False,
          upcast_attention=False,
          num_views: int = 1,
          cd_attention_last: bool = False,
          cd_attention_mid: bool = False,
          multiview_attention: bool = True,
          sparse_mv_attention: bool = False,
          selfattn_block: str = "custom",
          mvcd_attention: bool=False,
          use_dino: bool = False
      ):
          super().__init__()
          resnets = []
          attentions = []
  
          self.has_cross_attention = True
          self.num_attention_heads = num_attention_heads
  
          if selfattn_block == "custom":
              from .transformer_mv2d import TransformerMV2DModel
          elif selfattn_block == "rowwise":
              from .transformer_mv2d_rowwise import TransformerMV2DModel
          elif selfattn_block == "self_rowwise":
              from .transformer_mv2d_self_rowwise import TransformerMV2DModel
          else:
              raise NotImplementedError
          
          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(
                      TransformerMV2DModel(
                          num_attention_heads,
                          out_channels // num_attention_heads,
                          in_channels=out_channels,
                          num_layers=transformer_layers_per_block,
                          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,
                          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,
                          use_dino=use_dino
                      )
                  )
              else:
                  raise NotImplementedError
          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
  
      def forward(
          self,
          hidden_states: torch.FloatTensor,
          res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
          temb: Optional[torch.FloatTensor] = None,
          encoder_hidden_states: Optional[torch.FloatTensor] = None,
          cross_attention_kwargs: Optional[Dict[str, Any]] = None,
          upsample_size: Optional[int] = None,
          attention_mask: Optional[torch.FloatTensor] = None,
          encoder_attention_mask: Optional[torch.FloatTensor] = None,
          dino_feature: Optional[torch.FloatTensor] = None
      ):
          for resnet, attn in zip(self.resnets, self.attentions):
              # pop res hidden states
              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 self.training 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 = torch.utils.checkpoint.checkpoint(
                      create_custom_forward(attn, return_dict=False),
                      hidden_states,
                      encoder_hidden_states,
                      dino_feature,
                      None,  # timestep
                      None,  # class_labels
                      cross_attention_kwargs,
                      attention_mask,
                      encoder_attention_mask,
                      **ckpt_kwargs,
                  )[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,
                      dino_feature=dino_feature,
                      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 CrossAttnDownBlockMV2D(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: int = 1,
  