1433 lines
49 KiB
Python
1433 lines
49 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
|
|
|
|
import numpy as np
|
|
|
|
# limitations under the License.
|
|
import torch
|
|
from torch import nn
|
|
|
|
from .attention import AttentionBlockNew, SpatialTransformer
|
|
from .resnet import Downsample2D, FirDownsample2D, FirUpsample2D, ResnetBlock, 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,
|
|
cross_attention_dim=None,
|
|
downsample_padding=None,
|
|
):
|
|
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,
|
|
downsample_padding=downsample_padding,
|
|
)
|
|
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,
|
|
downsample_padding=downsample_padding,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
)
|
|
elif down_block_type == "CrossAttnDownBlock2D":
|
|
if cross_attention_dim is None:
|
|
raise ValueError("cross_attention_dim must be specified for CrossAttnUpBlock2D")
|
|
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,
|
|
downsample_padding=downsample_padding,
|
|
cross_attention_dim=cross_attention_dim,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
)
|
|
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,
|
|
)
|
|
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,
|
|
)
|
|
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,
|
|
downsample_padding=downsample_padding,
|
|
)
|
|
|
|
|
|
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,
|
|
cross_attention_dim=None,
|
|
):
|
|
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,
|
|
)
|
|
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,
|
|
cross_attention_dim=cross_attention_dim,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
)
|
|
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,
|
|
attn_num_head_channels=attn_num_head_channels,
|
|
)
|
|
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,
|
|
)
|
|
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,
|
|
)
|
|
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,
|
|
)
|
|
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,
|
|
attn_num_head_channels=1,
|
|
attention_type="default",
|
|
output_scale_factor=1.0,
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
|
|
self.attention_type = attention_type
|
|
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
|
|
|
# there is always at least one resnet
|
|
resnets = [
|
|
ResnetBlock(
|
|
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(
|
|
AttentionBlockNew(
|
|
in_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
num_groups=resnet_groups,
|
|
)
|
|
)
|
|
resnets.append(
|
|
ResnetBlock(
|
|
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_states=None):
|
|
hidden_states = self.resnets[0](hidden_states, temb)
|
|
for attn, resnet in zip(self.attentions, self.resnets[1:]):
|
|
if self.attention_type == "default":
|
|
hidden_states = attn(hidden_states)
|
|
else:
|
|
hidden_states = attn(hidden_states, encoder_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,
|
|
attention_type="default",
|
|
output_scale_factor=1.0,
|
|
cross_attention_dim=1280,
|
|
**kwargs,
|
|
):
|
|
super().__init__()
|
|
|
|
self.attention_type = attention_type
|
|
resnet_groups = resnet_groups if resnet_groups is not None else min(in_channels // 4, 32)
|
|
|
|
# there is always at least one resnet
|
|
resnets = [
|
|
ResnetBlock(
|
|
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(
|
|
SpatialTransformer(
|
|
in_channels,
|
|
attn_num_head_channels,
|
|
in_channels // attn_num_head_channels,
|
|
depth=1,
|
|
context_dim=cross_attention_dim,
|
|
)
|
|
)
|
|
resnets.append(
|
|
ResnetBlock(
|
|
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):
|
|
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)
|
|
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,
|
|
attention_type="default",
|
|
output_scale_factor=1.0,
|
|
downsample_padding=1,
|
|
add_downsample=True,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
self.attention_type = attention_type
|
|
|
|
for i in range(num_layers):
|
|
in_channels = in_channels if i == 0 else out_channels
|
|
resnets.append(
|
|
ResnetBlock(
|
|
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(
|
|
AttentionBlockNew(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
)
|
|
)
|
|
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_downsample:
|
|
self.downsamplers = nn.ModuleList(
|
|
[
|
|
Downsample2D(
|
|
in_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,
|
|
attention_type="default",
|
|
output_scale_factor=1.0,
|
|
downsample_padding=1,
|
|
add_downsample=True,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
self.attention_type = attention_type
|
|
|
|
for i in range(num_layers):
|
|
in_channels = in_channels if i == 0 else out_channels
|
|
resnets.append(
|
|
ResnetBlock(
|
|
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(
|
|
SpatialTransformer(
|
|
out_channels,
|
|
attn_num_head_channels,
|
|
out_channels // attn_num_head_channels,
|
|
depth=1,
|
|
context_dim=cross_attention_dim,
|
|
)
|
|
)
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_downsample:
|
|
self.downsamplers = nn.ModuleList(
|
|
[
|
|
Downsample2D(
|
|
in_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op"
|
|
)
|
|
]
|
|
)
|
|
else:
|
|
self.downsamplers = None
|
|
|
|
def forward(self, hidden_states, temb=None, encoder_hidden_states=None):
|
|
output_states = ()
|
|
|
|
for resnet, attn in zip(self.resnets, self.attentions):
|
|
hidden_states = resnet(hidden_states, temb)
|
|
hidden_states = attn(hidden_states, context=encoder_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 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(
|
|
ResnetBlock(
|
|
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(
|
|
in_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 in self.resnets:
|
|
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(
|
|
ResnetBlock(
|
|
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(
|
|
in_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(
|
|
ResnetBlock(
|
|
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(
|
|
AttentionBlockNew(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
num_groups=resnet_groups,
|
|
)
|
|
)
|
|
|
|
self.attentions = nn.ModuleList(attentions)
|
|
self.resnets = nn.ModuleList(resnets)
|
|
|
|
if add_downsample:
|
|
self.downsamplers = nn.ModuleList(
|
|
[
|
|
Downsample2D(
|
|
in_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,
|
|
attention_type="default",
|
|
output_scale_factor=np.sqrt(2.0),
|
|
downsample_padding=1,
|
|
add_downsample=True,
|
|
):
|
|
super().__init__()
|
|
self.attentions = nn.ModuleList([])
|
|
self.resnets = nn.ModuleList([])
|
|
|
|
self.attention_type = attention_type
|
|
|
|
for i in range(num_layers):
|
|
in_channels = in_channels if i == 0 else out_channels
|
|
self.resnets.append(
|
|
ResnetBlock(
|
|
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(
|
|
AttentionBlockNew(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
)
|
|
)
|
|
|
|
if add_downsample:
|
|
self.resnet_down = ResnetBlock(
|
|
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_nin_shortcut=True,
|
|
down=True,
|
|
kernel="fir",
|
|
)
|
|
self.downsamplers = nn.ModuleList([FirDownsample2D(in_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(
|
|
ResnetBlock(
|
|
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 = ResnetBlock(
|
|
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_nin_shortcut=True,
|
|
down=True,
|
|
kernel="fir",
|
|
)
|
|
self.downsamplers = nn.ModuleList([FirDownsample2D(in_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 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,
|
|
attention_type="default",
|
|
attn_num_head_channels=1,
|
|
output_scale_factor=1.0,
|
|
add_upsample=True,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
self.attention_type = attention_type
|
|
|
|
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(
|
|
ResnetBlock(
|
|
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(
|
|
AttentionBlockNew(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
)
|
|
)
|
|
|
|
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,
|
|
attention_type="default",
|
|
output_scale_factor=1.0,
|
|
downsample_padding=1,
|
|
add_upsample=True,
|
|
):
|
|
super().__init__()
|
|
resnets = []
|
|
attentions = []
|
|
|
|
self.attention_type = attention_type
|
|
|
|
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(
|
|
ResnetBlock(
|
|
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(
|
|
SpatialTransformer(
|
|
out_channels,
|
|
attn_num_head_channels,
|
|
out_channels // attn_num_head_channels,
|
|
depth=1,
|
|
context_dim=cross_attention_dim,
|
|
)
|
|
)
|
|
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, encoder_hidden_states=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, context=encoder_hidden_states)
|
|
|
|
if self.upsamplers is not None:
|
|
for upsampler in self.upsamplers:
|
|
hidden_states = upsampler(hidden_states)
|
|
|
|
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(
|
|
ResnetBlock(
|
|
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
|
|
|
|
def forward(self, hidden_states, res_hidden_states_tuple, temb=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 self.upsamplers is not None:
|
|
for upsampler in self.upsamplers:
|
|
hidden_states = upsampler(hidden_states)
|
|
|
|
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(
|
|
ResnetBlock(
|
|
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(
|
|
ResnetBlock(
|
|
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(
|
|
AttentionBlockNew(
|
|
out_channels,
|
|
num_head_channels=attn_num_head_channels,
|
|
rescale_output_factor=output_scale_factor,
|
|
eps=resnet_eps,
|
|
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,
|
|
attention_type="default",
|
|
output_scale_factor=np.sqrt(2.0),
|
|
upsample_padding=1,
|
|
add_upsample=True,
|
|
):
|
|
super().__init__()
|
|
self.attentions = nn.ModuleList([])
|
|
self.resnets = nn.ModuleList([])
|
|
|
|
self.attention_type = attention_type
|
|
|
|
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(
|
|
ResnetBlock(
|
|
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(
|
|
AttentionBlockNew(
|
|
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 = ResnetBlock(
|
|
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_nin_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(
|
|
ResnetBlock(
|
|
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 = ResnetBlock(
|
|
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_nin_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
|