Source code for mmagic.models.editors.animatediff.motion_module
# Copyright (c) OpenMMLab. All rights reserved.
# Adapted from https://github.com/huggingface/diffusers
import math
from dataclasses import dataclass
from typing import Optional
import torch
import torch.nn.functional as F
from diffusers.utils import BaseOutput
from diffusers.utils.import_utils import is_xformers_available
from einops import rearrange, repeat
from torch import nn
from mmagic.models.editors.ddpm.attention import GEGLU, ApproximateGELU
from .attention_3d import CrossAttention
[docs]def zero_module(module):
"""Zero out the parameters of a module and return it."""
for p in module.parameters():
p.detach().zero_()
return module
@dataclass
[docs]class TemporalTransformer3DModelOutput(BaseOutput):
"""Output of TemporalTransformer3DModel."""
if is_xformers_available():
import xformers
import xformers.ops
else:
[docs]def get_motion_module(in_channels, motion_module_type: str,
motion_module_kwargs: dict):
"""Get motion module."""
if motion_module_type == 'Vanilla':
return VanillaTemporalModule(
in_channels=in_channels,
**motion_module_kwargs,
)
else:
raise ValueError
[docs]class VanillaTemporalModule(nn.Module):
"""Module which uses transformer to handle 3d motion."""
def __init__(
self,
in_channels,
num_attention_heads=8,
num_transformer_block=2,
attention_block_types=('Temporal_Self', 'Temporal_Self'),
cross_frame_attention_mode=None,
temporal_position_encoding=False,
temporal_position_encoding_max_len=24,
temporal_attention_dim_div=1,
zero_initialize=True,
):
super().__init__()
temp_pos_max_len = temporal_position_encoding_max_len
self.temporal_transformer = TemporalTransformer3DModel(
in_channels=in_channels,
num_attention_heads=num_attention_heads,
attention_head_dim=in_channels // num_attention_heads //
temporal_attention_dim_div,
num_layers=num_transformer_block,
attention_block_types=attention_block_types,
cross_frame_attention_mode=cross_frame_attention_mode,
temporal_position_encoding=temporal_position_encoding,
temporal_position_encoding_max_len=temp_pos_max_len,
)
if zero_initialize:
self.temporal_transformer.proj_out = zero_module(
self.temporal_transformer.proj_out)
[docs] def forward(self,
input_tensor,
temb,
encoder_hidden_states,
attention_mask=None,
anchor_frame_idx=None):
"""forward with sample."""
hidden_states = input_tensor
hidden_states = self.temporal_transformer(hidden_states,
encoder_hidden_states,
attention_mask)
output = hidden_states
return output
[docs]class TemporalTransformer3DModel(nn.Module):
"""Module which uses implement 3D Transformer."""
def __init__(
self,
in_channels,
num_attention_heads,
attention_head_dim,
num_layers,
attention_block_types=(
'Temporal_Self',
'Temporal_Self',
),
dropout=0.0,
norm_num_groups=32,
cross_attention_dim=768,
activation_fn='geglu',
attention_bias=False,
upcast_attention=False,
cross_frame_attention_mode=None,
temporal_position_encoding=False,
temporal_position_encoding_max_len=24,
):
super().__init__()
inner_dim = num_attention_heads * attention_head_dim
self.norm = torch.nn.GroupNorm(
num_groups=norm_num_groups,
num_channels=in_channels,
eps=1e-6,
affine=True)
self.proj_in = nn.Linear(in_channels, inner_dim)
temp_pos_max_len = temporal_position_encoding_max_len
self.transformer_blocks = nn.ModuleList([
TemporalTransformerBlock(
dim=inner_dim,
num_attention_heads=num_attention_heads,
attention_head_dim=attention_head_dim,
attention_block_types=attention_block_types,
dropout=dropout,
norm_num_groups=norm_num_groups,
cross_attention_dim=cross_attention_dim,
activation_fn=activation_fn,
attention_bias=attention_bias,
upcast_attention=upcast_attention,
cross_frame_attention_mode=cross_frame_attention_mode,
temporal_position_encoding=temporal_position_encoding,
temporal_position_encoding_max_len=temp_pos_max_len,
) for d in range(num_layers)
])
self.proj_out = nn.Linear(inner_dim, in_channels)
[docs] def forward(self,
hidden_states,
encoder_hidden_states=None,
attention_mask=None):
"""forward with hidden states, encoder_hidden_states and
attention_mask."""
assert hidden_states.dim(
) == 5, f'{"Expected hidden_states to have ndim=5, "}'
f'but got ndim={hidden_states.dim()}.'
video_length = hidden_states.shape[2]
hidden_states = rearrange(hidden_states, 'b c f h w -> (b f) c h w')
batch, channel, height, weight = hidden_states.shape
residual = hidden_states
hidden_states = self.norm(hidden_states)
inner_dim = hidden_states.shape[1]
hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(
batch, height * weight, inner_dim)
hidden_states = self.proj_in(hidden_states)
# Transformer Blocks
for block in self.transformer_blocks:
hidden_states = block(
hidden_states,
encoder_hidden_states=encoder_hidden_states,
video_length=video_length)
# output
hidden_states = self.proj_out(hidden_states)
hidden_states = hidden_states.reshape(batch, height,
weight, inner_dim).permute(
0, 3, 1, 2).contiguous()
output = hidden_states + residual
output = rearrange(output, '(b f) c h w -> b c f h w', f=video_length)
return output
[docs]class TemporalTransformerBlock(nn.Module):
"""Module which is a component of Temporal 3D Transformer."""
def __init__(
self,
dim,
num_attention_heads,
attention_head_dim,
attention_block_types=(
'Temporal_Self',
'Temporal_Self',
),
dropout=0.0,
norm_num_groups=32,
cross_attention_dim=768,
activation_fn='geglu',
attention_bias=False,
upcast_attention=False,
cross_frame_attention_mode=None,
temporal_position_encoding=False,
temporal_position_encoding_max_len=24,
):
super().__init__()
attention_blocks = []
norms = []
temp_pos_max_len = temporal_position_encoding_max_len
for block_name in attention_block_types:
attention_blocks.append(
VersatileAttention(
attention_mode=block_name.split('_')[0],
cross_attention_dim=cross_attention_dim
if block_name.endswith('_Cross') else None,
query_dim=dim,
heads=num_attention_heads,
dim_head=attention_head_dim,
dropout=dropout,
bias=attention_bias,
upcast_attention=upcast_attention,
cross_frame_attention_mode=cross_frame_attention_mode,
temporal_position_encoding=temporal_position_encoding,
temporal_position_encoding_max_len=temp_pos_max_len,
))
norms.append(nn.LayerNorm(dim))
self.attention_blocks = nn.ModuleList(attention_blocks)
self.norms = nn.ModuleList(norms)
self.ff = FeedForward(
dim, dropout=dropout, activation_fn=activation_fn)
self.ff_norm = nn.LayerNorm(dim)
[docs] def forward(self,
hidden_states,
encoder_hidden_states=None,
attention_mask=None,
video_length=None):
"""forward with hidden states, encoder_hidden_states and
attention_mask."""
for attention_block, norm in zip(self.attention_blocks, self.norms):
norm_hidden_states = norm(hidden_states)
hidden_states = attention_block(
norm_hidden_states,
encoder_hidden_states=encoder_hidden_states
if attention_block.is_cross_attention else None,
video_length=video_length,
) + hidden_states
hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states
output = hidden_states
return output
[docs]class PositionalEncoding(nn.Module):
"""a implementation of PositionEncoding."""
def __init__(self, d_model, dropout=0., max_len=24):
super().__init__()
self.dropout = nn.Dropout(p=dropout)
position = torch.arange(max_len).unsqueeze(1)
div_term = torch.exp(
torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model))
pe = torch.zeros(1, max_len, d_model)
pe[0, :, 0::2] = torch.sin(position * div_term)
pe[0, :, 1::2] = torch.cos(position * div_term)
self.register_buffer('pe', pe)
[docs] def forward(self, x):
"""forward function."""
x = x + self.pe[:, :x.size(1)]
return self.dropout(x)
[docs]class VersatileAttention(CrossAttention):
"""a implementation of VersatileAttention."""
def __init__(self,
attention_mode=None,
cross_frame_attention_mode=None,
temporal_position_encoding=False,
temporal_position_encoding_max_len=24,
*args,
**kwargs):
super().__init__(*args, **kwargs)
assert attention_mode == 'Temporal'
self.attention_mode = attention_mode
self.is_cross_attention = kwargs['cross_attention_dim'] is not None
self.pos_encoder = PositionalEncoding(
kwargs['query_dim'],
dropout=0.,
max_len=temporal_position_encoding_max_len) if (
temporal_position_encoding
and attention_mode == 'Temporal') else None
self._use_memory_efficient_attention_xformers = False
[docs] def extra_repr(self):
"""return module information."""
return f'(Module Info) Attention_Mode: {self.attention_mode},\
Is_Cross_Attention: {self.is_cross_attention}'
[docs] def reshape_heads_to_batch_dim(self, tensor):
"""reshape heads num to batch dim."""
batch_size, seq_len, dim = tensor.shape
head_size = self.heads
tensor = tensor.reshape(batch_size, seq_len, head_size,
dim // head_size)
tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size * head_size,
seq_len, dim // head_size)
return tensor
[docs] def reshape_batch_dim_to_heads(self, tensor):
"""reshape batch dim to heads num."""
batch_size, seq_len, dim = tensor.shape
head_size = self.heads
tensor = tensor.reshape(batch_size // head_size, head_size, seq_len,
dim)
tensor = tensor.permute(0, 2, 1, 3).reshape(batch_size // head_size,
seq_len, dim * head_size)
return tensor
[docs] def _memory_efficient_attention_xformers(self, query, key, value,
attention_mask):
"""use xformers to save memory."""
# TODO attention_mask
query = query.contiguous()
key = key.contiguous()
value = value.contiguous()
hidden_states = xformers.ops.memory_efficient_attention(
query, key, value, attn_bias=attention_mask)
hidden_states = self.reshape_batch_dim_to_heads(hidden_states)
return hidden_states
[docs] def forward(self,
hidden_states,
encoder_hidden_states=None,
attention_mask=None,
video_length=None):
"""forward with hidden states, encoder_hidden_states and
attention_mask."""
batch_size, sequence_length, _ = hidden_states.shape
if self.attention_mode == 'Temporal':
d = hidden_states.shape[1]
hidden_states = rearrange(
hidden_states, '(b f) d c -> (b d) f c', f=video_length)
if self.pos_encoder is not None:
hidden_states = self.pos_encoder(hidden_states)
encoder_hidden_states = repeat(
encoder_hidden_states, 'b n c -> (b d) n c', d=d
) if encoder_hidden_states is not None else encoder_hidden_states
else:
raise NotImplementedError
encoder_hidden_states = encoder_hidden_states
if self.group_norm is not None:
hidden_states = self.group_norm(hidden_states.transpose(
1, 2)).transpose(1, 2)
query = self.to_q(hidden_states)
dim = query.shape[-1]
query = self.reshape_heads_to_batch_dim(query)
if self.added_kv_proj_dim is not None:
raise NotImplementedError
encoder_hidden_states = encoder_hidden_states \
if encoder_hidden_states is not None else hidden_states
key = self.to_k(encoder_hidden_states)
value = self.to_v(encoder_hidden_states)
key = self.reshape_heads_to_batch_dim(key)
value = self.reshape_heads_to_batch_dim(value)
if attention_mask is not None:
if attention_mask.shape[-1] != query.shape[1]:
target_length = query.shape[1]
attention_mask = F.pad(
attention_mask, (0, target_length), value=0.0)
attention_mask = attention_mask.repeat_interleave(
self.heads, dim=0)
# attention, what we cannot get enough of
if self._use_memory_efficient_attention_xformers:
hidden_states = self._memory_efficient_attention_xformers(
query, key, value, attention_mask)
# Some versions of xformers return output in fp32,
# cast it back to the dtype of the input
hidden_states = hidden_states.to(query.dtype)
else:
if self._slice_size is None or query.shape[
0] // self._slice_size == 1:
hidden_states = self._attention(query, key, value,
attention_mask)
else:
hidden_states = self._sliced_attention(query, key, value,
sequence_length, dim,
attention_mask)
# linear proj
hidden_states = self.to_out[0](hidden_states)
# dropout
hidden_states = self.to_out[1](hidden_states)
if self.attention_mode == 'Temporal':
hidden_states = rearrange(
hidden_states, '(b d) f c -> (b f) d c', d=d)
return hidden_states
[docs]class FeedForward(nn.Module):
r"""
A feed-forward layer.
Parameters:
dim (`int`): The number of channels in the input.
dim_out (`int`, *optional*):
The number of channels in the output. If not given, defaults to `dim`.
mult (`int`, *optional*, defaults to 4):
The multiplier to use for the hidden dimension.
dropout (`float`, *optional*, defaults to 0.0):
The dropout probability to use.
activation_fn (`str`, *optional*, defaults to `"geglu"`):
Activation function to be used in feed-forward.
"""
def __init__(
self,
dim: int,
dim_out: Optional[int] = None,
mult: int = 4,
dropout: float = 0.0,
activation_fn: str = 'geglu',
):
super().__init__()
inner_dim = int(dim * mult)
dim_out = dim_out if dim_out is not None else dim
if activation_fn == 'gelu':
act_fn = GELU(dim, inner_dim)
elif activation_fn == 'geglu':
act_fn = GEGLU(dim, inner_dim)
elif activation_fn == 'geglu-approximate':
act_fn = ApproximateGELU(dim, inner_dim)
self.net = nn.ModuleList([])
# project in
self.net.append(act_fn)
# project dropout
self.net.append(nn.Dropout(dropout))
# project out
self.net.append(nn.Linear(inner_dim, dim_out))
[docs] def forward(self, hidden_states):
for module in self.net:
hidden_states = module(hidden_states)
return hidden_states
[docs]class GELU(nn.Module):
r"""
GELU activation function
"""
def __init__(self, dim_in: int, dim_out: int):
super().__init__()
self.proj = nn.Linear(dim_in, dim_out)
[docs] def gelu(self, gate):
if gate.device.type != 'mps':
return F.gelu(gate)
# mps: gelu is not implemented for float16
return F.gelu(gate.to(dtype=torch.float32)).to(dtype=gate.dtype)
[docs] def forward(self, hidden_states):
hidden_states = self.proj(hidden_states)
hidden_states = self.gelu(hidden_states)
return hidden_states