mmagic.models.editors.animatediff.unet_block

Module Contents

Classes

UNetMidBlock3DCrossAttn	3D unet mid block built by cross attention.
CrossAttnDownBlock3D	Down block built by 3D cross attention.
DownBlock3D	Down block built by 3D resnet.
CrossAttnUpBlock3D	Up block built by 3D cross attention.
UpBlock3D	Up block built by 3D resnet.

Functions

get_down_block(down_block_type, num_layers, ...[, ...])	get unet down path block.
get_up_block(up_block_type, num_layers, in_channels, ...)	get unet up path block.

mmagic.models.editors.animatediff.unet_block.get_down_block(down_block_type, num_layers, in_channels, out_channels, temb_channels, add_downsample, resnet_eps, resnet_act_fn, attn_num_head_channels, resnet_groups=None, cross_attention_dim=None, downsample_padding=None, dual_cross_attention=False, use_linear_projection=False, only_cross_attention=False, upcast_attention=False, resnet_time_scale_shift='default', 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)

get unet down path block.

mmagic.models.editors.animatediff.unet_block.get_up_block(up_block_type, num_layers, in_channels, out_channels, prev_output_channel, temb_channels, add_upsample, resnet_eps, resnet_act_fn, attn_num_head_channels, resnet_groups=None, cross_attention_dim=None, dual_cross_attention=False, use_linear_projection=False, only_cross_attention=False, upcast_attention=False, resnet_time_scale_shift='default', 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)

get unet up path block.

class mmagic.models.editors.animatediff.unet_block.UNetMidBlock3DCrossAttn(in_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-06, resnet_time_scale_shift: str = 'default', resnet_act_fn: str = 'swish', resnet_groups: int = 32, resnet_pre_norm: bool = True, attn_num_head_channels=1, output_scale_factor=1.0, cross_attention_dim=1280, dual_cross_attention=False, use_linear_projection=False, upcast_attention=False, 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)

Bases: torch.nn.Module

3D unet mid block built by cross attention.

forward(hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None)

forward with hidden states.

class mmagic.models.editors.animatediff.unet_block.CrossAttnDownBlock3D(in_channels: int, out_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-06, resnet_time_scale_shift: str = 'default', resnet_act_fn: str = 'swish', resnet_groups: int = 32, resnet_pre_norm: bool = True, attn_num_head_channels=1, cross_attention_dim=1280, output_scale_factor=1.0, downsample_padding=1, add_downsample=True, dual_cross_attention=False, use_linear_projection=False, only_cross_attention=False, upcast_attention=False, 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)

Bases: torch.nn.Module

Down block built by 3D cross attention.

forward(hidden_states, temb=None, encoder_hidden_states=None, attention_mask=None)

forward with hidden states.

class mmagic.models.editors.animatediff.unet_block.DownBlock3D(in_channels: int, out_channels: int, temb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-06, 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)

Bases: torch.nn.Module

Down block built by 3D resnet.

forward(hidden_states, temb=None, encoder_hidden_states=None)

forward with hidden states.

class mmagic.models.editors.animatediff.unet_block.CrossAttnUpBlock3D(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-06, resnet_time_scale_shift: str = 'default', resnet_act_fn: str = 'swish', resnet_groups: int = 32, resnet_pre_norm: bool = True, attn_num_head_channels=1, cross_attention_dim=1280, output_scale_factor=1.0, add_upsample=True, dual_cross_attention=False, use_linear_projection=False, only_cross_attention=False, upcast_attention=False, 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)

Bases: torch.nn.Module

Up block built by 3D cross attention.

forward(hidden_states, res_hidden_states_tuple, temb=None, encoder_hidden_states=None, upsample_size=None, attention_mask=None)

forward with hidden states and res hidden states.

class mmagic.models.editors.animatediff.unet_block.UpBlock3D(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-06, 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_motion_module=None, motion_module_type=None, motion_module_kwargs=None, use_inflated_groupnorm=None)

Bases: torch.nn.Module

Up block built by 3D resnet.

forward(hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, encoder_hidden_states=None)

forward with hidden states and res hidden states.