MuseV/musev/models/unet_2d_blocks.py

# 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, Literal, Optional, Tuple, Union, List

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.utils.torch_utils import apply_freeu
from diffusers.models.activations import get_activation
from diffusers.models.attention_processor import (
 Attention,
 AttnAddedKVProcessor,
 AttnAddedKVProcessor2_0,
)
from diffusers.models.dual_transformer_2d import DualTransformer2DModel
from diffusers.models.normalization import AdaGroupNorm
from diffusers.models.resnet import (
 Downsample2D,
 FirDownsample2D,
 FirUpsample2D,
 KDownsample2D,
 KUpsample2D,
 ResnetBlock2D,
 Upsample2D,
)
from diffusers.models.unet_2d_blocks import (
 AttnDownBlock2D,
 AttnDownEncoderBlock2D,
 AttnSkipDownBlock2D,
 AttnSkipUpBlock2D,
 AttnUpBlock2D,
 AttnUpDecoderBlock2D,
 DownEncoderBlock2D,
 KCrossAttnDownBlock2D,
 KCrossAttnUpBlock2D,
 KDownBlock2D,
 KUpBlock2D,
 ResnetDownsampleBlock2D,
 ResnetUpsampleBlock2D,
 SimpleCrossAttnDownBlock2D,
 SimpleCrossAttnUpBlock2D,
 SkipDownBlock2D,
 SkipUpBlock2D,
 UpDecoderBlock2D,
)

from .transformer_2d import Transformer2DModel


logger = logging.get_logger(__name__) # pylint: disable=invalid-name


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 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,
 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" # default to '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_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 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,
 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" # 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,
 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 UNetMidBlock2D(nn.Module):
 """
 A 2D UNet mid-block [`UNetMidBlock2D`] with multiple residual blocks and optional attention blocks.

 Args:
 in_channels (`int`): The number of input channels.
 temb_channels (`int`): The number of temporal embedding channels.
 dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
 num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
 resnet_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
 resnet_time_scale_shift (`str`, *optional*, defaults to `default`):
 The type of normalization to apply to the time embeddings. This can help to improve the performance of the
 model on tasks with long-range temporal dependencies.
 resnet_act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks.
 resnet_groups (`int`, *optional*, defaults to 32):
 The number of groups to use in the group normalization layers of the resnet blocks.
 attn_groups (`Optional[int]`, *optional*, defaults to None): The number of groups for the attention blocks.
 resnet_pre_norm (`bool`, *optional*, defaults to `True`):
 Whether to use pre-normalization for the resnet blocks.
 add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks.
 attention_head_dim (`int`, *optional*, defaults to 1):
 Dimension of a single attention head. The number of attention heads is determined based on this value and
 the number of input channels.
 output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor.

 Returns:
 `torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size,
 in_channels, height, width)`.

 """

 def __init__(
 self,
 in_channels: int,
 temb_channels: int,
 dropout: float = 0.0,
 num_layers: int = 1,
 resnet_eps: float = 1e-6,
 resnet_time_scale_shift: str = "default", # default, spatial
 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
 )

 # 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 = []

 if attention_head_dim is None:
 logger.warn(
 f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {in_channels}."
 )
 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)

 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,
 self_attn_block_embs: Optional[List[torch.Tensor]] = None,
 self_attn_block_embs_mode: Literal["read", "write"] = "write",
 ) -> torch.FloatTensor:
 hidden_states = self.resnets[0](hidden_states, temb)
 for attn, resnet in zip(self.attentions, self.resnets[1:]):
 if attn is not None:
 hidden_states = attn(
 hidden_states,
 temb=temb,
 self_attn_block_embs=self_attn_block_embs,
 self_attn_block_embs_mode=self_attn_block_embs_mode,
 )
 hidden_states = resnet(hidden_states, temb)

 return hidden_states


class UNetMidBlock2DCrossAttn(nn.Module):
 def __init__(
 self,
 in_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-6,
 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,
 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__()

 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)
 )

 # support for variable transformer layers per block
 if isinstance(transformer_layers_per_block, int):
 transformer_layers_per_block = [transformer_layers_per_block] * num_layers

 # 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 i in range(num_layers):
 if not dual_cross_attention:
 attentions.append(
 Transformer2DModel(
 num_attention_heads,
 in_channels // num_attention_heads,
 in_channels=in_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,
 upcast_attention=upcast_attention,
 attention_type=attention_type,
 )
 )
 else:
 attentions.append(
 DualTransformer2DModel(
 num_attention_heads,
 in_channels // num_attention_heads,
 in_channels=in_channels,
 num_layers=1,
 cross_attention_dim=cross_attention_dim,
 norm_num_groups=resnet_groups,
 )
 )
 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.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,
 self_attn_block_embs: Optional[List[torch.Tensor]] = None,
 self_attn_block_embs_mode: Literal["read", "write"] = "write",
 ) -> torch.FloatTensor:
 lora_scale = (
 cross_attention_kwargs.get("scale", 1.0)
 if cross_attention_kwargs is not None
 else 1.0
 )
 hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)
 for attn, resnet in zip(self.attentions, self.resnets[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 = 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,
 self_attn_block_embs=self_attn_block_embs,
 self_attn_block_embs_mode=self_attn_block_embs_mode,
 )[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,
 self_attn_block_embs=self_attn_block_embs,
 )[0]
 hidden_states = resnet(hidden_states, temb, scale=lora_scale)

 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-6,
 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

 # 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,
 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.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,
 self_attn_block_embs: Optional[List[torch.Tensor]] = None,
 self_attn_block_embs_mode: Literal["read", "write"] = "write",
 ) -> torch.FloatTensor:
 cross_attention_kwargs = (
 cross_attention_kwargs if cross_attention_kwargs is not None else {}
 )
 lora_scale = cross_attention_kwargs.get("scale", 1.0)

 if attention_mask is None:
 # if encoder_hidden_states is defined: we are doing cross-attn, so we should use cross-attn mask.
 mask = None if encoder_hidden_states is None else encoder_attention_mask
 else:
 # when attention_mask is defined: we don't even check for encoder_attention_mask.
 # this is to maintain compatibility with UnCLIP, which uses 'attention_mask' param for cross-attn masks.
 # TODO: UnCLIP should express cross-attn mask via encoder_attention_mask param instead of via attention_mask.
 # then we can simplify this whole if/else block to:
 # mask = attention_mask if encoder_hidden_states is None else encoder_attention_mask
 mask = attention_mask

 hidden_states = self.resnets[0](hidden_states, temb, scale=lora_scale)
 for attn, resnet in zip(self.attentions, self.resnets[1:]):
 # attn
 hidden_states = attn(
 hidden_states,
 encoder_hidden_states=encoder_hidden_states,
 attention_mask=mask,
 **cross_attention_kwargs,
 self_attn_block_embs=self_attn_block_embs,
 self_attn_block_embs_mode=self_attn_block_embs_mode,
 )

 # resnet
 hidden_states = resnet(hidden_states, temb, scale=lora_scale)

 return hidden_states


class CrossAttnDownBlock2D(nn.Module):
 print_idx = 0

 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-6,
 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.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,
 additional_residuals: Optional[torch.FloatTensor] = None,
 self_attn_block_embs: Optional[List[torch.Tensor]] = None,
 self_attn_block_embs_mode: Literal["read", "write"] = "write",
 ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
 output_states = ()

 lora_scale = (
 cross_attention_kwargs.get("scale", 1.0)
 if cross_attention_kwargs is not None
 else 1.0
 )

 blocks = list(zip(self.resnets, self.attentions))

 for i, (resnet, attn) in enumerate(blocks):
 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,
 )
 if self.print_idx == 0:
 logger.debug(f"unet3d after resnet {hidden_states.mean()}")

 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,
 self_attn_block_embs=self_attn_block_embs,
 self_attn_block_embs_mode=self_attn_block_embs_mode,
 )[0]
 else:
 hidden_states = resnet(hidden_states, temb, scale=lora_scale)
 if self.print_idx == 0:
 logger.debug(f"unet3d after resnet {hidden_states.mean()}")
 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,
 self_attn_block_embs=self_attn_block_embs,
 self_attn_block_embs_mode=self_attn_block_embs_mode,
 )[0]

 # apply additional residuals to the output of the last pair of resnet and attention blocks
 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, scale=lora_scale)

 output_states = output_states + (hidden_states,)

 self.print_idx += 1
 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-6,
 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.FloatTensor,
 temb: Optional[torch.FloatTensor] = None,
 scale: float = 1.0,
 self_attn_block_embs: Optional[List[torch.Tensor]] = None,
 self_attn_block_embs_mode: Literal["read", "write"] = "write",
 ) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
 output_states = ()

 for resnet in self.resnets:
 if self.training 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, scale=scale)

 output_states = output_states + (hidden_states,)

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

 output_states = output_states + (hidden_states,)

 return hidden_states, output_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-6,
 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.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,
 self_attn_block_embs: Optional[List[torch.Tensor]] = None,
 self_attn_block_embs_mode: Literal["read", "write"] = "write",
 ) -> torch.FloatTensor:
 lora_scale = (
 cross_attention_kwargs.get("scale", 1.0)
 if cross_attention_kwargs is not None
 else 1.0
 )
 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):
 # pop res hidden states
 res_hidden_states = res_hidden_states_tuple[-1]
 res_hidden_states_tuple = res_hidden_states_tuple[:-1]

 # FreeU: Only operate on the first two stages
 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 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 = 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,
 self_attn_block_embs=self_attn_block_embs,
 self_attn_block_embs_mode=self_attn_block_embs_mode,
 )[0]
 else:
 hidden_states = resnet(hidden_states, temb, scale=lora_scale)
 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,
 self_attn_block_embs=self_attn_block_embs,
 )[0]

 if self.upsamplers is not None:
 for upsampler in self.upsamplers:
 hidden_states = upsampler(
 hidden_states, upsample_size, scale=lora_scale
 )

 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-6,
 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.FloatTensor,
 res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
 temb: Optional[torch.FloatTensor] = None,
 upsample_size: Optional[int] = None,
 scale: float = 1.0,
 self_attn_block_embs: Optional[List[torch.Tensor]] = None,
 self_attn_block_embs_mode: Literal["read", "write"] = "write",
 ) -> torch.FloatTensor:
 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:
 # pop res hidden states
 res_hidden_states = res_hidden_states_tuple[-1]
 res_hidden_states_tuple = res_hidden_states_tuple[:-1]

 # FreeU: Only operate on the first two stages
 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 self.training 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, scale=scale)

 if self.upsamplers is not None:
 for upsampler in self.upsamplers:
 hidden_states = upsampler(hidden_states, upsample_size, scale=scale)

 return hidden_states