Spaces:

lemonaddie
/

geowizard

Running on Zero

App Files Files Community

geowizard / models /unet_2d_blocks.py

lemonaddie

Upload 9 files

87f795e verified 3 months ago

raw history blame contribute delete

No virus

136 kB

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

	# Some modifications are reimplemented in public environments by Xiao Fu and Mu Hu

	from typing import Any, Dict, Optional, Tuple, Union

	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 models.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 AutoencoderTinyBlock(nn.Module):
	"""
	Tiny Autoencoder block used in [`AutoencoderTiny`]. It is a mini residual module consisting of plain conv + ReLU
	blocks.

	Args:
	in_channels (`int`): The number of input channels.
	out_channels (`int`): The number of output channels.
	act_fn (`str`):
	` The activation function to use. Supported values are `"swish"`, `"mish"`, `"gelu"`, and `"relu"`.

	Returns:
	`torch.FloatTensor`: A tensor with the same shape as the input tensor, but with the number of channels equal to
	`out_channels`.
	"""

	def __init__(self, in_channels: int, out_channels: int, act_fn: str):
	super().__init__()
	act_fn = get_activation(act_fn)
	self.conv = nn.Sequential(
	nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
	act_fn,
	nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
	act_fn,
	nn.Conv2d(out_channels, out_channels, kernel_size=3, padding=1),
	)
	self.skip = (
	nn.Conv2d(in_channels, out_channels, kernel_size=1, bias=False)
	if in_channels != out_channels
	else nn.Identity()
	)
	self.fuse = nn.ReLU()

	def forward(self, x: torch.FloatTensor) -> torch.FloatTensor:
	return self.fuse(self.conv(x) + self.skip(x))


	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) -> 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)
	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,
	) -> 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,
	)[0]
	hidden_states = torch.utils.checkpoint.checkpoint(
	create_custom_forward(resnet),
	hidden_states,
	temb,
	**ckpt_kwargs,
	)
	else:
	hidden_states = attn(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	cross_attention_kwargs=cross_attention_kwargs,
	attention_mask=attention_mask,
	encoder_attention_mask=encoder_attention_mask,
	return_dict=False,
	)[0]
	hidden_states = resnet(hidden_states, temb, 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,
	) -> 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,
	)

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

	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,
	attention_head_dim: int = 1,
	output_scale_factor: float = 1.0,
	downsample_padding: int = 1,
	downsample_type: str = "conv",
	):
	super().__init__()
	resnets = []
	attentions = []
	self.downsample_type = downsample_type

	if attention_head_dim is None:
	logger.warn(
	f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
	)
	attention_head_dim = out_channels

	for i in range(num_layers):
	in_channels = in_channels if i == 0 else out_channels
	resnets.append(
	ResnetBlock2D(
	in_channels=in_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	)
	)
	attentions.append(
	Attention(
	out_channels,
	heads=out_channels // attention_head_dim,
	dim_head=attention_head_dim,
	rescale_output_factor=output_scale_factor,
	eps=resnet_eps,
	norm_num_groups=resnet_groups,
	residual_connection=True,
	bias=True,
	upcast_softmax=True,
	_from_deprecated_attn_block=True,
	)
	)

	self.attentions = nn.ModuleList(attentions)
	self.resnets = nn.ModuleList(resnets)

	if downsample_type == "conv":
	self.downsamplers = nn.ModuleList(
	[
	Downsample2D(
	out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
	)
	]
	)
	elif downsample_type == "resnet":
	self.downsamplers = nn.ModuleList(
	[
	ResnetBlock2D(
	in_channels=out_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	down=True,
	)
	]
	)
	else:
	self.downsamplers = None

	def forward(
	self,
	hidden_states: torch.FloatTensor,
	temb: Optional[torch.FloatTensor] = None,
	upsample_size: Optional[int] = None,
	cross_attention_kwargs: Optional[Dict[str, Any]] = None,
	) -> Tuple[torch.FloatTensor, Tuple[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)

	output_states = ()

	for resnet, attn in zip(self.resnets, self.attentions):
	cross_attention_kwargs.update({"scale": lora_scale})
	hidden_states = resnet(hidden_states, temb, scale=lora_scale)
	hidden_states = attn(hidden_states, **cross_attention_kwargs)
	output_states = output_states + (hidden_states,)

	if self.downsamplers is not None:
	for downsampler in self.downsamplers:
	if self.downsample_type == "resnet":
	hidden_states = downsampler(hidden_states, temb=temb, scale=lora_scale)
	else:
	hidden_states = downsampler(hidden_states, scale=lora_scale)

	output_states += (hidden_states,)

	return hidden_states, output_states


	class CrossAttnDownBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	temb_channels: int,
	dropout: float = 0.0,
	num_layers: int = 1,
	transformer_layers_per_block: Union[int, Tuple[int]] = 1,
	resnet_eps: float = 1e-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,
	) -> 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,
	)
	hidden_states = attn(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	cross_attention_kwargs=cross_attention_kwargs,
	attention_mask=attention_mask,
	encoder_attention_mask=encoder_attention_mask,
	return_dict=False,
	)[0]
	else:
	hidden_states = resnet(hidden_states, temb, 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,
	)[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,)

	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
	) -> 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 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: 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=None,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	)
	)

	self.resnets = nn.ModuleList(resnets)

	if add_downsample:
	self.downsamplers = nn.ModuleList(
	[
	Downsample2D(
	out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
	)
	]
	)
	else:
	self.downsamplers = None

	def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
	for resnet in self.resnets:
	hidden_states = resnet(hidden_states, temb=None, scale=scale)

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

	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,
	attention_head_dim: int = 1,
	output_scale_factor: float = 1.0,
	add_downsample: bool = True,
	downsample_padding: int = 1,
	):
	super().__init__()
	resnets = []
	attentions = []

	if attention_head_dim is None:
	logger.warn(
	f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
	)
	attention_head_dim = out_channels

	for i in range(num_layers):
	in_channels = in_channels if i == 0 else out_channels
	resnets.append(
	ResnetBlock2D(
	in_channels=in_channels,
	out_channels=out_channels,
	temb_channels=None,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	)
	)
	attentions.append(
	Attention(
	out_channels,
	heads=out_channels // attention_head_dim,
	dim_head=attention_head_dim,
	rescale_output_factor=output_scale_factor,
	eps=resnet_eps,
	norm_num_groups=resnet_groups,
	residual_connection=True,
	bias=True,
	upcast_softmax=True,
	_from_deprecated_attn_block=True,
	)
	)

	self.attentions = nn.ModuleList(attentions)
	self.resnets = nn.ModuleList(resnets)

	if add_downsample:
	self.downsamplers = nn.ModuleList(
	[
	Downsample2D(
	out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
	)
	]
	)
	else:
	self.downsamplers = None

	def forward(self, hidden_states: torch.FloatTensor, scale: float = 1.0) -> torch.FloatTensor:
	for resnet, attn in zip(self.resnets, self.attentions):
	hidden_states = resnet(hidden_states, temb=None, scale=scale)
	cross_attention_kwargs = {"scale": scale}
	hidden_states = attn(hidden_states, **cross_attention_kwargs)

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

	return hidden_states


	class AttnSkipDownBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	temb_channels: int,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-6,
	resnet_time_scale_shift: str = "default",
	resnet_act_fn: str = "swish",
	resnet_pre_norm: bool = True,
	attention_head_dim: int = 1,
	output_scale_factor: float = np.sqrt(2.0),
	add_downsample: bool = True,
	):
	super().__init__()
	self.attentions = nn.ModuleList([])
	self.resnets = nn.ModuleList([])

	if attention_head_dim is None:
	logger.warn(
	f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `in_channels`: {out_channels}."
	)
	attention_head_dim = out_channels

	for i in range(num_layers):
	in_channels = in_channels if i == 0 else out_channels
	self.resnets.append(
	ResnetBlock2D(
	in_channels=in_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=min(in_channels // 4, 32),
	groups_out=min(out_channels // 4, 32),
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	)
	)
	self.attentions.append(
	Attention(
	out_channels,
	heads=out_channels // attention_head_dim,
	dim_head=attention_head_dim,
	rescale_output_factor=output_scale_factor,
	eps=resnet_eps,
	norm_num_groups=32,
	residual_connection=True,
	bias=True,
	upcast_softmax=True,
	_from_deprecated_attn_block=True,
	)
	)

	if add_downsample:
	self.resnet_down = ResnetBlock2D(
	in_channels=out_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=min(out_channels // 4, 32),
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	use_in_shortcut=True,
	down=True,
	kernel="fir",
	)
	self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])
	self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1))
	else:
	self.resnet_down = None
	self.downsamplers = None
	self.skip_conv = None

	def forward(
	self,
	hidden_states: torch.FloatTensor,
	temb: Optional[torch.FloatTensor] = None,
	skip_sample: Optional[torch.FloatTensor] = None,
	scale: float = 1.0,
	) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...], torch.FloatTensor]:
	output_states = ()

	for resnet, attn in zip(self.resnets, self.attentions):
	hidden_states = resnet(hidden_states, temb, scale=scale)
	cross_attention_kwargs = {"scale": scale}
	hidden_states = attn(hidden_states, **cross_attention_kwargs)
	output_states += (hidden_states,)

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

	hidden_states = self.skip_conv(skip_sample) + hidden_states

	output_states += (hidden_states,)

	return hidden_states, output_states, skip_sample


	class SkipDownBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	temb_channels: int,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-6,
	resnet_time_scale_shift: str = "default",
	resnet_act_fn: str = "swish",
	resnet_pre_norm: bool = True,
	output_scale_factor: float = np.sqrt(2.0),
	add_downsample: bool = True,
	downsample_padding: int = 1,
	):
	super().__init__()
	self.resnets = nn.ModuleList([])

	for i in range(num_layers):
	in_channels = in_channels if i == 0 else out_channels
	self.resnets.append(
	ResnetBlock2D(
	in_channels=in_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=min(in_channels // 4, 32),
	groups_out=min(out_channels // 4, 32),
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	)
	)

	if add_downsample:
	self.resnet_down = ResnetBlock2D(
	in_channels=out_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=min(out_channels // 4, 32),
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	use_in_shortcut=True,
	down=True,
	kernel="fir",
	)
	self.downsamplers = nn.ModuleList([FirDownsample2D(out_channels, out_channels=out_channels)])
	self.skip_conv = nn.Conv2d(3, out_channels, kernel_size=(1, 1), stride=(1, 1))
	else:
	self.resnet_down = None
	self.downsamplers = None
	self.skip_conv = None

	def forward(
	self,
	hidden_states: torch.FloatTensor,
	temb: Optional[torch.FloatTensor] = None,
	skip_sample: Optional[torch.FloatTensor] = None,
	scale: float = 1.0,
	) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...], torch.FloatTensor]:
	output_states = ()

	for resnet in self.resnets:
	hidden_states = resnet(hidden_states, temb, scale)
	output_states += (hidden_states,)

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

	hidden_states = self.skip_conv(skip_sample) + hidden_states

	output_states += (hidden_states,)

	return hidden_states, output_states, skip_sample


	class ResnetDownsampleBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	temb_channels: int,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-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,
	skip_time_act: bool = False,
	):
	super().__init__()
	resnets = []

	for i in range(num_layers):
	in_channels = in_channels if i == 0 else out_channels
	resnets.append(
	ResnetBlock2D(
	in_channels=in_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	skip_time_act=skip_time_act,
	)
	)

	self.resnets = nn.ModuleList(resnets)

	if add_downsample:
	self.downsamplers = nn.ModuleList(
	[
	ResnetBlock2D(
	in_channels=out_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	skip_time_act=skip_time_act,
	down=True,
	)
	]
	)
	else:
	self.downsamplers = None

	self.gradient_checkpointing = False

	def forward(
	self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
	) -> 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)

	output_states = output_states + (hidden_states,)

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

	output_states = output_states + (hidden_states,)

	return hidden_states, output_states


	class SimpleCrossAttnDownBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	temb_channels: int,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-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,
	cross_attention_dim: int = 1280,
	output_scale_factor: float = 1.0,
	add_downsample: bool = True,
	skip_time_act: bool = False,
	only_cross_attention: bool = False,
	cross_attention_norm: Optional[str] = None,
	):
	super().__init__()

	self.has_cross_attention = True

	resnets = []
	attentions = []

	self.attention_head_dim = attention_head_dim
	self.num_heads = out_channels // self.attention_head_dim

	for i in range(num_layers):
	in_channels = in_channels if i == 0 else out_channels
	resnets.append(
	ResnetBlock2D(
	in_channels=in_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	skip_time_act=skip_time_act,
	)
	)

	processor = (
	AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor()
	)

	attentions.append(
	Attention(
	query_dim=out_channels,
	cross_attention_dim=out_channels,
	heads=self.num_heads,
	dim_head=attention_head_dim,
	added_kv_proj_dim=cross_attention_dim,
	norm_num_groups=resnet_groups,
	bias=True,
	upcast_softmax=True,
	only_cross_attention=only_cross_attention,
	cross_attention_norm=cross_attention_norm,
	processor=processor,
	)
	)
	self.attentions = nn.ModuleList(attentions)
	self.resnets = nn.ModuleList(resnets)

	if add_downsample:
	self.downsamplers = nn.ModuleList(
	[
	ResnetBlock2D(
	in_channels=out_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	skip_time_act=skip_time_act,
	down=True,
	)
	]
	)
	else:
	self.downsamplers = None

	self.gradient_checkpointing = False

	def forward(
	self,
	hidden_states: torch.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,
	) -> Tuple[torch.FloatTensor, Tuple[torch.FloatTensor, ...]]:
	output_states = ()
	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

	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 = attn(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	attention_mask=mask,
	**cross_attention_kwargs,
	)
	else:
	hidden_states = resnet(hidden_states, temb, scale=lora_scale)

	hidden_states = attn(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	attention_mask=mask,
	**cross_attention_kwargs,
	)

	output_states = output_states + (hidden_states,)

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

	output_states = output_states + (hidden_states,)

	return hidden_states, output_states


	class KDownBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	temb_channels: int,
	dropout: float = 0.0,
	num_layers: int = 4,
	resnet_eps: float = 1e-5,
	resnet_act_fn: str = "gelu",
	resnet_group_size: int = 32,
	add_downsample: bool = False,
	):
	super().__init__()
	resnets = []

	for i in range(num_layers):
	in_channels = in_channels if i == 0 else out_channels
	groups = in_channels // resnet_group_size
	groups_out = out_channels // resnet_group_size

	resnets.append(
	ResnetBlock2D(
	in_channels=in_channels,
	out_channels=out_channels,
	dropout=dropout,
	temb_channels=temb_channels,
	groups=groups,
	groups_out=groups_out,
	eps=resnet_eps,
	non_linearity=resnet_act_fn,
	time_embedding_norm="ada_group",
	conv_shortcut_bias=False,
	)
	)

	self.resnets = nn.ModuleList(resnets)

	if add_downsample:
	# YiYi's comments- might be able to use FirDownsample2D, look into details later
	self.downsamplers = nn.ModuleList([KDownsample2D()])
	else:
	self.downsamplers = None

	self.gradient_checkpointing = False

	def forward(
	self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
	) -> 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)

	output_states += (hidden_states,)

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

	return hidden_states, output_states


	class KCrossAttnDownBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	temb_channels: int,
	cross_attention_dim: int,
	dropout: float = 0.0,
	num_layers: int = 4,
	resnet_group_size: int = 32,
	add_downsample: bool = True,
	attention_head_dim: int = 64,
	add_self_attention: bool = False,
	resnet_eps: float = 1e-5,
	resnet_act_fn: str = "gelu",
	):
	super().__init__()
	resnets = []
	attentions = []

	self.has_cross_attention = True

	for i in range(num_layers):
	in_channels = in_channels if i == 0 else out_channels
	groups = in_channels // resnet_group_size
	groups_out = out_channels // resnet_group_size

	resnets.append(
	ResnetBlock2D(
	in_channels=in_channels,
	out_channels=out_channels,
	dropout=dropout,
	temb_channels=temb_channels,
	groups=groups,
	groups_out=groups_out,
	eps=resnet_eps,
	non_linearity=resnet_act_fn,
	time_embedding_norm="ada_group",
	conv_shortcut_bias=False,
	)
	)
	attentions.append(
	KAttentionBlock(
	out_channels,
	out_channels // attention_head_dim,
	attention_head_dim,
	cross_attention_dim=cross_attention_dim,
	temb_channels=temb_channels,
	attention_bias=True,
	add_self_attention=add_self_attention,
	cross_attention_norm="layer_norm",
	group_size=resnet_group_size,
	)
	)

	self.resnets = nn.ModuleList(resnets)
	self.attentions = nn.ModuleList(attentions)

	if add_downsample:
	self.downsamplers = nn.ModuleList([KDownsample2D()])
	else:
	self.downsamplers = None

	self.gradient_checkpointing = False

	def forward(
	self,
	hidden_states: torch.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,
	) -> 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

	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

	ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
	hidden_states = torch.utils.checkpoint.checkpoint(
	create_custom_forward(resnet),
	hidden_states,
	temb,
	**ckpt_kwargs,
	)
	hidden_states = attn(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	emb=temb,
	attention_mask=attention_mask,
	cross_attention_kwargs=cross_attention_kwargs,
	encoder_attention_mask=encoder_attention_mask,
	)
	else:
	hidden_states = resnet(hidden_states, temb, scale=lora_scale)
	hidden_states = attn(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	emb=temb,
	attention_mask=attention_mask,
	cross_attention_kwargs=cross_attention_kwargs,
	encoder_attention_mask=encoder_attention_mask,
	)

	if self.downsamplers is None:
	output_states += (None,)
	else:
	output_states += (hidden_states,)

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

	return hidden_states, output_states


	class AttnUpBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	prev_output_channel: int,
	out_channels: int,
	temb_channels: int,
	resolution_idx: int = None,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-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,
	upsample_type: str = "conv",
	):
	super().__init__()
	resnets = []
	attentions = []

	self.upsample_type = upsample_type

	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`: {out_channels}."
	)
	attention_head_dim = out_channels

	for i in range(num_layers):
	res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
	resnet_in_channels = prev_output_channel if i == 0 else out_channels

	resnets.append(
	ResnetBlock2D(
	in_channels=resnet_in_channels + res_skip_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	)
	)
	attentions.append(
	Attention(
	out_channels,
	heads=out_channels // attention_head_dim,
	dim_head=attention_head_dim,
	rescale_output_factor=output_scale_factor,
	eps=resnet_eps,
	norm_num_groups=resnet_groups,
	residual_connection=True,
	bias=True,
	upcast_softmax=True,
	_from_deprecated_attn_block=True,
	)
	)

	self.attentions = nn.ModuleList(attentions)
	self.resnets = nn.ModuleList(resnets)

	if upsample_type == "conv":
	self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
	elif upsample_type == "resnet":
	self.upsamplers = nn.ModuleList(
	[
	ResnetBlock2D(
	in_channels=out_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	up=True,
	)
	]
	)
	else:
	self.upsamplers = None

	self.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,
	) -> torch.FloatTensor:
	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)

	hidden_states = resnet(hidden_states, temb, scale=scale)
	cross_attention_kwargs = {"scale": scale}
	hidden_states = attn(hidden_states, **cross_attention_kwargs)

	if self.upsamplers is not None:
	for upsampler in self.upsamplers:
	if self.upsample_type == "resnet":
	hidden_states = upsampler(hidden_states, temb=temb, scale=scale)
	else:
	hidden_states = upsampler(hidden_states, scale=scale)

	return hidden_states


	class CrossAttnUpBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	prev_output_channel: int,
	temb_channels: int,
	resolution_idx: Optional[int] = None,
	dropout: float = 0.0,
	num_layers: int = 1,
	transformer_layers_per_block: Union[int, Tuple[int]] = 1,
	resnet_eps: float = 1e-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,
	) -> 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,
	)[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,
	)[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,
	) -> 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


	class UpDecoderBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	resolution_idx: Optional[int] = None,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-6,
	resnet_time_scale_shift: str = "default", # default, spatial
	resnet_act_fn: str = "swish",
	resnet_groups: int = 32,
	resnet_pre_norm: bool = True,
	output_scale_factor: float = 1.0,
	add_upsample: bool = True,
	temb_channels: Optional[int] = None,
	):
	super().__init__()
	resnets = []

	for i in range(num_layers):
	input_channels = in_channels if i == 0 else out_channels

	resnets.append(
	ResnetBlock2D(
	in_channels=input_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	)
	)

	self.resnets = nn.ModuleList(resnets)

	if add_upsample:
	self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
	else:
	self.upsamplers = None

	self.resolution_idx = resolution_idx

	def forward(
	self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
	) -> torch.FloatTensor:
	for resnet in self.resnets:
	hidden_states = resnet(hidden_states, temb=temb, scale=scale)

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

	return hidden_states


	class AttnUpDecoderBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	resolution_idx: Optional[int] = None,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-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,
	add_upsample: bool = True,
	temb_channels: Optional[int] = None,
	):
	super().__init__()
	resnets = []
	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 `out_channels`: {out_channels}."
	)
	attention_head_dim = out_channels

	for i in range(num_layers):
	input_channels = in_channels if i == 0 else out_channels

	resnets.append(
	ResnetBlock2D(
	in_channels=input_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	)
	)
	attentions.append(
	Attention(
	out_channels,
	heads=out_channels // attention_head_dim,
	dim_head=attention_head_dim,
	rescale_output_factor=output_scale_factor,
	eps=resnet_eps,
	norm_num_groups=resnet_groups if resnet_time_scale_shift != "spatial" else None,
	spatial_norm_dim=temb_channels if resnet_time_scale_shift == "spatial" else None,
	residual_connection=True,
	bias=True,
	upcast_softmax=True,
	_from_deprecated_attn_block=True,
	)
	)

	self.attentions = nn.ModuleList(attentions)
	self.resnets = nn.ModuleList(resnets)

	if add_upsample:
	self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
	else:
	self.upsamplers = None

	self.resolution_idx = resolution_idx

	def forward(
	self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, scale: float = 1.0
	) -> torch.FloatTensor:
	for resnet, attn in zip(self.resnets, self.attentions):
	hidden_states = resnet(hidden_states, temb=temb, scale=scale)
	cross_attention_kwargs = {"scale": scale}
	hidden_states = attn(hidden_states, temb=temb, **cross_attention_kwargs)

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

	return hidden_states


	class AttnSkipUpBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	prev_output_channel: int,
	out_channels: int,
	temb_channels: int,
	resolution_idx: Optional[int] = None,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-6,
	resnet_time_scale_shift: str = "default",
	resnet_act_fn: str = "swish",
	resnet_pre_norm: bool = True,
	attention_head_dim: int = 1,
	output_scale_factor: float = np.sqrt(2.0),
	add_upsample: bool = True,
	):
	super().__init__()
	self.attentions = nn.ModuleList([])
	self.resnets = nn.ModuleList([])

	for i in range(num_layers):
	res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
	resnet_in_channels = prev_output_channel if i == 0 else out_channels

	self.resnets.append(
	ResnetBlock2D(
	in_channels=resnet_in_channels + res_skip_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=min(resnet_in_channels + res_skip_channels // 4, 32),
	groups_out=min(out_channels // 4, 32),
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	)
	)

	if attention_head_dim is None:
	logger.warn(
	f"It is not recommend to pass `attention_head_dim=None`. Defaulting `attention_head_dim` to `out_channels`: {out_channels}."
	)
	attention_head_dim = out_channels

	self.attentions.append(
	Attention(
	out_channels,
	heads=out_channels // attention_head_dim,
	dim_head=attention_head_dim,
	rescale_output_factor=output_scale_factor,
	eps=resnet_eps,
	norm_num_groups=32,
	residual_connection=True,
	bias=True,
	upcast_softmax=True,
	_from_deprecated_attn_block=True,
	)
	)

	self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels)
	if add_upsample:
	self.resnet_up = ResnetBlock2D(
	in_channels=out_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=min(out_channels // 4, 32),
	groups_out=min(out_channels // 4, 32),
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	use_in_shortcut=True,
	up=True,
	kernel="fir",
	)
	self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
	self.skip_norm = torch.nn.GroupNorm(
	num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True
	)
	self.act = nn.SiLU()
	else:
	self.resnet_up = None
	self.skip_conv = None
	self.skip_norm = None
	self.act = None

	self.resolution_idx = resolution_idx

	def forward(
	self,
	hidden_states: torch.FloatTensor,
	res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
	temb: Optional[torch.FloatTensor] = None,
	skip_sample=None,
	scale: float = 1.0,
	) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
	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]
	hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

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

	cross_attention_kwargs = {"scale": scale}
	hidden_states = self.attentions[0](hidden_states, **cross_attention_kwargs)

	if skip_sample is not None:
	skip_sample = self.upsampler(skip_sample)
	else:
	skip_sample = 0

	if self.resnet_up is not None:
	skip_sample_states = self.skip_norm(hidden_states)
	skip_sample_states = self.act(skip_sample_states)
	skip_sample_states = self.skip_conv(skip_sample_states)

	skip_sample = skip_sample + skip_sample_states

	hidden_states = self.resnet_up(hidden_states, temb, scale=scale)

	return hidden_states, skip_sample


	class SkipUpBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	prev_output_channel: int,
	out_channels: int,
	temb_channels: int,
	resolution_idx: Optional[int] = None,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-6,
	resnet_time_scale_shift: str = "default",
	resnet_act_fn: str = "swish",
	resnet_pre_norm: bool = True,
	output_scale_factor: float = np.sqrt(2.0),
	add_upsample: bool = True,
	upsample_padding: int = 1,
	):
	super().__init__()
	self.resnets = nn.ModuleList([])

	for i in range(num_layers):
	res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
	resnet_in_channels = prev_output_channel if i == 0 else out_channels

	self.resnets.append(
	ResnetBlock2D(
	in_channels=resnet_in_channels + res_skip_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=min((resnet_in_channels + res_skip_channels) // 4, 32),
	groups_out=min(out_channels // 4, 32),
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	)
	)

	self.upsampler = FirUpsample2D(in_channels, out_channels=out_channels)
	if add_upsample:
	self.resnet_up = ResnetBlock2D(
	in_channels=out_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=min(out_channels // 4, 32),
	groups_out=min(out_channels // 4, 32),
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	use_in_shortcut=True,
	up=True,
	kernel="fir",
	)
	self.skip_conv = nn.Conv2d(out_channels, 3, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1))
	self.skip_norm = torch.nn.GroupNorm(
	num_groups=min(out_channels // 4, 32), num_channels=out_channels, eps=resnet_eps, affine=True
	)
	self.act = nn.SiLU()
	else:
	self.resnet_up = None
	self.skip_conv = None
	self.skip_norm = None
	self.act = None

	self.resolution_idx = resolution_idx

	def forward(
	self,
	hidden_states: torch.FloatTensor,
	res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
	temb: Optional[torch.FloatTensor] = None,
	skip_sample=None,
	scale: float = 1.0,
	) -> Tuple[torch.FloatTensor, torch.FloatTensor]:
	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]
	hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)

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

	if skip_sample is not None:
	skip_sample = self.upsampler(skip_sample)
	else:
	skip_sample = 0

	if self.resnet_up is not None:
	skip_sample_states = self.skip_norm(hidden_states)
	skip_sample_states = self.act(skip_sample_states)
	skip_sample_states = self.skip_conv(skip_sample_states)

	skip_sample = skip_sample + skip_sample_states

	hidden_states = self.resnet_up(hidden_states, temb, scale=scale)

	return hidden_states, skip_sample


	class ResnetUpsampleBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	prev_output_channel: int,
	out_channels: int,
	temb_channels: int,
	resolution_idx: Optional[int] = None,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-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,
	skip_time_act: bool = False,
	):
	super().__init__()
	resnets = []

	for i in range(num_layers):
	res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
	resnet_in_channels = prev_output_channel if i == 0 else out_channels

	resnets.append(
	ResnetBlock2D(
	in_channels=resnet_in_channels + res_skip_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	skip_time_act=skip_time_act,
	)
	)

	self.resnets = nn.ModuleList(resnets)

	if add_upsample:
	self.upsamplers = nn.ModuleList(
	[
	ResnetBlock2D(
	in_channels=out_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	skip_time_act=skip_time_act,
	up=True,
	)
	]
	)
	else:
	self.upsamplers = None

	self.gradient_checkpointing = False
	self.resolution_idx = resolution_idx

	def forward(
	self,
	hidden_states: torch.FloatTensor,
	res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
	temb: Optional[torch.FloatTensor] = None,
	upsample_size: Optional[int] = None,
	scale: float = 1.0,
	) -> torch.FloatTensor:
	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]
	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, temb, scale=scale)

	return hidden_states


	class SimpleCrossAttnUpBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	prev_output_channel: int,
	temb_channels: int,
	resolution_idx: Optional[int] = None,
	dropout: float = 0.0,
	num_layers: int = 1,
	resnet_eps: float = 1e-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,
	cross_attention_dim: int = 1280,
	output_scale_factor: float = 1.0,
	add_upsample: bool = True,
	skip_time_act: bool = False,
	only_cross_attention: bool = False,
	cross_attention_norm: Optional[str] = None,
	):
	super().__init__()
	resnets = []
	attentions = []

	self.has_cross_attention = True
	self.attention_head_dim = attention_head_dim

	self.num_heads = out_channels // self.attention_head_dim

	for i in range(num_layers):
	res_skip_channels = in_channels if (i == num_layers - 1) else out_channels
	resnet_in_channels = prev_output_channel if i == 0 else out_channels

	resnets.append(
	ResnetBlock2D(
	in_channels=resnet_in_channels + res_skip_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	skip_time_act=skip_time_act,
	)
	)

	processor = (
	AttnAddedKVProcessor2_0() if hasattr(F, "scaled_dot_product_attention") else AttnAddedKVProcessor()
	)

	attentions.append(
	Attention(
	query_dim=out_channels,
	cross_attention_dim=out_channels,
	heads=self.num_heads,
	dim_head=self.attention_head_dim,
	added_kv_proj_dim=cross_attention_dim,
	norm_num_groups=resnet_groups,
	bias=True,
	upcast_softmax=True,
	only_cross_attention=only_cross_attention,
	cross_attention_norm=cross_attention_norm,
	processor=processor,
	)
	)
	self.attentions = nn.ModuleList(attentions)
	self.resnets = nn.ModuleList(resnets)

	if add_upsample:
	self.upsamplers = nn.ModuleList(
	[
	ResnetBlock2D(
	in_channels=out_channels,
	out_channels=out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=resnet_groups,
	dropout=dropout,
	time_embedding_norm=resnet_time_scale_shift,
	non_linearity=resnet_act_fn,
	output_scale_factor=output_scale_factor,
	pre_norm=resnet_pre_norm,
	skip_time_act=skip_time_act,
	up=True,
	)
	]
	)
	else:
	self.upsamplers = None

	self.gradient_checkpointing = False
	self.resolution_idx = resolution_idx

	def forward(
	self,
	hidden_states: torch.FloatTensor,
	res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
	temb: Optional[torch.FloatTensor] = None,
	encoder_hidden_states: Optional[torch.FloatTensor] = None,
	upsample_size: Optional[int] = None,
	attention_mask: Optional[torch.FloatTensor] = None,
	cross_attention_kwargs: Optional[Dict[str, Any]] = None,
	encoder_attention_mask: Optional[torch.FloatTensor] = None,
	) -> 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

	for resnet, attn in zip(self.resnets, self.attentions):
	# resnet
	# 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

	hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
	hidden_states = attn(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	attention_mask=mask,
	**cross_attention_kwargs,
	)
	else:
	hidden_states = resnet(hidden_states, temb, scale=lora_scale)

	hidden_states = attn(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	attention_mask=mask,
	**cross_attention_kwargs,
	)

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

	return hidden_states


	class KUpBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	temb_channels: int,
	resolution_idx: int,
	dropout: float = 0.0,
	num_layers: int = 5,
	resnet_eps: float = 1e-5,
	resnet_act_fn: str = "gelu",
	resnet_group_size: Optional[int] = 32,
	add_upsample: bool = True,
	):
	super().__init__()
	resnets = []
	k_in_channels = 2 * out_channels
	k_out_channels = in_channels
	num_layers = num_layers - 1

	for i in range(num_layers):
	in_channels = k_in_channels if i == 0 else out_channels
	groups = in_channels // resnet_group_size
	groups_out = out_channels // resnet_group_size

	resnets.append(
	ResnetBlock2D(
	in_channels=in_channels,
	out_channels=k_out_channels if (i == num_layers - 1) else out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=groups,
	groups_out=groups_out,
	dropout=dropout,
	non_linearity=resnet_act_fn,
	time_embedding_norm="ada_group",
	conv_shortcut_bias=False,
	)
	)

	self.resnets = nn.ModuleList(resnets)

	if add_upsample:
	self.upsamplers = nn.ModuleList([KUpsample2D()])
	else:
	self.upsamplers = None

	self.gradient_checkpointing = False
	self.resolution_idx = resolution_idx

	def forward(
	self,
	hidden_states: torch.FloatTensor,
	res_hidden_states_tuple: Tuple[torch.FloatTensor, ...],
	temb: Optional[torch.FloatTensor] = None,
	upsample_size: Optional[int] = None,
	scale: float = 1.0,
	) -> torch.FloatTensor:
	res_hidden_states_tuple = res_hidden_states_tuple[-1]
	if res_hidden_states_tuple is not None:
	hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1)

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

	return hidden_states


	class KCrossAttnUpBlock2D(nn.Module):
	def __init__(
	self,
	in_channels: int,
	out_channels: int,
	temb_channels: int,
	resolution_idx: int,
	dropout: float = 0.0,
	num_layers: int = 4,
	resnet_eps: float = 1e-5,
	resnet_act_fn: str = "gelu",
	resnet_group_size: int = 32,
	attention_head_dim: int = 1, # attention dim_head
	cross_attention_dim: int = 768,
	add_upsample: bool = True,
	upcast_attention: bool = False,
	):
	super().__init__()
	resnets = []
	attentions = []

	is_first_block = in_channels == out_channels == temb_channels
	is_middle_block = in_channels != out_channels
	add_self_attention = True if is_first_block else False

	self.has_cross_attention = True
	self.attention_head_dim = attention_head_dim

	# in_channels, and out_channels for the block (k-unet)
	k_in_channels = out_channels if is_first_block else 2 * out_channels
	k_out_channels = in_channels

	num_layers = num_layers - 1

	for i in range(num_layers):
	in_channels = k_in_channels if i == 0 else out_channels
	groups = in_channels // resnet_group_size
	groups_out = out_channels // resnet_group_size

	if is_middle_block and (i == num_layers - 1):
	conv_2d_out_channels = k_out_channels
	else:
	conv_2d_out_channels = None

	resnets.append(
	ResnetBlock2D(
	in_channels=in_channels,
	out_channels=out_channels,
	conv_2d_out_channels=conv_2d_out_channels,
	temb_channels=temb_channels,
	eps=resnet_eps,
	groups=groups,
	groups_out=groups_out,
	dropout=dropout,
	non_linearity=resnet_act_fn,
	time_embedding_norm="ada_group",
	conv_shortcut_bias=False,
	)
	)
	attentions.append(
	KAttentionBlock(
	k_out_channels if (i == num_layers - 1) else out_channels,
	k_out_channels // attention_head_dim
	if (i == num_layers - 1)
	else out_channels // attention_head_dim,
	attention_head_dim,
	cross_attention_dim=cross_attention_dim,
	temb_channels=temb_channels,
	attention_bias=True,
	add_self_attention=add_self_attention,
	cross_attention_norm="layer_norm",
	upcast_attention=upcast_attention,
	)
	)

	self.resnets = nn.ModuleList(resnets)
	self.attentions = nn.ModuleList(attentions)

	if add_upsample:
	self.upsamplers = nn.ModuleList([KUpsample2D()])
	else:
	self.upsamplers = None

	self.gradient_checkpointing = False
	self.resolution_idx = resolution_idx

	def forward(
	self,
	hidden_states: torch.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,
	) -> torch.FloatTensor:
	res_hidden_states_tuple = res_hidden_states_tuple[-1]
	if res_hidden_states_tuple is not None:
	hidden_states = torch.cat([hidden_states, res_hidden_states_tuple], dim=1)

	lora_scale = cross_attention_kwargs.get("scale", 1.0) if cross_attention_kwargs is not None else 1.0
	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

	ckpt_kwargs: Dict[str, Any] = {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {}
	hidden_states = torch.utils.checkpoint.checkpoint(
	create_custom_forward(resnet),
	hidden_states,
	temb,
	**ckpt_kwargs,
	)
	hidden_states = attn(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	emb=temb,
	attention_mask=attention_mask,
	cross_attention_kwargs=cross_attention_kwargs,
	encoder_attention_mask=encoder_attention_mask,
	)
	else:
	hidden_states = resnet(hidden_states, temb, scale=lora_scale)
	hidden_states = attn(
	hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	emb=temb,
	attention_mask=attention_mask,
	cross_attention_kwargs=cross_attention_kwargs,
	encoder_attention_mask=encoder_attention_mask,
	)

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

	return hidden_states


	# can potentially later be renamed to `No-feed-forward` attention
	class KAttentionBlock(nn.Module):
	r"""
	A basic Transformer block.

	Parameters:
	dim (`int`): The number of channels in the input and output.
	num_attention_heads (`int`): The number of heads to use for multi-head attention.
	attention_head_dim (`int`): The number of channels in each head.
	dropout (`float`, optional, defaults to 0.0): The dropout probability to use.
	cross_attention_dim (`int`, optional): The size of the encoder_hidden_states vector for cross attention.
	attention_bias (`bool`, optional, defaults to `False`):
	Configure if the attention layers should contain a bias parameter.
	upcast_attention (`bool`, optional, defaults to `False`):
	Set to `True` to upcast the attention computation to `float32`.
	temb_channels (`int`, optional, defaults to 768):
	The number of channels in the token embedding.
	add_self_attention (`bool`, optional, defaults to `False`):
	Set to `True` to add self-attention to the block.
	cross_attention_norm (`str`, optional, defaults to `None`):
	The type of normalization to use for the cross attention. Can be `None`, `layer_norm`, or `group_norm`.
	group_size (`int`, optional, defaults to 32):
	The number of groups to separate the channels into for group normalization.
	"""

	def __init__(
	self,
	dim: int,
	num_attention_heads: int,
	attention_head_dim: int,
	dropout: float = 0.0,
	cross_attention_dim: Optional[int] = None,
	attention_bias: bool = False,
	upcast_attention: bool = False,
	temb_channels: int = 768, # for ada_group_norm
	add_self_attention: bool = False,
	cross_attention_norm: Optional[str] = None,
	group_size: int = 32,
	):
	super().__init__()
	self.add_self_attention = add_self_attention

	# 1. Self-Attn
	if add_self_attention:
	self.norm1 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size))
	self.attn1 = Attention(
	query_dim=dim,
	heads=num_attention_heads,
	dim_head=attention_head_dim,
	dropout=dropout,
	bias=attention_bias,
	cross_attention_dim=None,
	cross_attention_norm=None,
	)

	# 2. Cross-Attn
	self.norm2 = AdaGroupNorm(temb_channels, dim, max(1, dim // group_size))
	self.attn2 = Attention(
	query_dim=dim,
	cross_attention_dim=cross_attention_dim,
	heads=num_attention_heads,
	dim_head=attention_head_dim,
	dropout=dropout,
	bias=attention_bias,
	upcast_attention=upcast_attention,
	cross_attention_norm=cross_attention_norm,
	)

	def _to_3d(self, hidden_states: torch.FloatTensor, height: int, weight: int) -> torch.FloatTensor:
	return hidden_states.permute(0, 2, 3, 1).reshape(hidden_states.shape[0], height * weight, -1)

	def _to_4d(self, hidden_states: torch.FloatTensor, height: int, weight: int) -> torch.FloatTensor:
	return hidden_states.permute(0, 2, 1).reshape(hidden_states.shape[0], -1, height, weight)

	def forward(
	self,
	hidden_states: torch.FloatTensor,
	encoder_hidden_states: Optional[torch.FloatTensor] = None,
	# TODO: mark emb as non-optional (self.norm2 requires it).
	# requires assessing impact of change to positional param interface.
	emb: 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,
	) -> torch.FloatTensor:
	cross_attention_kwargs = cross_attention_kwargs if cross_attention_kwargs is not None else {}

	# 1. Self-Attention
	if self.add_self_attention:
	norm_hidden_states = self.norm1(hidden_states, emb)

	height, weight = norm_hidden_states.shape[2:]
	norm_hidden_states = self._to_3d(norm_hidden_states, height, weight)

	attn_output = self.attn1(
	norm_hidden_states,
	encoder_hidden_states=None,
	attention_mask=attention_mask,
	**cross_attention_kwargs,
	)
	attn_output = self._to_4d(attn_output, height, weight)

	hidden_states = attn_output + hidden_states

	# 2. Cross-Attention/None
	norm_hidden_states = self.norm2(hidden_states, emb)

	height, weight = norm_hidden_states.shape[2:]
	norm_hidden_states = self._to_3d(norm_hidden_states, height, weight)
	attn_output = self.attn2(
	norm_hidden_states,
	encoder_hidden_states=encoder_hidden_states,
	attention_mask=attention_mask if encoder_hidden_states is None else encoder_attention_mask,
	**cross_attention_kwargs,
	)
	attn_output = self._to_4d(attn_output, height, weight)

	hidden_states = attn_output + hidden_states

	return hidden_states