|
|
|
|
|
|
|
""" |
|
This module defines various 3D UNet blocks used in the video model. |
|
|
|
The blocks include: |
|
- UNetMidBlock3DCrossAttn: The middle block of the UNet with cross attention. |
|
- CrossAttnDownBlock3D: The downsampling block with cross attention. |
|
- DownBlock3D: The standard downsampling block without cross attention. |
|
- CrossAttnUpBlock3D: The upsampling block with cross attention. |
|
- UpBlock3D: The standard upsampling block without cross attention. |
|
|
|
These blocks are used to construct the 3D UNet architecture for video-related tasks. |
|
""" |
|
|
|
import torch |
|
from einops import rearrange |
|
from torch import nn |
|
|
|
from .motion_module import get_motion_module |
|
from .resnet import Downsample3D, ResnetBlock3D, Upsample3D |
|
from .transformer_3d import Transformer3DModel |
|
|
|
|
|
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, |
|
audio_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", |
|
unet_use_cross_frame_attention=None, |
|
unet_use_temporal_attention=None, |
|
use_inflated_groupnorm=None, |
|
use_motion_module=None, |
|
motion_module_type=None, |
|
motion_module_kwargs=None, |
|
use_audio_module=None, |
|
depth=0, |
|
stack_enable_blocks_name=None, |
|
stack_enable_blocks_depth=None, |
|
): |
|
""" |
|
Factory function to instantiate a down-block module for the 3D UNet architecture. |
|
|
|
Down blocks are used in the downsampling part of the U-Net to reduce the spatial dimensions |
|
of the feature maps while increasing the depth. This function can create blocks with or without |
|
cross attention based on the specified parameters. |
|
|
|
Parameters: |
|
- down_block_type (str): The type of down block to instantiate. |
|
- num_layers (int): The number of layers in the block. |
|
- in_channels (int): The number of input channels. |
|
- out_channels (int): The number of output channels. |
|
- temb_channels (int): The number of token embedding channels. |
|
- add_downsample (bool): Flag to add a downsampling layer. |
|
- resnet_eps (float): Epsilon for residual block stability. |
|
- resnet_act_fn (callable): Activation function for the residual block. |
|
- ... (remaining parameters): Additional parameters for configuring the block. |
|
|
|
Returns: |
|
- nn.Module: An instance of a down-sampling block module. |
|
""" |
|
down_block_type = ( |
|
down_block_type[7:] |
|
if down_block_type.startswith("UNetRes") |
|
else down_block_type |
|
) |
|
if down_block_type == "DownBlock3D": |
|
return DownBlock3D( |
|
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, |
|
use_inflated_groupnorm=use_inflated_groupnorm, |
|
use_motion_module=use_motion_module, |
|
motion_module_type=motion_module_type, |
|
motion_module_kwargs=motion_module_kwargs, |
|
) |
|
|
|
if down_block_type == "CrossAttnDownBlock3D": |
|
if cross_attention_dim is None: |
|
raise ValueError( |
|
"cross_attention_dim must be specified for CrossAttnDownBlock3D" |
|
) |
|
return CrossAttnDownBlock3D( |
|
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, |
|
audio_attention_dim=audio_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, |
|
unet_use_cross_frame_attention=unet_use_cross_frame_attention, |
|
unet_use_temporal_attention=unet_use_temporal_attention, |
|
use_inflated_groupnorm=use_inflated_groupnorm, |
|
use_motion_module=use_motion_module, |
|
motion_module_type=motion_module_type, |
|
motion_module_kwargs=motion_module_kwargs, |
|
use_audio_module=use_audio_module, |
|
depth=depth, |
|
stack_enable_blocks_name=stack_enable_blocks_name, |
|
stack_enable_blocks_depth=stack_enable_blocks_depth, |
|
) |
|
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, |
|
audio_attention_dim=None, |
|
dual_cross_attention=False, |
|
use_linear_projection=False, |
|
only_cross_attention=False, |
|
upcast_attention=False, |
|
resnet_time_scale_shift="default", |
|
unet_use_cross_frame_attention=None, |
|
unet_use_temporal_attention=None, |
|
use_inflated_groupnorm=None, |
|
use_motion_module=None, |
|
motion_module_type=None, |
|
motion_module_kwargs=None, |
|
use_audio_module=None, |
|
depth=0, |
|
stack_enable_blocks_name=None, |
|
stack_enable_blocks_depth=None, |
|
): |
|
""" |
|
Factory function to instantiate an up-block module for the 3D UNet architecture. |
|
|
|
Up blocks are used in the upsampling part of the U-Net to increase the spatial dimensions |
|
of the feature maps while decreasing the depth. This function can create blocks with or without |
|
cross attention based on the specified parameters. |
|
|
|
Parameters: |
|
- up_block_type (str): The type of up block to instantiate. |
|
- num_layers (int): The number of layers in the block. |
|
- in_channels (int): The number of input channels. |
|
- out_channels (int): The number of output channels. |
|
- prev_output_channel (int): The number of channels from the previous layer's output. |
|
- temb_channels (int): The number of token embedding channels. |
|
- add_upsample (bool): Flag to add an upsampling layer. |
|
- resnet_eps (float): Epsilon for residual block stability. |
|
- resnet_act_fn (callable): Activation function for the residual block. |
|
- ... (remaining parameters): Additional parameters for configuring the block. |
|
|
|
Returns: |
|
- nn.Module: An instance of an up-sampling block module. |
|
""" |
|
up_block_type = ( |
|
up_block_type[7:] if up_block_type.startswith("UNetRes") else up_block_type |
|
) |
|
if up_block_type == "UpBlock3D": |
|
return UpBlock3D( |
|
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, |
|
use_inflated_groupnorm=use_inflated_groupnorm, |
|
use_motion_module=use_motion_module, |
|
motion_module_type=motion_module_type, |
|
motion_module_kwargs=motion_module_kwargs, |
|
) |
|
|
|
if up_block_type == "CrossAttnUpBlock3D": |
|
if cross_attention_dim is None: |
|
raise ValueError( |
|
"cross_attention_dim must be specified for CrossAttnUpBlock3D" |
|
) |
|
return CrossAttnUpBlock3D( |
|
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, |
|
audio_attention_dim=audio_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, |
|
unet_use_cross_frame_attention=unet_use_cross_frame_attention, |
|
unet_use_temporal_attention=unet_use_temporal_attention, |
|
use_inflated_groupnorm=use_inflated_groupnorm, |
|
use_motion_module=use_motion_module, |
|
motion_module_type=motion_module_type, |
|
motion_module_kwargs=motion_module_kwargs, |
|
use_audio_module=use_audio_module, |
|
depth=depth, |
|
stack_enable_blocks_name=stack_enable_blocks_name, |
|
stack_enable_blocks_depth=stack_enable_blocks_depth, |
|
) |
|
raise ValueError(f"{up_block_type} does not exist.") |
|
|
|
|
|
class UNetMidBlock3DCrossAttn(nn.Module): |
|
""" |
|
A 3D UNet middle block with cross attention mechanism. This block is part of the U-Net architecture |
|
and is used for feature extraction in the middle of the downsampling path. |
|
|
|
Parameters: |
|
- in_channels (int): Number of input channels. |
|
- temb_channels (int): Number of token embedding channels. |
|
- dropout (float): Dropout rate. |
|
- num_layers (int): Number of layers in the block. |
|
- resnet_eps (float): Epsilon for residual block. |
|
- resnet_time_scale_shift (str): Time scale shift for time embedding normalization. |
|
- resnet_act_fn (str): Activation function for the residual block. |
|
- resnet_groups (int): Number of groups for the convolutions in the residual block. |
|
- resnet_pre_norm (bool): Whether to use pre-normalization in the residual block. |
|
- attn_num_head_channels (int): Number of attention heads. |
|
- cross_attention_dim (int): Dimensionality of the cross attention layers. |
|
- audio_attention_dim (int): Dimensionality of the audio attention layers. |
|
- dual_cross_attention (bool): Whether to use dual cross attention. |
|
- use_linear_projection (bool): Whether to use linear projection in attention. |
|
- upcast_attention (bool): Whether to upcast attention to the original input dimension. |
|
- unet_use_cross_frame_attention (bool): Whether to use cross frame attention in U-Net. |
|
- unet_use_temporal_attention (bool): Whether to use temporal attention in U-Net. |
|
- use_inflated_groupnorm (bool): Whether to use inflated group normalization. |
|
- use_motion_module (bool): Whether to use motion module. |
|
- motion_module_type (str): Type of motion module. |
|
- motion_module_kwargs (dict): Keyword arguments for the motion module. |
|
- use_audio_module (bool): Whether to use audio module. |
|
- depth (int): Depth of the block in the network. |
|
- stack_enable_blocks_name (str): Name of the stack enable blocks. |
|
- stack_enable_blocks_depth (int): Depth of the stack enable blocks. |
|
|
|
Forward method: |
|
The forward method applies the residual blocks, cross attention, and optional motion and audio modules |
|
to the input hidden states. It returns the transformed hidden states. |
|
""" |
|
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, |
|
audio_attention_dim=1024, |
|
dual_cross_attention=False, |
|
use_linear_projection=False, |
|
upcast_attention=False, |
|
unet_use_cross_frame_attention=None, |
|
unet_use_temporal_attention=None, |
|
use_inflated_groupnorm=None, |
|
use_motion_module=None, |
|
motion_module_type=None, |
|
motion_module_kwargs=None, |
|
use_audio_module=None, |
|
depth=0, |
|
stack_enable_blocks_name=None, |
|
stack_enable_blocks_depth=None, |
|
): |
|
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) |
|
) |
|
|
|
|
|
resnets = [ |
|
ResnetBlock3D( |
|
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, |
|
use_inflated_groupnorm=use_inflated_groupnorm, |
|
) |
|
] |
|
attentions = [] |
|
motion_modules = [] |
|
audio_modules = [] |
|
|
|
for _ in range(num_layers): |
|
if dual_cross_attention: |
|
raise NotImplementedError |
|
attentions.append( |
|
Transformer3DModel( |
|
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, |
|
unet_use_cross_frame_attention=unet_use_cross_frame_attention, |
|
unet_use_temporal_attention=unet_use_temporal_attention, |
|
) |
|
) |
|
audio_modules.append( |
|
Transformer3DModel( |
|
attn_num_head_channels, |
|
in_channels // attn_num_head_channels, |
|
in_channels=in_channels, |
|
num_layers=1, |
|
cross_attention_dim=audio_attention_dim, |
|
norm_num_groups=resnet_groups, |
|
use_linear_projection=use_linear_projection, |
|
upcast_attention=upcast_attention, |
|
use_audio_module=use_audio_module, |
|
depth=depth, |
|
unet_block_name="mid", |
|
stack_enable_blocks_name=stack_enable_blocks_name, |
|
stack_enable_blocks_depth=stack_enable_blocks_depth, |
|
) |
|
if use_audio_module |
|
else None |
|
) |
|
|
|
motion_modules.append( |
|
get_motion_module( |
|
in_channels=in_channels, |
|
motion_module_type=motion_module_type, |
|
motion_module_kwargs=motion_module_kwargs, |
|
) |
|
if use_motion_module |
|
else None |
|
) |
|
resnets.append( |
|
ResnetBlock3D( |
|
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, |
|
use_inflated_groupnorm=use_inflated_groupnorm, |
|
) |
|
) |
|
|
|
self.attentions = nn.ModuleList(attentions) |
|
self.resnets = nn.ModuleList(resnets) |
|
self.audio_modules = nn.ModuleList(audio_modules) |
|
self.motion_modules = nn.ModuleList(motion_modules) |
|
|
|
def forward( |
|
self, |
|
hidden_states, |
|
temb=None, |
|
encoder_hidden_states=None, |
|
attention_mask=None, |
|
full_mask=None, |
|
face_mask=None, |
|
lip_mask=None, |
|
audio_embedding=None, |
|
motion_scale=None, |
|
): |
|
""" |
|
Forward pass for the UNetMidBlock3DCrossAttn class. |
|
|
|
Args: |
|
self (UNetMidBlock3DCrossAttn): An instance of the UNetMidBlock3DCrossAttn class. |
|
hidden_states (Tensor): The input hidden states tensor. |
|
temb (Tensor, optional): The input temporal embedding tensor. Defaults to None. |
|
encoder_hidden_states (Tensor, optional): The encoder hidden states tensor. Defaults to None. |
|
attention_mask (Tensor, optional): The attention mask tensor. Defaults to None. |
|
full_mask (Tensor, optional): The full mask tensor. Defaults to None. |
|
face_mask (Tensor, optional): The face mask tensor. Defaults to None. |
|
lip_mask (Tensor, optional): The lip mask tensor. Defaults to None. |
|
audio_embedding (Tensor, optional): The audio embedding tensor. Defaults to None. |
|
|
|
Returns: |
|
Tensor: The output tensor after passing through the UNetMidBlock3DCrossAttn layers. |
|
""" |
|
hidden_states = self.resnets[0](hidden_states, temb) |
|
for attn, resnet, audio_module, motion_module in zip( |
|
self.attentions, self.resnets[1:], self.audio_modules, self.motion_modules |
|
): |
|
hidden_states, motion_frame = attn( |
|
hidden_states, |
|
encoder_hidden_states=encoder_hidden_states, |
|
return_dict=False, |
|
) |
|
if len(motion_frame[0]) > 0: |
|
|
|
motion_frames = motion_frame[0][0] |
|
motion_frames = rearrange( |
|
motion_frames, |
|
"b f (d1 d2) c -> b c f d1 d2", |
|
d1=hidden_states.size(-1), |
|
) |
|
|
|
else: |
|
motion_frames = torch.zeros( |
|
hidden_states.shape[0], |
|
hidden_states.shape[1], |
|
4, |
|
hidden_states.shape[3], |
|
hidden_states.shape[4], |
|
) |
|
|
|
n_motion_frames = motion_frames.size(2) |
|
if audio_module is not None: |
|
hidden_states = ( |
|
audio_module( |
|
hidden_states, |
|
encoder_hidden_states=audio_embedding, |
|
attention_mask=attention_mask, |
|
full_mask=full_mask, |
|
face_mask=face_mask, |
|
lip_mask=lip_mask, |
|
motion_scale=motion_scale, |
|
return_dict=False, |
|
) |
|
)[0] |
|
if motion_module is not None: |
|
motion_frames = motion_frames.to( |
|
device=hidden_states.device, dtype=hidden_states.dtype |
|
) |
|
|
|
_hidden_states = ( |
|
torch.cat([motion_frames, hidden_states], dim=2) |
|
if n_motion_frames > 0 |
|
else hidden_states |
|
) |
|
hidden_states = motion_module( |
|
_hidden_states, encoder_hidden_states=encoder_hidden_states |
|
) |
|
hidden_states = hidden_states[:, :, n_motion_frames:] |
|
|
|
hidden_states = resnet(hidden_states, temb) |
|
|
|
return hidden_states |
|
|
|
|
|
class CrossAttnDownBlock3D(nn.Module): |
|
""" |
|
A 3D downsampling block with cross attention for the U-Net architecture. |
|
|
|
Parameters: |
|
- (same as above, refer to the constructor for details) |
|
|
|
Forward method: |
|
The forward method downsamples the input hidden states using residual blocks and cross attention. |
|
It also applies optional motion and audio modules. The method supports gradient checkpointing |
|
to save memory during training. |
|
""" |
|
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, |
|
audio_attention_dim=1024, |
|
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, |
|
unet_use_cross_frame_attention=None, |
|
unet_use_temporal_attention=None, |
|
use_inflated_groupnorm=None, |
|
use_motion_module=None, |
|
motion_module_type=None, |
|
motion_module_kwargs=None, |
|
use_audio_module=None, |
|
depth=0, |
|
stack_enable_blocks_name=None, |
|
stack_enable_blocks_depth=None, |
|
): |
|
super().__init__() |
|
resnets = [] |
|
attentions = [] |
|
audio_modules = [] |
|
motion_modules = [] |
|
|
|
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( |
|
ResnetBlock3D( |
|
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, |
|
use_inflated_groupnorm=use_inflated_groupnorm, |
|
) |
|
) |
|
if dual_cross_attention: |
|
raise NotImplementedError |
|
attentions.append( |
|
Transformer3DModel( |
|
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, |
|
unet_use_cross_frame_attention=unet_use_cross_frame_attention, |
|
unet_use_temporal_attention=unet_use_temporal_attention, |
|
) |
|
) |
|
|
|
audio_modules.append( |
|
Transformer3DModel( |
|
attn_num_head_channels, |
|
in_channels // attn_num_head_channels, |
|
in_channels=out_channels, |
|
num_layers=1, |
|
cross_attention_dim=audio_attention_dim, |
|
norm_num_groups=resnet_groups, |
|
use_linear_projection=use_linear_projection, |
|
only_cross_attention=only_cross_attention, |
|
upcast_attention=upcast_attention, |
|
use_audio_module=use_audio_module, |
|
depth=depth, |
|
unet_block_name="down", |
|
stack_enable_blocks_name=stack_enable_blocks_name, |
|
stack_enable_blocks_depth=stack_enable_blocks_depth, |
|
) |
|
if use_audio_module |
|
else None |
|
) |
|
motion_modules.append( |
|
get_motion_module( |
|
in_channels=out_channels, |
|
motion_module_type=motion_module_type, |
|
motion_module_kwargs=motion_module_kwargs, |
|
) |
|
if use_motion_module |
|
else None |
|
) |
|
|
|
self.attentions = nn.ModuleList(attentions) |
|
self.resnets = nn.ModuleList(resnets) |
|
self.audio_modules = nn.ModuleList(audio_modules) |
|
self.motion_modules = nn.ModuleList(motion_modules) |
|
|
|
if add_downsample: |
|
self.downsamplers = nn.ModuleList( |
|
[ |
|
Downsample3D( |
|
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, |
|
full_mask=None, |
|
face_mask=None, |
|
lip_mask=None, |
|
audio_embedding=None, |
|
motion_scale=None, |
|
): |
|
""" |
|
Defines the forward pass for the CrossAttnDownBlock3D class. |
|
|
|
Parameters: |
|
- hidden_states : torch.Tensor |
|
The input tensor to the block. |
|
temb : torch.Tensor, optional |
|
The token embeddings from the previous block. |
|
encoder_hidden_states : torch.Tensor, optional |
|
The hidden states from the encoder. |
|
attention_mask : torch.Tensor, optional |
|
The attention mask for the cross-attention mechanism. |
|
full_mask : torch.Tensor, optional |
|
The full mask for the cross-attention mechanism. |
|
face_mask : torch.Tensor, optional |
|
The face mask for the cross-attention mechanism. |
|
lip_mask : torch.Tensor, optional |
|
The lip mask for the cross-attention mechanism. |
|
audio_embedding : torch.Tensor, optional |
|
The audio embedding for the cross-attention mechanism. |
|
|
|
Returns: |
|
-- torch.Tensor |
|
The output tensor from the block. |
|
""" |
|
output_states = () |
|
|
|
for _, (resnet, attn, audio_module, motion_module) in enumerate( |
|
zip(self.resnets, self.attentions, self.audio_modules, self.motion_modules) |
|
): |
|
|
|
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) |
|
|
|
return module(*inputs) |
|
|
|
return custom_forward |
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint( |
|
create_custom_forward(resnet), hidden_states, temb |
|
) |
|
|
|
motion_frames = [] |
|
hidden_states, motion_frame = torch.utils.checkpoint.checkpoint( |
|
create_custom_forward(attn, return_dict=False), |
|
hidden_states, |
|
encoder_hidden_states, |
|
) |
|
if len(motion_frame[0]) > 0: |
|
motion_frames = motion_frame[0][0] |
|
|
|
motion_frames = rearrange( |
|
motion_frames, |
|
"b f (d1 d2) c -> b c f d1 d2", |
|
d1=hidden_states.size(-1), |
|
) |
|
|
|
else: |
|
motion_frames = torch.zeros( |
|
hidden_states.shape[0], |
|
hidden_states.shape[1], |
|
4, |
|
hidden_states.shape[3], |
|
hidden_states.shape[4], |
|
) |
|
|
|
n_motion_frames = motion_frames.size(2) |
|
|
|
if audio_module is not None: |
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint( |
|
create_custom_forward(audio_module, return_dict=False), |
|
hidden_states, |
|
audio_embedding, |
|
attention_mask, |
|
full_mask, |
|
face_mask, |
|
lip_mask, |
|
motion_scale, |
|
)[0] |
|
|
|
|
|
if motion_module is not None: |
|
motion_frames = motion_frames.to( |
|
device=hidden_states.device, dtype=hidden_states.dtype |
|
) |
|
_hidden_states = torch.cat( |
|
[motion_frames, hidden_states], dim=2 |
|
) |
|
hidden_states = torch.utils.checkpoint.checkpoint( |
|
create_custom_forward(motion_module), |
|
_hidden_states, |
|
encoder_hidden_states, |
|
) |
|
hidden_states = hidden_states[:, :, n_motion_frames:] |
|
|
|
else: |
|
hidden_states = resnet(hidden_states, temb) |
|
hidden_states = attn( |
|
hidden_states, |
|
encoder_hidden_states=encoder_hidden_states, |
|
).sample |
|
if audio_module is not None: |
|
hidden_states = audio_module( |
|
hidden_states, |
|
audio_embedding, |
|
attention_mask=attention_mask, |
|
full_mask=full_mask, |
|
face_mask=face_mask, |
|
lip_mask=lip_mask, |
|
return_dict=False, |
|
)[0] |
|
|
|
if motion_module is not None: |
|
hidden_states = motion_module( |
|
hidden_states, encoder_hidden_states=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 DownBlock3D(nn.Module): |
|
""" |
|
A 3D downsampling block for the U-Net architecture. This block performs downsampling operations |
|
using residual blocks and an optional motion module. |
|
|
|
Parameters: |
|
- in_channels (int): Number of input channels. |
|
- out_channels (int): Number of output channels. |
|
- temb_channels (int): Number of token embedding channels. |
|
- dropout (float): Dropout rate for the block. |
|
- num_layers (int): Number of layers in the block. |
|
- resnet_eps (float): Epsilon for residual block stability. |
|
- resnet_time_scale_shift (str): Time scale shift for the residual block's time embedding. |
|
- resnet_act_fn (str): Activation function used in the residual block. |
|
- resnet_groups (int): Number of groups for the convolutions in the residual block. |
|
- resnet_pre_norm (bool): Whether to use pre-normalization in the residual block. |
|
- output_scale_factor (float): Scaling factor for the block's output. |
|
- add_downsample (bool): Whether to add a downsampling layer. |
|
- downsample_padding (int): Padding for the downsampling layer. |
|
- use_inflated_groupnorm (bool): Whether to use inflated group normalization. |
|
- use_motion_module (bool): Whether to include a motion module. |
|
- motion_module_type (str): Type of motion module to use. |
|
- motion_module_kwargs (dict): Keyword arguments for the motion module. |
|
|
|
Forward method: |
|
The forward method processes the input hidden states through the residual blocks and optional |
|
motion modules, followed by an optional downsampling step. It supports gradient checkpointing |
|
during training to reduce memory usage. |
|
""" |
|
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, |
|
use_inflated_groupnorm=None, |
|
use_motion_module=None, |
|
motion_module_type=None, |
|
motion_module_kwargs=None, |
|
): |
|
super().__init__() |
|
resnets = [] |
|
motion_modules = [] |
|
|
|
|
|
for i in range(num_layers): |
|
in_channels = in_channels if i == 0 else out_channels |
|
resnets.append( |
|
ResnetBlock3D( |
|
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, |
|
use_inflated_groupnorm=use_inflated_groupnorm, |
|
) |
|
) |
|
motion_modules.append( |
|
get_motion_module( |
|
in_channels=out_channels, |
|
motion_module_type=motion_module_type, |
|
motion_module_kwargs=motion_module_kwargs, |
|
) |
|
if use_motion_module |
|
else None |
|
) |
|
|
|
self.resnets = nn.ModuleList(resnets) |
|
self.motion_modules = nn.ModuleList(motion_modules) |
|
|
|
if add_downsample: |
|
self.downsamplers = nn.ModuleList( |
|
[ |
|
Downsample3D( |
|
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, |
|
): |
|
""" |
|
forward method for the DownBlock3D class. |
|
|
|
Args: |
|
hidden_states (Tensor): The input tensor to the DownBlock3D layer. |
|
temb (Tensor, optional): The token embeddings, if using transformer. |
|
encoder_hidden_states (Tensor, optional): The hidden states from the encoder. |
|
|
|
Returns: |
|
Tensor: The output tensor after passing through the DownBlock3D layer. |
|
""" |
|
output_states = () |
|
|
|
for resnet, motion_module in zip(self.resnets, self.motion_modules): |
|
|
|
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) |
|
|
|
|
|
hidden_states = ( |
|
motion_module( |
|
hidden_states, encoder_hidden_states=encoder_hidden_states |
|
) |
|
if motion_module is not None |
|
else 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 CrossAttnUpBlock3D(nn.Module): |
|
""" |
|
Standard 3D downsampling block for the U-Net architecture. This block performs downsampling |
|
operations in the U-Net using residual blocks and an optional motion module. |
|
|
|
Parameters: |
|
- in_channels (int): Number of input channels. |
|
- out_channels (int): Number of output channels. |
|
- temb_channels (int): Number of channels for the temporal embedding. |
|
- dropout (float): Dropout rate for the block. |
|
- num_layers (int): Number of layers in the block. |
|
- resnet_eps (float): Epsilon for residual block stability. |
|
- resnet_time_scale_shift (str): Time scale shift for the residual block's time embedding. |
|
- resnet_act_fn (str): Activation function used in the residual block. |
|
- resnet_groups (int): Number of groups for the convolutions in the residual block. |
|
- resnet_pre_norm (bool): Whether to use pre-normalization in the residual block. |
|
- output_scale_factor (float): Scaling factor for the block's output. |
|
- add_downsample (bool): Whether to add a downsampling layer. |
|
- downsample_padding (int): Padding for the downsampling layer. |
|
- use_inflated_groupnorm (bool): Whether to use inflated group normalization. |
|
- use_motion_module (bool): Whether to include a motion module. |
|
- motion_module_type (str): Type of motion module to use. |
|
- motion_module_kwargs (dict): Keyword arguments for the motion module. |
|
|
|
Forward method: |
|
The forward method processes the input hidden states through the residual blocks and optional |
|
motion modules, followed by an optional downsampling step. It supports gradient checkpointing |
|
during training to reduce memory usage. |
|
""" |
|
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, |
|
audio_attention_dim=1024, |
|
output_scale_factor=1.0, |
|
add_upsample=True, |
|
dual_cross_attention=False, |
|
use_linear_projection=False, |
|
only_cross_attention=False, |
|
upcast_attention=False, |
|
unet_use_cross_frame_attention=None, |
|
unet_use_temporal_attention=None, |
|
use_motion_module=None, |
|
use_inflated_groupnorm=None, |
|
motion_module_type=None, |
|
motion_module_kwargs=None, |
|
use_audio_module=None, |
|
depth=0, |
|
stack_enable_blocks_name=None, |
|
stack_enable_blocks_depth=None, |
|
): |
|
super().__init__() |
|
resnets = [] |
|
attentions = [] |
|
audio_modules = [] |
|
motion_modules = [] |
|
|
|
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( |
|
ResnetBlock3D( |
|
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, |
|
use_inflated_groupnorm=use_inflated_groupnorm, |
|
) |
|
) |
|
|
|
if dual_cross_attention: |
|
raise NotImplementedError |
|
attentions.append( |
|
Transformer3DModel( |
|
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, |
|
unet_use_cross_frame_attention=unet_use_cross_frame_attention, |
|
unet_use_temporal_attention=unet_use_temporal_attention, |
|
) |
|
) |
|
audio_modules.append( |
|
Transformer3DModel( |
|
attn_num_head_channels, |
|
in_channels // attn_num_head_channels, |
|
in_channels=out_channels, |
|
num_layers=1, |
|
cross_attention_dim=audio_attention_dim, |
|
norm_num_groups=resnet_groups, |
|
use_linear_projection=use_linear_projection, |
|
only_cross_attention=only_cross_attention, |
|
upcast_attention=upcast_attention, |
|
use_audio_module=use_audio_module, |
|
depth=depth, |
|
unet_block_name="up", |
|
stack_enable_blocks_name=stack_enable_blocks_name, |
|
stack_enable_blocks_depth=stack_enable_blocks_depth, |
|
) |
|
if use_audio_module |
|
else None |
|
) |
|
motion_modules.append( |
|
get_motion_module( |
|
in_channels=out_channels, |
|
motion_module_type=motion_module_type, |
|
motion_module_kwargs=motion_module_kwargs, |
|
) |
|
if use_motion_module |
|
else None |
|
) |
|
|
|
self.attentions = nn.ModuleList(attentions) |
|
self.resnets = nn.ModuleList(resnets) |
|
self.audio_modules = nn.ModuleList(audio_modules) |
|
self.motion_modules = nn.ModuleList(motion_modules) |
|
|
|
if add_upsample: |
|
self.upsamplers = nn.ModuleList( |
|
[Upsample3D(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, |
|
full_mask=None, |
|
face_mask=None, |
|
lip_mask=None, |
|
audio_embedding=None, |
|
motion_scale=None, |
|
): |
|
""" |
|
Forward pass for the CrossAttnUpBlock3D class. |
|
|
|
Args: |
|
self (CrossAttnUpBlock3D): An instance of the CrossAttnUpBlock3D class. |
|
hidden_states (Tensor): The input hidden states tensor. |
|
res_hidden_states_tuple (Tuple[Tensor]): A tuple of residual hidden states tensors. |
|
temb (Tensor, optional): The token embeddings tensor. Defaults to None. |
|
encoder_hidden_states (Tensor, optional): The encoder hidden states tensor. Defaults to None. |
|
upsample_size (int, optional): The upsample size. Defaults to None. |
|
attention_mask (Tensor, optional): The attention mask tensor. Defaults to None. |
|
full_mask (Tensor, optional): The full mask tensor. Defaults to None. |
|
face_mask (Tensor, optional): The face mask tensor. Defaults to None. |
|
lip_mask (Tensor, optional): The lip mask tensor. Defaults to None. |
|
audio_embedding (Tensor, optional): The audio embedding tensor. Defaults to None. |
|
|
|
Returns: |
|
Tensor: The output tensor after passing through the CrossAttnUpBlock3D. |
|
""" |
|
for _, (resnet, attn, audio_module, motion_module) in enumerate( |
|
zip(self.resnets, self.attentions, self.audio_modules, self.motion_modules) |
|
): |
|
|
|
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) |
|
|
|
return module(*inputs) |
|
|
|
return custom_forward |
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint( |
|
create_custom_forward(resnet), hidden_states, temb |
|
) |
|
|
|
motion_frames = [] |
|
hidden_states, motion_frame = torch.utils.checkpoint.checkpoint( |
|
create_custom_forward(attn, return_dict=False), |
|
hidden_states, |
|
encoder_hidden_states, |
|
) |
|
if len(motion_frame[0]) > 0: |
|
motion_frames = motion_frame[0][0] |
|
|
|
motion_frames = rearrange( |
|
motion_frames, |
|
"b f (d1 d2) c -> b c f d1 d2", |
|
d1=hidden_states.size(-1), |
|
) |
|
else: |
|
motion_frames = torch.zeros( |
|
hidden_states.shape[0], |
|
hidden_states.shape[1], |
|
4, |
|
hidden_states.shape[3], |
|
hidden_states.shape[4], |
|
) |
|
|
|
n_motion_frames = motion_frames.size(2) |
|
|
|
if audio_module is not None: |
|
|
|
hidden_states = torch.utils.checkpoint.checkpoint( |
|
create_custom_forward(audio_module, return_dict=False), |
|
hidden_states, |
|
audio_embedding, |
|
attention_mask, |
|
full_mask, |
|
face_mask, |
|
lip_mask, |
|
motion_scale, |
|
)[0] |
|
|
|
|
|
if motion_module is not None: |
|
motion_frames = motion_frames.to( |
|
device=hidden_states.device, dtype=hidden_states.dtype |
|
) |
|
|
|
_hidden_states = ( |
|
torch.cat([motion_frames, hidden_states], dim=2) |
|
if n_motion_frames > 0 |
|
else hidden_states |
|
) |
|
hidden_states = torch.utils.checkpoint.checkpoint( |
|
create_custom_forward(motion_module), |
|
_hidden_states, |
|
encoder_hidden_states, |
|
) |
|
hidden_states = hidden_states[:, :, n_motion_frames:] |
|
else: |
|
hidden_states = resnet(hidden_states, temb) |
|
hidden_states = attn( |
|
hidden_states, |
|
encoder_hidden_states=encoder_hidden_states, |
|
).sample |
|
|
|
if audio_module is not None: |
|
|
|
hidden_states = ( |
|
audio_module( |
|
hidden_states, |
|
encoder_hidden_states=audio_embedding, |
|
attention_mask=attention_mask, |
|
full_mask=full_mask, |
|
face_mask=face_mask, |
|
lip_mask=lip_mask, |
|
) |
|
).sample |
|
|
|
hidden_states = ( |
|
motion_module( |
|
hidden_states, encoder_hidden_states=encoder_hidden_states |
|
) |
|
if motion_module is not None |
|
else hidden_states |
|
) |
|
|
|
if self.upsamplers is not None: |
|
for upsampler in self.upsamplers: |
|
hidden_states = upsampler(hidden_states, upsample_size) |
|
|
|
return hidden_states |
|
|
|
|
|
class UpBlock3D(nn.Module): |
|
""" |
|
3D upsampling block with cross attention for the U-Net architecture. This block performs |
|
upsampling operations and incorporates cross attention mechanisms, which allow the model to |
|
focus on different parts of the input when upscaling. |
|
|
|
Parameters: |
|
- in_channels (int): Number of input channels. |
|
- out_channels (int): Number of output channels. |
|
- prev_output_channel (int): Number of channels from the previous layer's output. |
|
- temb_channels (int): Number of channels for the temporal embedding. |
|
- dropout (float): Dropout rate for the block. |
|
- num_layers (int): Number of layers in the block. |
|
- resnet_eps (float): Epsilon for residual block stability. |
|
- resnet_time_scale_shift (str): Time scale shift for the residual block's time embedding. |
|
- resnet_act_fn (str): Activation function used in the residual block. |
|
- resnet_groups (int): Number of groups for the convolutions in the residual block. |
|
- resnet_pre_norm (bool): Whether to use pre-normalization in the residual block. |
|
- attn_num_head_channels (int): Number of attention heads for the cross attention mechanism. |
|
- cross_attention_dim (int): Dimensionality of the cross attention layers. |
|
- audio_attention_dim (int): Dimensionality of the audio attention layers. |
|
- output_scale_factor (float): Scaling factor for the block's output. |
|
- add_upsample (bool): Whether to add an upsampling layer. |
|
- dual_cross_attention (bool): Whether to use dual cross attention (not implemented). |
|
- use_linear_projection (bool): Whether to use linear projection in the cross attention. |
|
- only_cross_attention (bool): Whether to use only cross attention (no self-attention). |
|
- upcast_attention (bool): Whether to upcast attention to the original input dimension. |
|
- unet_use_cross_frame_attention (bool): Whether to use cross frame attention in U-Net. |
|
- unet_use_temporal_attention (bool): Whether to use temporal attention in U-Net. |
|
- use_motion_module (bool): Whether to include a motion module. |
|
- use_inflated_groupnorm (bool): Whether to use inflated group normalization. |
|
- motion_module_type (str): Type of motion module to use. |
|
- motion_module_kwargs (dict): Keyword arguments for the motion module. |
|
- use_audio_module (bool): Whether to include an audio module. |
|
- depth (int): Depth of the block in the network. |
|
- stack_enable_blocks_name (str): Name of the stack enable blocks. |
|
- stack_enable_blocks_depth (int): Depth of the stack enable blocks. |
|
|
|
Forward method: |
|
The forward method upsamples the input hidden states and residual hidden states, processes |
|
them through the residual and cross attention blocks, and optional motion and audio modules. |
|
It supports gradient checkpointing during training. |
|
""" |
|
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, |
|
use_inflated_groupnorm=None, |
|
use_motion_module=None, |
|
motion_module_type=None, |
|
motion_module_kwargs=None, |
|
): |
|
super().__init__() |
|
resnets = [] |
|
motion_modules = [] |
|
|
|
|
|
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( |
|
ResnetBlock3D( |
|
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, |
|
use_inflated_groupnorm=use_inflated_groupnorm, |
|
) |
|
) |
|
motion_modules.append( |
|
get_motion_module( |
|
in_channels=out_channels, |
|
motion_module_type=motion_module_type, |
|
motion_module_kwargs=motion_module_kwargs, |
|
) |
|
if use_motion_module |
|
else None |
|
) |
|
|
|
self.resnets = nn.ModuleList(resnets) |
|
self.motion_modules = nn.ModuleList(motion_modules) |
|
|
|
if add_upsample: |
|
self.upsamplers = nn.ModuleList( |
|
[Upsample3D(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, |
|
encoder_hidden_states=None, |
|
): |
|
""" |
|
Forward pass for the UpBlock3D class. |
|
|
|
Args: |
|
self (UpBlock3D): An instance of the UpBlock3D class. |
|
hidden_states (Tensor): The input hidden states tensor. |
|
res_hidden_states_tuple (Tuple[Tensor]): A tuple of residual hidden states tensors. |
|
temb (Tensor, optional): The token embeddings tensor. Defaults to None. |
|
upsample_size (int, optional): The upsample size. Defaults to None. |
|
encoder_hidden_states (Tensor, optional): The encoder hidden states tensor. Defaults to None. |
|
|
|
Returns: |
|
Tensor: The output tensor after passing through the UpBlock3D layers. |
|
""" |
|
for resnet, motion_module in zip(self.resnets, self.motion_modules): |
|
|
|
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) |
|
hidden_states = ( |
|
motion_module( |
|
hidden_states, encoder_hidden_states=encoder_hidden_states |
|
) |
|
if motion_module is not None |
|
else hidden_states |
|
) |
|
|
|
if self.upsamplers is not None: |
|
for upsampler in self.upsamplers: |
|
hidden_states = upsampler(hidden_states, upsample_size) |
|
|
|
return hidden_states |
|
|