343 lines
12 KiB
Python
343 lines
12 KiB
Python
from enum import Enum
|
|
from typing import Optional
|
|
|
|
import math
|
|
import torch
|
|
from torch import nn
|
|
from einops import rearrange
|
|
|
|
from ldm_patched.ldm.modules.attention import FeedForward
|
|
from ldm_patched.modules.ops import disable_weight_init
|
|
|
|
|
|
class MotionModuleType(Enum):
|
|
AnimateDiffV1 = "AnimateDiff V1, Yuwei Guo, Shanghai AI Lab"
|
|
AnimateDiffV2 = "AnimateDiff V2, Yuwei Guo, Shanghai AI Lab"
|
|
AnimateDiffV3 = "AnimateDiff V3, Yuwei Guo, Shanghai AI Lab"
|
|
AnimateDiffXL = "AnimateDiff SDXL, Yuwei Guo, Shanghai AI Lab"
|
|
SparseCtrl = "SparseCtrl, Yuwei Guo, Shanghai AI Lab"
|
|
HotShotXL = "HotShot-XL, John Mullan, Natural Synthetics Inc"
|
|
|
|
|
|
@staticmethod
|
|
def get_mm_type(state_dict: dict[str, torch.Tensor]):
|
|
keys = list(state_dict.keys())
|
|
if any(["mid_block" in k for k in keys]):
|
|
return MotionModuleType.AnimateDiffV2
|
|
elif any(["down_blocks.3" in k for k in keys]):
|
|
if 32 in next((state_dict[key] for key in state_dict if 'pe' in key), None).shape:
|
|
return MotionModuleType.AnimateDiffV3
|
|
else:
|
|
return MotionModuleType.AnimateDiffV1
|
|
else:
|
|
if 32 in next((state_dict[key] for key in state_dict if 'pe' in key), None).shape:
|
|
return MotionModuleType.AnimateDiffXL
|
|
else:
|
|
return MotionModuleType.HotShotXL
|
|
|
|
|
|
def zero_module(module):
|
|
# Zero out the parameters of a module and return it.
|
|
for p in module.parameters():
|
|
p.detach().zero_()
|
|
return module
|
|
|
|
|
|
class MotionWrapper(nn.Module):
|
|
def __init__(self, mm_name: str, mm_hash: str, mm_type: MotionModuleType, operations = disable_weight_init):
|
|
super().__init__()
|
|
self.mm_name = mm_name
|
|
self.mm_type = mm_type
|
|
self.mm_hash = mm_hash
|
|
max_len = 24 if self.enable_gn_hack() else 32
|
|
in_channels = (320, 640, 1280) if self.is_xl else (320, 640, 1280, 1280)
|
|
self.down_blocks = nn.ModuleList([])
|
|
self.up_blocks = nn.ModuleList([])
|
|
for c in in_channels:
|
|
if mm_type in [MotionModuleType.SparseCtrl]:
|
|
self.down_blocks.append(MotionModule(c, num_mm=2, max_len=max_len, attention_block_types=("Temporal_Self", ), operations=operations))
|
|
else:
|
|
self.down_blocks.append(MotionModule(c, num_mm=2, max_len=max_len, operations=operations))
|
|
self.up_blocks.insert(0,MotionModule(c, num_mm=3, max_len=max_len, operations=operations))
|
|
if mm_type in [MotionModuleType.AnimateDiffV2]:
|
|
self.mid_block = MotionModule(1280, num_mm=1, max_len=max_len, operations=operations)
|
|
|
|
|
|
def enable_gn_hack(self):
|
|
return self.mm_type in [MotionModuleType.AnimateDiffV1, MotionModuleType.HotShotXL]
|
|
|
|
|
|
@property
|
|
def is_xl(self):
|
|
return self.mm_type in [MotionModuleType.AnimateDiffXL, MotionModuleType.HotShotXL]
|
|
|
|
|
|
@property
|
|
def is_adxl(self):
|
|
return self.mm_type == MotionModuleType.AnimateDiffXL
|
|
|
|
|
|
class MotionModule(nn.Module):
|
|
def __init__(self, in_channels, num_mm, max_len, attention_block_types=("Temporal_Self", "Temporal_Self"), operations = disable_weight_init):
|
|
super().__init__()
|
|
self.motion_modules = nn.ModuleList([
|
|
VanillaTemporalModule(
|
|
in_channels=in_channels,
|
|
temporal_position_encoding_max_len=max_len,
|
|
attention_block_types=attention_block_types,
|
|
operations=operations,)
|
|
for _ in range(num_mm)])
|
|
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
for mm in self.motion_modules:
|
|
x = mm(x)
|
|
return x
|
|
|
|
|
|
class VanillaTemporalModule(nn.Module):
|
|
def __init__(
|
|
self,
|
|
in_channels,
|
|
num_attention_heads = 8,
|
|
num_transformer_block = 1,
|
|
attention_block_types =( "Temporal_Self", "Temporal_Self" ),
|
|
temporal_position_encoding_max_len = 24,
|
|
temporal_attention_dim_div = 1,
|
|
zero_initialize = True,
|
|
operations = disable_weight_init,
|
|
):
|
|
super().__init__()
|
|
|
|
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,
|
|
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
|
operations=operations,
|
|
)
|
|
|
|
if zero_initialize:
|
|
self.temporal_transformer.proj_out = zero_module(self.temporal_transformer.proj_out)
|
|
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
return self.temporal_transformer(x)
|
|
|
|
|
|
class TemporalTransformer3DModel(nn.Module):
|
|
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,
|
|
activation_fn = "geglu",
|
|
attention_bias = False,
|
|
upcast_attention = False,
|
|
|
|
temporal_position_encoding_max_len = 24,
|
|
|
|
operations = disable_weight_init,
|
|
):
|
|
super().__init__()
|
|
|
|
inner_dim = num_attention_heads * attention_head_dim
|
|
|
|
self.norm = operations.GroupNorm(num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True)
|
|
self.proj_in = operations.Linear(in_channels, inner_dim)
|
|
|
|
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,
|
|
activation_fn=activation_fn,
|
|
attention_bias=attention_bias,
|
|
upcast_attention=upcast_attention,
|
|
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
|
operations=operations,
|
|
)
|
|
for _ in range(num_layers)
|
|
]
|
|
)
|
|
self.proj_out = operations.Linear(inner_dim, in_channels)
|
|
|
|
def forward(self, hidden_states: torch.Tensor):
|
|
_, _, height, _ = hidden_states.shape
|
|
residual = hidden_states
|
|
|
|
hidden_states = self.norm(hidden_states).type(hidden_states.dtype)
|
|
hidden_states = rearrange(hidden_states, "b c h w -> b (h w) c")
|
|
hidden_states = self.proj_in(hidden_states)
|
|
|
|
# Transformer Blocks
|
|
for block in self.transformer_blocks:
|
|
hidden_states = block(hidden_states)
|
|
|
|
# output
|
|
hidden_states = self.proj_out(hidden_states)
|
|
hidden_states = rearrange(hidden_states, "b (h w) c -> b c h w", h=height)
|
|
|
|
output = hidden_states + residual
|
|
return output
|
|
|
|
|
|
class TemporalTransformerBlock(nn.Module):
|
|
def __init__(
|
|
self,
|
|
dim,
|
|
num_attention_heads,
|
|
attention_head_dim,
|
|
attention_block_types = ( "Temporal_Self", "Temporal_Self", ),
|
|
dropout = 0.0,
|
|
activation_fn = "geglu",
|
|
attention_bias = False,
|
|
upcast_attention = False,
|
|
temporal_position_encoding_max_len = 24,
|
|
operations = disable_weight_init,
|
|
):
|
|
super().__init__()
|
|
|
|
attention_blocks = []
|
|
norms = []
|
|
|
|
for _ in attention_block_types:
|
|
attention_blocks.append(
|
|
VersatileAttention(
|
|
query_dim=dim,
|
|
heads=num_attention_heads,
|
|
dim_head=attention_head_dim,
|
|
dropout=dropout,
|
|
bias=attention_bias,
|
|
upcast_attention=upcast_attention,
|
|
temporal_position_encoding_max_len=temporal_position_encoding_max_len,
|
|
operations=operations,
|
|
)
|
|
)
|
|
norms.append(operations.LayerNorm(dim))
|
|
|
|
self.attention_blocks = nn.ModuleList(attention_blocks)
|
|
self.norms = nn.ModuleList(norms)
|
|
|
|
self.ff = FeedForward(dim, dropout=dropout, glu=(activation_fn=='geglu'), operations=operations)
|
|
self.ff_norm = operations.LayerNorm(dim)
|
|
|
|
|
|
def forward(self, hidden_states: torch.Tensor):
|
|
for attention_block, norm in zip(self.attention_blocks, self.norms):
|
|
norm_hidden_states = norm(hidden_states).type(hidden_states.dtype)
|
|
hidden_states = attention_block(norm_hidden_states) + hidden_states
|
|
|
|
hidden_states = self.ff(self.ff_norm(hidden_states).type(hidden_states.dtype)) + hidden_states
|
|
|
|
output = hidden_states
|
|
return output
|
|
|
|
|
|
class PositionalEncoding(nn.Module):
|
|
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)
|
|
|
|
def forward(self, x):
|
|
x = x + self.pe[:, :x.size(1)].type(x.dtype)
|
|
return self.dropout(x)
|
|
|
|
|
|
class CrossAttention(nn.Module):
|
|
r"""
|
|
A cross attention layer.
|
|
|
|
Parameters:
|
|
query_dim (`int`): The number of channels in the query.
|
|
cross_attention_dim (`int`, *optional*):
|
|
The number of channels in the encoder_hidden_states. If not given, defaults to `query_dim`.
|
|
heads (`int`, *optional*, defaults to 8): The number of heads to use for multi-head attention.
|
|
dim_head (`int`, *optional*, defaults to 64): The number of channels in each head.
|
|
dropout (`float`, *optional*, defaults to 0.0): The dropout probability to use.
|
|
bias (`bool`, *optional*, defaults to False):
|
|
Set to `True` for the query, key, and value linear layers to contain a bias parameter.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
query_dim: int,
|
|
cross_attention_dim: Optional[int] = None,
|
|
heads: int = 8,
|
|
dim_head: int = 64,
|
|
dropout: float = 0.0,
|
|
bias=False,
|
|
upcast_attention: bool = False,
|
|
upcast_softmax: bool = False,
|
|
operations = disable_weight_init,
|
|
):
|
|
super().__init__()
|
|
inner_dim = dim_head * heads
|
|
cross_attention_dim = cross_attention_dim if cross_attention_dim is not None else query_dim
|
|
self.upcast_attention = upcast_attention
|
|
self.upcast_softmax = upcast_softmax
|
|
self.scale = dim_head**-0.5
|
|
self.heads = heads
|
|
|
|
self.to_q = operations.Linear(query_dim, inner_dim, bias=bias)
|
|
self.to_k = operations.Linear(cross_attention_dim, inner_dim, bias=bias)
|
|
self.to_v = operations.Linear(cross_attention_dim, inner_dim, bias=bias)
|
|
|
|
self.to_out = nn.Sequential(operations.Linear(inner_dim, query_dim), nn.Dropout(dropout))
|
|
|
|
|
|
class VersatileAttention(CrossAttention):
|
|
def __init__(
|
|
self,
|
|
temporal_position_encoding_max_len = 24,
|
|
*args, **kwargs
|
|
):
|
|
super().__init__(*args, **kwargs)
|
|
|
|
self.pos_encoder = PositionalEncoding(
|
|
kwargs["query_dim"],
|
|
max_len=temporal_position_encoding_max_len)
|
|
|
|
|
|
def forward(self, x: torch.Tensor):
|
|
from scripts.animatediff_mm import mm_animatediff
|
|
video_length = mm_animatediff.ad_params.batch_size
|
|
|
|
d = x.shape[1]
|
|
x = rearrange(x, "(b f) d c -> (b d) f c", f=video_length)
|
|
x = self.pos_encoder(x)
|
|
|
|
q = self.to_q(x)
|
|
k = self.to_k(x)
|
|
v = self.to_v(x)
|
|
|
|
q, k, v = map(lambda t: rearrange(t, 'b s (h d) -> (b h) s d', h=self.heads), (q, k, v))
|
|
x = torch.nn.functional.scaled_dot_product_attention(q, k, v)
|
|
x = rearrange(x, '(b h) s d -> b s (h d)', h=self.heads)
|
|
|
|
x = self.to_out(x) # linear proj and dropout
|
|
x = rearrange(x, "(b d) f c -> (b f) d c", d=d)
|
|
|
|
return x
|