2178 lines
77 KiB
Python
2178 lines
77 KiB
Python
# Copyright 2022 The HuggingFace Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
import numpy as np
|
|
import torch
|
|
from torch import nn
|
|
|
|
from .attention import AttentionBlock, CrossAttention, DualTransformer2DModel, Transformer2DModel
|
|
from .resnet import Downsample2D, FirDownsample2D, FirUpsample2D, ResnetBlock2D, Upsample2D
|
|
|
|
|
|
def get_down_block(
|
|
down_block_type,
|
|
num_layers,
|
|
in_channels,
|
|
out_channels,
|
|
temb_channels,
|
|
add_downsample,
|
|
resnet_eps,
|
|
resnet_act_fn,
|
|
attn_num_head_channels,
|
|
resnet_groups=None,
|
|
cross_attention_dim=None,
|
|
downsample_padding=None,
|
|
dual_cross_attention=False,
|
|
use_linear_projection=False,
|
|
only_cross_attention=False,
|
|
upcast_attention=False,
|
|
resnet_time_scale_shift="default",
|
|
):
|
|
down_block_type = down_block_type[7:] if down_block_type.startswith("UNetRes") else down_block_type
|
|
if down_block_type == "DownBlock2D":
|
|
return DownBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
add_downsample=add_downsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
downsample_padding=downsample_padding,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif down_block_type == "ResnetDownsampleBlock2D":
|
|
return ResnetDownsampleBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
add_downsample=add_downsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif down_block_type == "AttnDownBlock2D":
|
|
return AttnDownBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
add_downsample=add_downsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
downsample_padding=downsample_padding,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif down_block_type == "CrossAttnDownBlock2D":
|
|
if cross_attention_dim is None:
|
|
raise ValueError("cross_attention_dim must be specified for CrossAttnDownBlock2D")
|
|
return CrossAttnDownBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
add_downsample=add_downsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
downsample_padding=downsample_padding,
|
|
cross_attention_dim=cross_attention_dim,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
dual_cross_attention=dual_cross_attention,
|
|
use_linear_projection=use_linear_projection,
|
|
only_cross_attention=only_cross_attention,
|
|
upcast_attention=upcast_attention,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif down_block_type == "SimpleCrossAttnDownBlock2D":
|
|
if cross_attention_dim is None:
|
|
raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnDownBlock2D")
|
|
return SimpleCrossAttnDownBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
add_downsample=add_downsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
cross_attention_dim=cross_attention_dim,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif down_block_type == "SkipDownBlock2D":
|
|
return SkipDownBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
add_downsample=add_downsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
downsample_padding=downsample_padding,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif down_block_type == "AttnSkipDownBlock2D":
|
|
return AttnSkipDownBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
add_downsample=add_downsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
downsample_padding=downsample_padding,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif down_block_type == "DownEncoderBlock2D":
|
|
return DownEncoderBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_downsample=add_downsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
downsample_padding=downsample_padding,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif down_block_type == "AttnDownEncoderBlock2D":
|
|
return AttnDownEncoderBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_downsample=add_downsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
downsample_padding=downsample_padding,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
raise ValueError(f"{down_block_type} does not exist.")
|
|
|
|
|
|
def get_up_block(
|
|
up_block_type,
|
|
num_layers,
|
|
in_channels,
|
|
out_channels,
|
|
prev_output_channel,
|
|
temb_channels,
|
|
add_upsample,
|
|
resnet_eps,
|
|
resnet_act_fn,
|
|
attn_num_head_channels,
|
|
resnet_groups=None,
|
|
cross_attention_dim=None,
|
|
dual_cross_attention=False,
|
|
use_linear_projection=False,
|
|
only_cross_attention=False,
|
|
upcast_attention=False,
|
|
resnet_time_scale_shift="default",
|
|
):
|
|
up_block_type = up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type
|
|
if up_block_type == "UpBlock2D":
|
|
return UpBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
prev_output_channel=prev_output_channel,
|
|
temb_channels=temb_channels,
|
|
add_upsample=add_upsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif up_block_type == "ResnetUpsampleBlock2D":
|
|
return ResnetUpsampleBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
prev_output_channel=prev_output_channel,
|
|
temb_channels=temb_channels,
|
|
add_upsample=add_upsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif up_block_type == "CrossAttnUpBlock2D":
|
|
if cross_attention_dim is None:
|
|
raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D")
|
|
return CrossAttnUpBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
prev_output_channel=prev_output_channel,
|
|
temb_channels=temb_channels,
|
|
add_upsample=add_upsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
cross_attention_dim=cross_attention_dim,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
dual_cross_attention=dual_cross_attention,
|
|
use_linear_projection=use_linear_projection,
|
|
only_cross_attention=only_cross_attention,
|
|
upcast_attention=upcast_attention,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif up_block_type == "SimpleCrossAttnUpBlock2D":
|
|
if cross_attention_dim is None:
|
|
raise ValueError("cross_attention_dim must be specified for SimpleCrossAttnUpBlock2D")
|
|
return SimpleCrossAttnUpBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
prev_output_channel=prev_output_channel,
|
|
temb_channels=temb_channels,
|
|
add_upsample=add_upsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
cross_attention_dim=cross_attention_dim,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif up_block_type == "AttnUpBlock2D":
|
|
return AttnUpBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
prev_output_channel=prev_output_channel,
|
|
temb_channels=temb_channels,
|
|
add_upsample=add_upsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif up_block_type == "SkipUpBlock2D":
|
|
return SkipUpBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
prev_output_channel=prev_output_channel,
|
|
temb_channels=temb_channels,
|
|
add_upsample=add_upsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif up_block_type == "AttnSkipUpBlock2D":
|
|
return AttnSkipUpBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
prev_output_channel=prev_output_channel,
|
|
temb_channels=temb_channels,
|
|
add_upsample=add_upsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif up_block_type == "UpDecoderBlock2D":
|
|
return UpDecoderBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_upsample=add_upsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
elif up_block_type == "AttnUpDecoderBlock2D":
|
|
return AttnUpDecoderBlock2D(
|
|
num_layers=num_layers,
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
add_upsample=add_upsample,
|
|
resnet_eps=resnet_eps,
|
|
resnet_act_fn=resnet_act_fn,
|
|
resnet_groups=resnet_groups,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
resnet_time_scale_shift=resnet_time_scale_shift,
|
|
)
|
|
raise ValueError(f"{up_block_type} does not exist.")
|
|
|
|
|
|
class UNetMidBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
temb_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
add_attention: bool = True,
|
|
attn_num_head_channels=1,
|
|
output_scale_factor=1.0,
|
|
):
|
|
super().__init__()
|
|
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
|
self.add_attention = add_attention
|
|
|
|
# there is always at least one resnet
|
|
resnets = [
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
temb_channels=temb_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
]
|
|
attentions = []
|
|
|
|
for _ in range(num_layers):
|
|
if self.add_attention:
|
|
attentions.append(
|
|
AttentionBlock(
|
|
in_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
norm_num_groups=resnet_groups,
|
|
)
|
|
)
|
|
else:
|
|
attentions.append(None)
|
|
|
|
resnets.append(
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
temb_channels=temb_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
)
|
|
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
def forward(self, hidden_states, temb=None):
|
|
hidden_states = self.resnets[0](hidden_states, temb)
|
|
for attn, resnet in zip(self.attentions, self.resnets[1:]):
|
|
if attn is not None:
|
|
hidden_states = attn(hidden_states)
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class UNetMidBlock2DCrossAttn(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
temb_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
attn_num_head_channels=1,
|
|
output_scale_factor=1.0,
|
|
cross_attention_dim=1280,
|
|
dual_cross_attention=False,
|
|
use_linear_projection=False,
|
|
upcast_attention=False,
|
|
):
|
|
super().__init__()
|
|
|
|
self.has_cross_attention = True
|
|
self.attn_num_head_channels = attn_num_head_channels
|
|
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
|
|
|
# there is always at least one resnet
|
|
resnets = [
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
temb_channels=temb_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
]
|
|
attentions = []
|
|
|
|
for _ in range(num_layers):
|
|
if not dual_cross_attention:
|
|
attentions.append(
|
|
Transformer2DModel(
|
|
attn_num_head_channels,
|
|
in_channels // attn_num_head_channels,
|
|
in_channels=in_channels,
|
|
num_layers=1,
|
|
cross_attention_dim=cross_attention_dim,
|
|
norm_num_groups=resnet_groups,
|
|
use_linear_projection=use_linear_projection,
|
|
upcast_attention=upcast_attention,
|
|
)
|
|
)
|
|
else:
|
|
attentions.append(
|
|
DualTransformer2DModel(
|
|
attn_num_head_channels,
|
|
in_channels // attn_num_head_channels,
|
|
in_channels=in_channels,
|
|
num_layers=1,
|
|
cross_attention_dim=cross_attention_dim,
|
|
norm_num_groups=resnet_groups,
|
|
)
|
|
)
|
|
resnets.append(
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
temb_channels=temb_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
)
|
|
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):
|
|
# TODO(Patrick, William) - attention_mask is currently not used. Implement once used
|
|
hidden_states = self.resnets[0](hidden_states, temb)
|
|
for attn, resnet in zip(self.attentions, self.resnets[1:]):
|
|
hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class UNetMidBlock2DSimpleCrossAttn(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
temb_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
attn_num_head_channels=1,
|
|
output_scale_factor=1.0,
|
|
cross_attention_dim=1280,
|
|
):
|
|
super().__init__()
|
|
|
|
self.has_cross_attention = True
|
|
|
|
self.attn_num_head_channels = attn_num_head_channels
|
|
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
|
|
|
self.num_heads = in_channels // self.attn_num_head_channels
|
|
|
|
# there is always at least one resnet
|
|
resnets = [
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
temb_channels=temb_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
]
|
|
attentions = []
|
|
|
|
for _ in range(num_layers):
|
|
attentions.append(
|
|
CrossAttention(
|
|
query_dim=in_channels,
|
|
cross_attention_dim=in_channels,
|
|
heads=self.num_heads,
|
|
dim_head=attn_num_head_channels,
|
|
added_kv_proj_dim=cross_attention_dim,
|
|
norm_num_groups=resnet_groups,
|
|
bias=True,
|
|
upcast_softmax=True,
|
|
)
|
|
)
|
|
resnets.append(
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
temb_channels=temb_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
)
|
|
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):
|
|
hidden_states = self.resnets[0](hidden_states, temb)
|
|
for attn, resnet in zip(self.attentions, self.resnets[1:]):
|
|
# attn
|
|
residual = hidden_states
|
|
hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)
|
|
hidden_states = attn(
|
|
hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states.transpose(1, 2),
|
|
attention_mask=attention_mask,
|
|
)
|
|
hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)
|
|
hidden_states = hidden_states + residual
|
|
|
|
# resnet
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class AttnDownBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
temb_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
attn_num_head_channels=1,
|
|
output_scale_factor=1.0,
|
|
downsample_padding=1,
|
|
add_downsample=True,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
for i in range(num_layers):
|
|
in_channels = in_channels if i == 0 else out_channels
|
|
resnets.append(
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
)
|
|
attentions.append(
|
|
AttentionBlock(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
norm_num_groups=resnet_groups,
|
|
)
|
|
)
|
|
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_downsample:
|
|
self.downsamplers = nn.ModuleList(
|
|
[
|
|
Downsample2D(
|
|
out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
self.downsamplers = None
|
|
|
|
def forward(self, hidden_states, temb=None):
|
|
output_states = ()
|
|
|
|
for resnet, attn in zip(self.resnets, self.attentions):
|
|
hidden_states = resnet(hidden_states, temb)
|
|
hidden_states = attn(hidden_states)
|
|
output_states += (hidden_states,)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states)
|
|
|
|
output_states += (hidden_states,)
|
|
|
|
return hidden_states, output_states
|
|
|
|
|
|
class CrossAttnDownBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
temb_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
attn_num_head_channels=1,
|
|
cross_attention_dim=1280,
|
|
output_scale_factor=1.0,
|
|
downsample_padding=1,
|
|
add_downsample=True,
|
|
dual_cross_attention=False,
|
|
use_linear_projection=False,
|
|
only_cross_attention=False,
|
|
upcast_attention=False,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
self.has_cross_attention = True
|
|
self.attn_num_head_channels = attn_num_head_channels
|
|
|
|
for i in range(num_layers):
|
|
in_channels = in_channels if i == 0 else out_channels
|
|
resnets.append(
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
)
|
|
if not dual_cross_attention:
|
|
attentions.append(
|
|
Transformer2DModel(
|
|
attn_num_head_channels,
|
|
out_channels // attn_num_head_channels,
|
|
in_channels=out_channels,
|
|
num_layers=1,
|
|
cross_attention_dim=cross_attention_dim,
|
|
norm_num_groups=resnet_groups,
|
|
use_linear_projection=use_linear_projection,
|
|
only_cross_attention=only_cross_attention,
|
|
upcast_attention=upcast_attention,
|
|
)
|
|
)
|
|
else:
|
|
attentions.append(
|
|
DualTransformer2DModel(
|
|
attn_num_head_channels,
|
|
out_channels // attn_num_head_channels,
|
|
in_channels=out_channels,
|
|
num_layers=1,
|
|
cross_attention_dim=cross_attention_dim,
|
|
norm_num_groups=resnet_groups,
|
|
)
|
|
)
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_downsample:
|
|
self.downsamplers = nn.ModuleList(
|
|
[
|
|
Downsample2D(
|
|
out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
self.downsamplers = None
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):
|
|
# TODO(Patrick, William) - attention mask is not used
|
|
output_states = ()
|
|
|
|
for resnet, attn in zip(self.resnets, self.attentions):
|
|
if self.training and self.gradient_checkpointing:
|
|
|
|
def create_custom_forward(module, return_dict=None):
|
|
def custom_forward(*inputs):
|
|
if return_dict is not None:
|
|
return module(*inputs, return_dict=return_dict)
|
|
else:
|
|
return module(*inputs)
|
|
|
|
return custom_forward
|
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
|
|
hidden_states = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(attn, return_dict=False),
|
|
hidden_states,
|
|
encoder_hidden_states,
|
|
)[0]
|
|
else:
|
|
hidden_states = resnet(hidden_states, temb)
|
|
hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
|
|
|
|
output_states += (hidden_states,)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states)
|
|
|
|
output_states += (hidden_states,)
|
|
|
|
return hidden_states, output_states
|
|
|
|
|
|
class DownBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
temb_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
output_scale_factor=1.0,
|
|
add_downsample=True,
|
|
downsample_padding=1,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
|
|
for i in range(num_layers):
|
|
in_channels = in_channels if i == 0 else out_channels
|
|
resnets.append(
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
)
|
|
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_downsample:
|
|
self.downsamplers = nn.ModuleList(
|
|
[
|
|
Downsample2D(
|
|
out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
self.downsamplers = None
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(self, hidden_states, temb=None):
|
|
output_states = ()
|
|
|
|
for resnet in self.resnets:
|
|
if self.training and self.gradient_checkpointing:
|
|
|
|
def create_custom_forward(module):
|
|
def custom_forward(*inputs):
|
|
return module(*inputs)
|
|
|
|
return custom_forward
|
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
|
|
else:
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
output_states += (hidden_states,)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states)
|
|
|
|
output_states += (hidden_states,)
|
|
|
|
return hidden_states, output_states
|
|
|
|
|
|
class DownEncoderBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
output_scale_factor=1.0,
|
|
add_downsample=True,
|
|
downsample_padding=1,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
|
|
for i in range(num_layers):
|
|
in_channels = in_channels if i == 0 else out_channels
|
|
resnets.append(
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=None,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
)
|
|
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_downsample:
|
|
self.downsamplers = nn.ModuleList(
|
|
[
|
|
Downsample2D(
|
|
out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
self.downsamplers = None
|
|
|
|
def forward(self, hidden_states):
|
|
for resnet in self.resnets:
|
|
hidden_states = resnet(hidden_states, temb=None)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class AttnDownEncoderBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
attn_num_head_channels=1,
|
|
output_scale_factor=1.0,
|
|
add_downsample=True,
|
|
downsample_padding=1,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
for i in range(num_layers):
|
|
in_channels = in_channels if i == 0 else out_channels
|
|
resnets.append(
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=None,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
)
|
|
attentions.append(
|
|
AttentionBlock(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
norm_num_groups=resnet_groups,
|
|
)
|
|
)
|
|
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_downsample:
|
|
self.downsamplers = nn.ModuleList(
|
|
[
|
|
Downsample2D(
|
|
out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
self.downsamplers = None
|
|
|
|
def forward(self, hidden_states):
|
|
for resnet, attn in zip(self.resnets, self.attentions):
|
|
hidden_states = resnet(hidden_states, temb=None)
|
|
hidden_states = attn(hidden_states)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
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,
|
|
attn_num_head_channels=1,
|
|
output_scale_factor=np.sqrt(2.0),
|
|
downsample_padding=1,
|
|
add_downsample=True,
|
|
):
|
|
super().__init__()
|
|
self.attentions = nn.ModuleList([])
|
|
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,
|
|
)
|
|
)
|
|
self.attentions.append(
|
|
AttentionBlock(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
)
|
|
)
|
|
|
|
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, temb=None, skip_sample=None):
|
|
output_states = ()
|
|
|
|
for resnet, attn in zip(self.resnets, self.attentions):
|
|
hidden_states = resnet(hidden_states, temb)
|
|
hidden_states = attn(hidden_states)
|
|
output_states += (hidden_states,)
|
|
|
|
if self.downsamplers is not None:
|
|
hidden_states = self.resnet_down(hidden_states, temb)
|
|
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=np.sqrt(2.0),
|
|
add_downsample=True,
|
|
downsample_padding=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, temb=None, skip_sample=None):
|
|
output_states = ()
|
|
|
|
for resnet in self.resnets:
|
|
hidden_states = resnet(hidden_states, temb)
|
|
output_states += (hidden_states,)
|
|
|
|
if self.downsamplers is not None:
|
|
hidden_states = self.resnet_down(hidden_states, temb)
|
|
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=1.0,
|
|
add_downsample=True,
|
|
):
|
|
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(
|
|
[
|
|
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
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(self, hidden_states, temb=None):
|
|
output_states = ()
|
|
|
|
for resnet in self.resnets:
|
|
if self.training and self.gradient_checkpointing:
|
|
|
|
def create_custom_forward(module):
|
|
def custom_forward(*inputs):
|
|
return module(*inputs)
|
|
|
|
return custom_forward
|
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
|
|
else:
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
output_states += (hidden_states,)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states, temb)
|
|
|
|
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,
|
|
attn_num_head_channels=1,
|
|
cross_attention_dim=1280,
|
|
output_scale_factor=1.0,
|
|
add_downsample=True,
|
|
):
|
|
super().__init__()
|
|
|
|
self.has_cross_attention = True
|
|
|
|
resnets = []
|
|
attentions = []
|
|
|
|
self.attn_num_head_channels = attn_num_head_channels
|
|
self.num_heads = out_channels // self.attn_num_head_channels
|
|
|
|
for i in range(num_layers):
|
|
in_channels = in_channels if i == 0 else out_channels
|
|
resnets.append(
|
|
ResnetBlock2D(
|
|
in_channels=in_channels,
|
|
out_channels=out_channels,
|
|
temb_channels=temb_channels,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
)
|
|
attentions.append(
|
|
CrossAttention(
|
|
query_dim=out_channels,
|
|
cross_attention_dim=out_channels,
|
|
heads=self.num_heads,
|
|
dim_head=attn_num_head_channels,
|
|
added_kv_proj_dim=cross_attention_dim,
|
|
norm_num_groups=resnet_groups,
|
|
bias=True,
|
|
upcast_softmax=True,
|
|
)
|
|
)
|
|
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,
|
|
down=True,
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
self.downsamplers = None
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(self, hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None):
|
|
output_states = ()
|
|
|
|
for resnet, attn in zip(self.resnets, self.attentions):
|
|
# resnet
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
# attn
|
|
residual = hidden_states
|
|
hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)
|
|
hidden_states = attn(
|
|
hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states.transpose(1, 2),
|
|
attention_mask=attention_mask,
|
|
)
|
|
hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)
|
|
hidden_states = hidden_states + residual
|
|
|
|
output_states += (hidden_states,)
|
|
|
|
if self.downsamplers is not None:
|
|
for downsampler in self.downsamplers:
|
|
hidden_states = downsampler(hidden_states, temb)
|
|
|
|
output_states += (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,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
attn_num_head_channels=1,
|
|
output_scale_factor=1.0,
|
|
add_upsample=True,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
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(
|
|
AttentionBlock(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
norm_num_groups=resnet_groups,
|
|
)
|
|
)
|
|
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_upsample:
|
|
self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
|
|
else:
|
|
self.upsamplers = None
|
|
|
|
def forward(self, hidden_states, res_hidden_states_tuple, temb=None):
|
|
for resnet, attn in zip(self.resnets, self.attentions):
|
|
# pop res hidden states
|
|
res_hidden_states = res_hidden_states_tuple[-1]
|
|
res_hidden_states_tuple = res_hidden_states_tuple[:-1]
|
|
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
|
|
|
|
hidden_states = resnet(hidden_states, temb)
|
|
hidden_states = attn(hidden_states)
|
|
|
|
if self.upsamplers is not None:
|
|
for upsampler in self.upsamplers:
|
|
hidden_states = upsampler(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class CrossAttnUpBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
prev_output_channel: int,
|
|
temb_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
attn_num_head_channels=1,
|
|
cross_attention_dim=1280,
|
|
output_scale_factor=1.0,
|
|
add_upsample=True,
|
|
dual_cross_attention=False,
|
|
use_linear_projection=False,
|
|
only_cross_attention=False,
|
|
upcast_attention=False,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
self.has_cross_attention = True
|
|
self.attn_num_head_channels = attn_num_head_channels
|
|
|
|
for i in range(num_layers):
|
|
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(
|
|
attn_num_head_channels,
|
|
out_channels // attn_num_head_channels,
|
|
in_channels=out_channels,
|
|
num_layers=1,
|
|
cross_attention_dim=cross_attention_dim,
|
|
norm_num_groups=resnet_groups,
|
|
use_linear_projection=use_linear_projection,
|
|
only_cross_attention=only_cross_attention,
|
|
upcast_attention=upcast_attention,
|
|
)
|
|
)
|
|
else:
|
|
attentions.append(
|
|
DualTransformer2DModel(
|
|
attn_num_head_channels,
|
|
out_channels // attn_num_head_channels,
|
|
in_channels=out_channels,
|
|
num_layers=1,
|
|
cross_attention_dim=cross_attention_dim,
|
|
norm_num_groups=resnet_groups,
|
|
)
|
|
)
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_upsample:
|
|
self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
|
|
else:
|
|
self.upsamplers = None
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states,
|
|
res_hidden_states_tuple,
|
|
temb=None,
|
|
encoder_hidden_states=None,
|
|
upsample_size=None,
|
|
attention_mask=None,
|
|
):
|
|
# TODO(Patrick, William) - attention mask is not used
|
|
for resnet, attn in zip(self.resnets, self.attentions):
|
|
# pop res hidden states
|
|
res_hidden_states = res_hidden_states_tuple[-1]
|
|
res_hidden_states_tuple = res_hidden_states_tuple[:-1]
|
|
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
|
|
|
|
if self.training and self.gradient_checkpointing:
|
|
|
|
def create_custom_forward(module, return_dict=None):
|
|
def custom_forward(*inputs):
|
|
if return_dict is not None:
|
|
return module(*inputs, return_dict=return_dict)
|
|
else:
|
|
return module(*inputs)
|
|
|
|
return custom_forward
|
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
|
|
hidden_states = torch.utils.checkpoint.checkpoint(
|
|
create_custom_forward(attn, return_dict=False),
|
|
hidden_states,
|
|
encoder_hidden_states,
|
|
)[0]
|
|
else:
|
|
hidden_states = resnet(hidden_states, temb)
|
|
hidden_states = attn(hidden_states, encoder_hidden_states=encoder_hidden_states).sample
|
|
|
|
if self.upsamplers is not None:
|
|
for upsampler in self.upsamplers:
|
|
hidden_states = upsampler(hidden_states, upsample_size)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class UpBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
prev_output_channel: int,
|
|
out_channels: int,
|
|
temb_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
output_scale_factor=1.0,
|
|
add_upsample=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
|
|
|
|
def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=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]
|
|
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
|
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
|
|
else:
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
if self.upsamplers is not None:
|
|
for upsampler in self.upsamplers:
|
|
hidden_states = upsampler(hidden_states, upsample_size)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class UpDecoderBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
output_scale_factor=1.0,
|
|
add_upsample=True,
|
|
):
|
|
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=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_upsample:
|
|
self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
|
|
else:
|
|
self.upsamplers = None
|
|
|
|
def forward(self, hidden_states):
|
|
for resnet in self.resnets:
|
|
hidden_states = resnet(hidden_states, temb=None)
|
|
|
|
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,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
attn_num_head_channels=1,
|
|
output_scale_factor=1.0,
|
|
add_upsample=True,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
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=None,
|
|
eps=resnet_eps,
|
|
groups=resnet_groups,
|
|
dropout=dropout,
|
|
time_embedding_norm=resnet_time_scale_shift,
|
|
non_linearity=resnet_act_fn,
|
|
output_scale_factor=output_scale_factor,
|
|
pre_norm=resnet_pre_norm,
|
|
)
|
|
)
|
|
attentions.append(
|
|
AttentionBlock(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
norm_num_groups=resnet_groups,
|
|
)
|
|
)
|
|
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_upsample:
|
|
self.upsamplers = nn.ModuleList([Upsample2D(out_channels, use_conv=True, out_channels=out_channels)])
|
|
else:
|
|
self.upsamplers = None
|
|
|
|
def forward(self, hidden_states):
|
|
for resnet, attn in zip(self.resnets, self.attentions):
|
|
hidden_states = resnet(hidden_states, temb=None)
|
|
hidden_states = attn(hidden_states)
|
|
|
|
if self.upsamplers is not None:
|
|
for upsampler in self.upsamplers:
|
|
hidden_states = upsampler(hidden_states)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class AttnSkipUpBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
prev_output_channel: 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,
|
|
attn_num_head_channels=1,
|
|
output_scale_factor=np.sqrt(2.0),
|
|
upsample_padding=1,
|
|
add_upsample=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,
|
|
)
|
|
)
|
|
|
|
self.attentions.append(
|
|
AttentionBlock(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
)
|
|
)
|
|
|
|
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
|
|
|
|
def forward(self, hidden_states, res_hidden_states_tuple, temb=None, skip_sample=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]
|
|
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
|
|
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
hidden_states = self.attentions[0](hidden_states)
|
|
|
|
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)
|
|
|
|
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,
|
|
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=np.sqrt(2.0),
|
|
add_upsample=True,
|
|
upsample_padding=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
|
|
|
|
def forward(self, hidden_states, res_hidden_states_tuple, temb=None, skip_sample=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]
|
|
hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1)
|
|
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
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)
|
|
|
|
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,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
output_scale_factor=1.0,
|
|
add_upsample=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(
|
|
[
|
|
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.gradient_checkpointing = False
|
|
|
|
def forward(self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=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]
|
|
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
|
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint(create_custom_forward(resnet), hidden_states, temb)
|
|
else:
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
if self.upsamplers is not None:
|
|
for upsampler in self.upsamplers:
|
|
hidden_states = upsampler(hidden_states, temb)
|
|
|
|
return hidden_states
|
|
|
|
|
|
class SimpleCrossAttnUpBlock2D(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
out_channels: int,
|
|
prev_output_channel: int,
|
|
temb_channels: int,
|
|
dropout: float = 0.0,
|
|
num_layers: int = 1,
|
|
resnet_eps: float = 1e-6,
|
|
resnet_time_scale_shift: str = "default",
|
|
resnet_act_fn: str = "swish",
|
|
resnet_groups: int = 32,
|
|
resnet_pre_norm: bool = True,
|
|
attn_num_head_channels=1,
|
|
cross_attention_dim=1280,
|
|
output_scale_factor=1.0,
|
|
add_upsample=True,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
self.has_cross_attention = True
|
|
self.attn_num_head_channels = attn_num_head_channels
|
|
|
|
self.num_heads = out_channels // self.attn_num_head_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(
|
|
CrossAttention(
|
|
query_dim=out_channels,
|
|
cross_attention_dim=out_channels,
|
|
heads=self.num_heads,
|
|
dim_head=attn_num_head_channels,
|
|
added_kv_proj_dim=cross_attention_dim,
|
|
norm_num_groups=resnet_groups,
|
|
bias=True,
|
|
upcast_softmax=True,
|
|
)
|
|
)
|
|
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,
|
|
up=True,
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
self.upsamplers = None
|
|
|
|
self.gradient_checkpointing = False
|
|
|
|
def forward(
|
|
self,
|
|
hidden_states,
|
|
res_hidden_states_tuple,
|
|
temb=None,
|
|
encoder_hidden_states=None,
|
|
upsample_size=None,
|
|
attention_mask=None,
|
|
):
|
|
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)
|
|
|
|
hidden_states = resnet(hidden_states, temb)
|
|
|
|
# attn
|
|
residual = hidden_states
|
|
hidden_states = hidden_states.view(hidden_states.shape[0], hidden_states.shape[1], -1).transpose(1, 2)
|
|
hidden_states = attn(
|
|
hidden_states,
|
|
encoder_hidden_states=encoder_hidden_states.transpose(1, 2),
|
|
attention_mask=attention_mask,
|
|
)
|
|
hidden_states = hidden_states.transpose(-1, -2).reshape(residual.shape)
|
|
hidden_states = hidden_states + residual
|
|
|
|
if self.upsamplers is not None:
|
|
for upsampler in self.upsamplers:
|
|
hidden_states = upsampler(hidden_states, temb)
|
|
|
|
return hidden_states
|