# Copyright 2022 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.

import numpy as np

import torch

from torch import nn

from .attention import AttentionBlock, CrossAttention, DualTransformer2DModel, Transformer2DModel

from .resnet import Downsample2D, FirDownsample2D, FirUpsample2D, ResnetBlock2D, Upsample2D

def get_down_block(

down_block_type,

num_layers,

in_channels,

out_channels,

temb_channels,

add_downsample,

resnet_eps,

resnet_act_fn,

attn_num_head_channels,

resnet_groups=None,

cross_attention_dim=None,

downsample_padding=None,

dual_cross_attention=False,

use_linear_projection=False,

only_cross_attention=False,

upcast_attention=False,

resnet_time_scale_shift="default",

):

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,

)

elif down_block_type == "AttnDownBlock2D":

return AttnDownBlock2D(

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,

attn_num_head_channels=attn_num_head_channels,

resnet_time_scale_shift=resnet_time_scale_shift,

)

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,

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,

attn_num_head_channels=attn_num_head_channels,

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,

)

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,

attn_num_head_channels=attn_num_head_channels,

resnet_time_scale_shift=resnet_time_scale_shift,

)

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,

downsample_padding=downsample_padding,

attn_num_head_channels=attn_num_head_channels,

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,

attn_num_head_channels=attn_num_head_channels,

resnet_time_scale_shift=resnet_time_scale_shift,

)

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,

attn_num_head_channels,

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

):

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,

)

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,

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,

attn_num_head_channels=attn_num_head_channels,

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,

)

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,

attn_num_head_channels=attn_num_head_channels,

resnet_time_scale_shift=resnet_time_scale_shift,

)

elif up_block_type == "AttnUpBlock2D":

return AttnUpBlock2D(

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,

attn_num_head_channels=attn_num_head_channels,

resnet_time_scale_shift=resnet_time_scale_shift,

)

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,

attn_num_head_channels=attn_num_head_channels,

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,

)

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,

attn_num_head_channels=attn_num_head_channels,

resnet_time_scale_shift=resnet_time_scale_shift,

)

raise ValueError(f"{up_block_type} does not exist.")

class UNetMidBlock2D(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,

add_attention: bool = True,

attn_num_head_channels=1,

output_scale_factor=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

# 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 self.add_attention:

attentions.append(

AttentionBlock(

in_channels,

num_head_channels=attn_num_head_channels,

rescale_output_factor=output_scale_factor,

eps=resnet_eps,

norm_num_groups=resnet_groups,

)

)

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, temb=None):

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)

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,

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,

attn_num_head_channels=1,

output_scale_factor=1.0,

cross_attention_dim=1280,

dual_cross_attention=False,

use_linear_projection=False,

upcast_attention=False,

):

super().__init__()

self.has_cross_attention = True

self.attn_num_head_channels = attn_num_head_channels

resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)

# 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(

Transformer2DModel(

attn_num_head_channels,

in_channels // attn_num_head_channels,

in_channels=in_channels,

num_layers=1,

cross_attention_dim=cross_attention_dim,

norm_num_groups=resnet_groups,

use_linear_projection=use_linear_projection,

upcast_attention=upcast_attention,

)

)

else:

attentions.append(

DualTransformer2DModel(

attn_num_head_channels,

in_channels // attn_num_head_channels,

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)

def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):

# TODO(Patrick, William) - attention_mask is currently not used. Implement once used

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

hidden_states = resnet(hidden_states, temb)

return hidden_states

class UNetMidBlock2DSimpleCrossAttn(nn.Module):

def __init__(

self,

in_channels: int,

temb_channels: int,

dropout: float = 0.0,

num_layers: int = 1,

resnet_eps: float = 1e-6,

resnet_time_scale_shift: str = "default",

resnet_act_fn: str = "swish",

resnet_groups: int = 32,

resnet_pre_norm: bool = True,

attn_num_head_channels=1,

output_scale_factor=1.0,

cross_attention_dim=1280,

):

super().__init__()

self.has_cross_attention = True

self.attn_num_head_channels = attn_num_head_channels

resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)

self.num_heads = in_channels // self.attn_num_head_channels

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

attentions.append(

CrossAttention(

query_dim=in_channels,

cross_attention_dim=in_channels,

heads=self.num_heads,

dim_head=attn_num_head_channels,

added_kv_proj_dim=cross_attention_dim,

norm_num_groups=resnet_groups,

bias=True,

upcast_softmax=True,

)

)

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, temb=None, encoder_hidden_states=None, attention_mask=None):

hidden_states = self.resnets[0](hidden_states, temb)

for attn, resnet in zip(self.attentions, self.resnets[1:]):

# attn

residual = hidden_states

hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)

hidden_states = attn(

hidden_states,

encoder_hidden_states=encoder_hidden_states.transpose(1, 2),

attention_mask=attention_mask,

)

hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)

hidden_states = hidden_states + residual

# resnet

hidden_states = resnet(hidden_states, temb)

return hidden_states

class AttnDownBlock2D(nn.Module):

def __init__(

self,

in_channels: int,

out_channels: int,

temb_channels: int,

dropout: float = 0.0,

num_layers: int = 1,

resnet_eps: float = 1e-6,

resnet_time_scale_shift: str = "default",

resnet_act_fn: str = "swish",

resnet_groups: int = 32,

resnet_pre_norm: bool = True,

attn_num_head_channels=1,

output_scale_factor=1.0,

downsample_padding=1,

add_downsample=True,

):

super().__init__()

resnets = []

attentions = []

for i in range(num_layers):

in_channels = in_channels if i == 0 else out_channels

resnets.append(

ResnetBlock2D(

in_channels=in_channels,

out_channels=out_channels,

temb_channels=temb_channels,

eps=resnet_eps,

groups=resnet_groups,

dropout=dropout,

time_embedding_norm=resnet_time_scale_shift,

non_linearity=resnet_act_fn,

output_scale_factor=output_scale_factor,

pre_norm=resnet_pre_norm,

)

)

attentions.append(

AttentionBlock(

out_channels,

num_head_channels=attn_num_head_channels,

rescale_output_factor=output_scale_factor,

eps=resnet_eps,

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

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

output_states = ()

for resnet, attn in zip(self.resnets, self.attentions):

hidden_states = resnet(hidden_states, temb)

hidden_states = attn(hidden_states)

output_states += (hidden_states,)

if self.downsamplers is not None:

for downsampler in self.downsamplers:

hidden_states = downsampler(hidden_states)

output_states += (hidden_states,)

return hidden_states, output_states

class CrossAttnDownBlock2D(nn.Module):

def __init__(

self,

in_channels: int,

out_channels: int,

temb_channels: int,

dropout: float = 0.0,

num_layers: int = 1,

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,

attn_num_head_channels=1,

cross_attention_dim=1280,

output_scale_factor=1.0,

downsample_padding=1,

add_downsample=True,

dual_cross_attention=False,

use_linear_projection=False,

only_cross_attention=False,

upcast_attention=False,

):

super().__init__()

resnets = []

attentions = []

self.has_cross_attention = True

self.attn_num_head_channels = attn_num_head_channels

for i in range(num_layers):

in_channels = in_channels if i == 0 else out_channels

resnets.append(

ResnetBlock2D(

in_channels=in_channels,

out_channels=out_channels,

temb_channels=temb_channels,

eps=resnet_eps,

groups=resnet_groups,

dropout=dropout,

time_embedding_norm=resnet_time_scale_shift,

non_linearity=resnet_act_fn,

output_scale_factor=output_scale_factor,

pre_norm=resnet_pre_norm,

)

)

if not dual_cross_attention:

attentions.append(

Transformer2DModel(

attn_num_head_channels,

out_channels // attn_num_head_channels,

in_channels=out_channels,

num_layers=1,

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,

)

)

else:

attentions.append(

DualTransformer2DModel(

attn_num_head_channels,

out_channels // attn_num_head_channels,

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, temb=None, encoder_hidden_states=None, attention_mask=None):

# TODO(Patrick, William) - attention mask is not used

output_states = ()

for resnet, attn in zip(self.resnets, self.attentions):

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

hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)

hidden_states = torch.utils.checkpoint.checkpoint(

create_custom_forward(attn, return_dict=False),

hidden_states,

encoder_hidden_states,

)[0]

else:

hidden_states = resnet(hidden_states, temb)

hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample

output_states += (hidden_states,)

if self.downsamplers is not None:

for downsampler in self.downsamplers:

hidden_states = downsampler(hidden_states)

output_states += (hidden_states,)

return hidden_states, output_states

class DownBlock2D(nn.Module):

def __init__(

self,

in_channels: int,

out_channels: int,

temb_channels: int,

dropout: float = 0.0,

num_layers: int = 1,

resnet_eps: float = 1e-6,

resnet_time_scale_shift: str = "default",

resnet_act_fn: str = "swish",

resnet_groups: int = 32,

resnet_pre_norm: bool = True,

output_scale_factor=1.0,

add_downsample=True,

downsample_padding=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, temb=None):

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

hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)

else:

hidden_states = resnet(hidden_states, temb)

output_states += (hidden_states,)

if self.downsamplers is not None:

for downsampler in self.downsamplers:

hidden_states = downsampler(hidden_states)

output_states += (hidden_states,)

return hidden_states, output_states

class DownEncoderBlock2D(nn.Module):

def __init__(

self,

in_channels: int,

out_channels: int,

dropout: float = 0.0,

num_layers: int = 1,

resnet_eps: float = 1e-6,

resnet_time_scale_shift: str = "default",

resnet_act_fn: str = "swish",

resnet_groups: int = 32,

resnet_pre_norm: bool = True,

output_scale_factor=1.0,

add_downsample=True,

downsample_padding=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=None,

eps=resnet_eps,

groups=resnet_groups,

dropout=dropout,

time_embedding_norm=resnet_time_scale_shift,

non_linearity=resnet_act_fn,

output_scale_factor=output_scale_factor,

pre_norm=resnet_pre_norm,

)

)

self.resnets = nn.ModuleList(resnets)

if add_downsample:

self.downsamplers = nn.ModuleList(

[

Downsample2D(

out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"

)

]

)

else:

self.downsamplers = None

def forward(self, hidden_states):

for resnet in self.resnets:

hidden_states = resnet(hidden_states, temb=None)

if self.downsamplers is not None:

for downsampler in self.downsamplers:

hidden_states = downsampler(hidden_states)

return hidden_states

class AttnDownEncoderBlock2D(nn.Module):

def __init__(

self,

in_channels: int,

out_channels: int,

dropout: float = 0.0,

num_layers: int = 1,

resnet_eps: float = 1e-6,

resnet_time_scale_shift: str = "default",

resnet_act_fn: str = "swish",

resnet_groups: int = 32,

resnet_pre_norm: bool = True,

attn_num_head_channels=1,

output_scale_factor=1.0,

add_downsample=True,

downsample_padding=1,

):

super().__init__()

resnets = []

attentions = []

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=None,

eps=resnet_eps,

groups=resnet_groups,

dropout=dropout,

time_embedding_norm=resnet_time_scale_shift,

non_linearity=resnet_act_fn,

output_scale_factor=output_scale_factor,

pre_norm=resnet_pre_norm,

)

)

attentions.append(

AttentionBlock(

out_channels,

num_head_channels=attn_num_head_channels,

rescale_output_factor=output_scale_factor,

eps=resnet_eps,

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

def forward(self, hidden_states):

for resnet, attn in zip(self.resnets, self.attentions):

hidden_states = resnet(hidden_states, temb=None)

hidden_states = attn(hidden_states)

if self.downsamplers is not None:

for downsampler in self.downsamplers:

hidden_states = downsampler(hidden_states)

return hidden_states