197 lines
7.9 KiB
Python
Executable File
197 lines
7.9 KiB
Python
Executable File
import torch
|
|
import torch.nn as nn
|
|
|
|
from .attention import Attention3D, SpatialAttention, TemporalAttention
|
|
from .common import ResidualBlock3D
|
|
|
|
|
|
def get_mid_block(
|
|
mid_block_type: str,
|
|
in_channels: int,
|
|
num_layers: int,
|
|
act_fn: str,
|
|
norm_num_groups: int = 32,
|
|
norm_eps: float = 1e-6,
|
|
dropout: float = 0.0,
|
|
add_attention: bool = True,
|
|
attention_type: str = "3d",
|
|
num_attention_heads: int = 1,
|
|
output_scale_factor: float = 1.0,
|
|
) -> nn.Module:
|
|
if mid_block_type == "MidBlock3D":
|
|
return MidBlock3D(
|
|
in_channels=in_channels,
|
|
num_layers=num_layers,
|
|
act_fn=act_fn,
|
|
norm_num_groups=norm_num_groups,
|
|
norm_eps=norm_eps,
|
|
dropout=dropout,
|
|
add_attention=add_attention,
|
|
attention_type=attention_type,
|
|
attention_head_dim=in_channels // num_attention_heads,
|
|
output_scale_factor=output_scale_factor,
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown mid block type: {mid_block_type}")
|
|
|
|
|
|
class MidBlock3D(nn.Module):
|
|
"""
|
|
A 3D UNet mid-block [`MidBlock3D`] with multiple residual blocks and optional attention blocks.
|
|
|
|
Args:
|
|
in_channels (`int`): The number of input channels.
|
|
num_layers (`int`, *optional*, defaults to 1): The number of residual blocks.
|
|
act_fn (`str`, *optional*, defaults to `swish`): The activation function for the resnet blocks.
|
|
norm_num_groups (`int`, *optional*, defaults to 32):
|
|
The number of groups to use in the group normalization layers of the resnet blocks.
|
|
norm_eps (`float`, *optional*, 1e-6 ): The epsilon value for the resnet blocks.
|
|
dropout (`float`, *optional*, defaults to 0.0): The dropout rate.
|
|
add_attention (`bool`, *optional*, defaults to `True`): Whether to add attention blocks.
|
|
attention_type: (`str`, *optional*, defaults to `3d`): The type of attention to use. Defaults to `3d`.
|
|
attention_head_dim (`int`, *optional*, defaults to 1):
|
|
Dimension of a single attention head. The number of attention heads is determined based on this value and
|
|
the number of input channels.
|
|
output_scale_factor (`float`, *optional*, defaults to 1.0): The output scale factor.
|
|
|
|
Returns:
|
|
`torch.FloatTensor`: The output of the last residual block, which is a tensor of shape `(batch_size,
|
|
in_channels, temporal_length, height, width)`.
|
|
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
in_channels: int,
|
|
num_layers: int = 1,
|
|
act_fn: str = "silu",
|
|
norm_num_groups: int = 32,
|
|
norm_eps: float = 1e-6,
|
|
dropout: float = 0.0,
|
|
add_attention: bool = True,
|
|
attention_type: str = "3d",
|
|
attention_head_dim: int = 1,
|
|
output_scale_factor: float = 1.0,
|
|
):
|
|
super().__init__()
|
|
|
|
self.attention_type = attention_type
|
|
|
|
norm_num_groups = norm_num_groups if norm_num_groups is not None else min(in_channels // 4, 32)
|
|
|
|
self.convs = nn.ModuleList([
|
|
ResidualBlock3D(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
non_linearity=act_fn,
|
|
norm_num_groups=norm_num_groups,
|
|
norm_eps=norm_eps,
|
|
dropout=dropout,
|
|
output_scale_factor=output_scale_factor,
|
|
)
|
|
])
|
|
|
|
self.attentions = nn.ModuleList([])
|
|
for _ in range(num_layers - 1):
|
|
if add_attention:
|
|
if attention_type == "3d":
|
|
self.attentions.append(
|
|
Attention3D(
|
|
in_channels,
|
|
nheads=in_channels // attention_head_dim,
|
|
head_dim=attention_head_dim,
|
|
bias=True,
|
|
upcast_softmax=True,
|
|
norm_num_groups=norm_num_groups,
|
|
eps=norm_eps,
|
|
rescale_output_factor=output_scale_factor,
|
|
residual_connection=True,
|
|
)
|
|
)
|
|
elif attention_type == "spatial_temporal":
|
|
self.attentions.append(
|
|
nn.ModuleList([
|
|
SpatialAttention(
|
|
in_channels,
|
|
nheads=in_channels // attention_head_dim,
|
|
head_dim=attention_head_dim,
|
|
bias=True,
|
|
upcast_softmax=True,
|
|
norm_num_groups=norm_num_groups,
|
|
eps=norm_eps,
|
|
rescale_output_factor=output_scale_factor,
|
|
residual_connection=True,
|
|
),
|
|
TemporalAttention(
|
|
in_channels,
|
|
nheads=in_channels // attention_head_dim,
|
|
head_dim=attention_head_dim,
|
|
bias=True,
|
|
upcast_softmax=True,
|
|
norm_num_groups=norm_num_groups,
|
|
eps=norm_eps,
|
|
rescale_output_factor=output_scale_factor,
|
|
residual_connection=True,
|
|
),
|
|
])
|
|
)
|
|
elif attention_type == "spatial":
|
|
self.attentions.append(
|
|
SpatialAttention(
|
|
in_channels,
|
|
nheads=in_channels // attention_head_dim,
|
|
head_dim=attention_head_dim,
|
|
bias=True,
|
|
upcast_softmax=True,
|
|
norm_num_groups=norm_num_groups,
|
|
eps=norm_eps,
|
|
rescale_output_factor=output_scale_factor,
|
|
residual_connection=True,
|
|
)
|
|
)
|
|
elif attention_type == "temporal":
|
|
self.attentions.append(
|
|
TemporalAttention(
|
|
in_channels,
|
|
nheads=in_channels // attention_head_dim,
|
|
head_dim=attention_head_dim,
|
|
bias=True,
|
|
upcast_softmax=True,
|
|
norm_num_groups=norm_num_groups,
|
|
eps=norm_eps,
|
|
rescale_output_factor=output_scale_factor,
|
|
residual_connection=True,
|
|
)
|
|
)
|
|
else:
|
|
raise ValueError(f"Unknown attention type: {attention_type}")
|
|
else:
|
|
self.attentions.append(None)
|
|
|
|
self.convs.append(
|
|
ResidualBlock3D(
|
|
in_channels=in_channels,
|
|
out_channels=in_channels,
|
|
non_linearity=act_fn,
|
|
norm_num_groups=norm_num_groups,
|
|
norm_eps=norm_eps,
|
|
dropout=dropout,
|
|
output_scale_factor=output_scale_factor,
|
|
)
|
|
)
|
|
|
|
def forward(self, hidden_states: torch.FloatTensor) -> torch.FloatTensor:
|
|
hidden_states = self.convs[0](hidden_states)
|
|
|
|
for attn, resnet in zip(self.attentions, self.convs[1:]):
|
|
if attn is not None:
|
|
if self.attention_type == "spatial_temporal":
|
|
spatial_attn, temporal_attn = attn
|
|
hidden_states = spatial_attn(hidden_states)
|
|
hidden_states = temporal_attn(hidden_states)
|
|
else:
|
|
hidden_states = attn(hidden_states)
|
|
hidden_states = resnet(hidden_states)
|
|
|
|
return hidden_states
|