Spaces:
Running
Running
| 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 | |