multimodalart/EchoMimic-zero
8
1# Adapt from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/models/motion_module.py2import math3from dataclasses import dataclass4from typing import Callable, Optional5 6import torch7from diffusers.models.attention import FeedForward8from diffusers.models.attention_processor import Attention, AttnProcessor9from diffusers.utils import BaseOutput10from diffusers.utils.import_utils import is_xformers_available11from einops import rearrange, repeat12from torch import nn13 14 15def zero_module(module):16 # Zero out the parameters of a module and return it.17 for p in module.parameters():18 p.detach().zero_()19 return module20 21 22@dataclass23class TemporalTransformer3DModelOutput(BaseOutput):24 sample: torch.FloatTensor25 26 27if is_xformers_available():28 import xformers29 import xformers.ops30else:31 xformers = None32 33 34def get_motion_module(in_channels, motion_module_type: str, motion_module_kwargs: dict):35 if motion_module_type == "Vanilla":36 return VanillaTemporalModule(37 in_channels=in_channels,38 **motion_module_kwargs,39 )40 else:41 raise ValueError42 43 44class VanillaTemporalModule(nn.Module):45 def __init__(46 self,47 in_channels,48 num_attention_heads=8,49 num_transformer_block=2,50 attention_block_types=("Temporal_Self", "Temporal_Self"),51 cross_frame_attention_mode=None,52 temporal_position_encoding=False,53 temporal_position_encoding_max_len=24,54 temporal_attention_dim_div=1,55 zero_initialize=True,56 ):57 super().__init__()58 59 self.temporal_transformer = TemporalTransformer3DModel(60 in_channels=in_channels,61 num_attention_heads=num_attention_heads,62 attention_head_dim=in_channels63 // num_attention_heads64 // temporal_attention_dim_div,65 num_layers=num_transformer_block,66 attention_block_types=attention_block_types,67 cross_frame_attention_mode=cross_frame_attention_mode,68 temporal_position_encoding=temporal_position_encoding,69 temporal_position_encoding_max_len=temporal_position_encoding_max_len,70 )71 72 if zero_initialize:73 self.temporal_transformer.proj_out = zero_module(74 self.temporal_transformer.proj_out75 )76 77 def forward(78 self,79 input_tensor,80 temb,81 encoder_hidden_states,82 attention_mask=None,83 anchor_frame_idx=None,84 ):85 hidden_states = input_tensor86 hidden_states = self.temporal_transformer(87 hidden_states, encoder_hidden_states, attention_mask88 )89 90 output = hidden_states91 return output92 93 94class TemporalTransformer3DModel(nn.Module):95 def __init__(96 self,97 in_channels,98 num_attention_heads,99 attention_head_dim,100 num_layers,101 attention_block_types=(102 "Temporal_Self",103 "Temporal_Self",104 ),105 dropout=0.0,106 norm_num_groups=32,107 cross_attention_dim=768,108 activation_fn="geglu",109 attention_bias=False,110 upcast_attention=False,111 cross_frame_attention_mode=None,112 temporal_position_encoding=False,113 temporal_position_encoding_max_len=24,114 ):115 super().__init__()116 117 inner_dim = num_attention_heads * attention_head_dim118 119 self.norm = torch.nn.GroupNorm(120 num_groups=norm_num_groups, num_channels=in_channels, eps=1e-6, affine=True121 )122 self.proj_in = nn.Linear(in_channels, inner_dim)123 124 self.transformer_blocks = nn.ModuleList(125 [126 TemporalTransformerBlock(127 dim=inner_dim,128 num_attention_heads=num_attention_heads,129 attention_head_dim=attention_head_dim,130 attention_block_types=attention_block_types,131 dropout=dropout,132 norm_num_groups=norm_num_groups,133 cross_attention_dim=cross_attention_dim,134 activation_fn=activation_fn,135 attention_bias=attention_bias,136 upcast_attention=upcast_attention,137 cross_frame_attention_mode=cross_frame_attention_mode,138 temporal_position_encoding=temporal_position_encoding,139 temporal_position_encoding_max_len=temporal_position_encoding_max_len,140 )141 for d in range(num_layers)142 ]143 )144 self.proj_out = nn.Linear(inner_dim, in_channels)145 146 def forward(self, hidden_states, encoder_hidden_states=None, attention_mask=None):147 assert (148 hidden_states.dim() == 5149 ), f"Expected hidden_states to have ndim=5, but got ndim={hidden_states.dim()}."150 video_length = hidden_states.shape[2]151 hidden_states = rearrange(hidden_states, "b c f h w -> (b f) c h w")152 153 batch, channel, height, weight = hidden_states.shape154 residual = hidden_states155 156 hidden_states = self.norm(hidden_states)157 inner_dim = hidden_states.shape[1]158 hidden_states = hidden_states.permute(0, 2, 3, 1).reshape(159 batch, height * weight, inner_dim160 )161 hidden_states = self.proj_in(hidden_states)162 163 # Transformer Blocks164 for block in self.transformer_blocks:165 hidden_states = block(166 hidden_states,167 encoder_hidden_states=encoder_hidden_states,168 video_length=video_length,169 )170 171 # output172 hidden_states = self.proj_out(hidden_states)173 hidden_states = (174 hidden_states.reshape(batch, height, weight, inner_dim)175 .permute(0, 3, 1, 2)176 .contiguous()177 )178 179 output = hidden_states + residual180 output = rearrange(output, "(b f) c h w -> b c f h w", f=video_length)181 182 return output183 184 185class TemporalTransformerBlock(nn.Module):186 def __init__(187 self,188 dim,189 num_attention_heads,190 attention_head_dim,191 attention_block_types=(192 "Temporal_Self",193 "Temporal_Self",194 ),195 dropout=0.0,196 norm_num_groups=32,197 cross_attention_dim=768,198 activation_fn="geglu",199 attention_bias=False,200 upcast_attention=False,201 cross_frame_attention_mode=None,202 temporal_position_encoding=False,203 temporal_position_encoding_max_len=24,204 ):205 super().__init__()206 207 attention_blocks = []208 norms = []209 210 for block_name in attention_block_types:211 attention_blocks.append(212 VersatileAttention(213 attention_mode=block_name.split("_")[0],214 cross_attention_dim=cross_attention_dim215 if block_name.endswith("_Cross")216 else None,217 query_dim=dim,218 heads=num_attention_heads,219 dim_head=attention_head_dim,220 dropout=dropout,221 bias=attention_bias,222 upcast_attention=upcast_attention,223 cross_frame_attention_mode=cross_frame_attention_mode,224 temporal_position_encoding=temporal_position_encoding,225 temporal_position_encoding_max_len=temporal_position_encoding_max_len,226 )227 )228 norms.append(nn.LayerNorm(dim))229 230 self.attention_blocks = nn.ModuleList(attention_blocks)231 self.norms = nn.ModuleList(norms)232 233 self.ff = FeedForward(dim, dropout=dropout, activation_fn=activation_fn)234 self.ff_norm = nn.LayerNorm(dim)235 236 def forward(237 self,238 hidden_states,239 encoder_hidden_states=None,240 attention_mask=None,241 video_length=None,242 ):243 for attention_block, norm in zip(self.attention_blocks, self.norms):244 norm_hidden_states = norm(hidden_states)245 hidden_states = (246 attention_block(247 norm_hidden_states,248 encoder_hidden_states=encoder_hidden_states249 if attention_block.is_cross_attention250 else None,251 video_length=video_length,252 )253 + hidden_states254 )255 256 hidden_states = self.ff(self.ff_norm(hidden_states)) + hidden_states257 258 output = hidden_states259 return output260 261 262class PositionalEncoding(nn.Module):263 def __init__(self, d_model, dropout=0.0, max_len=24):264 super().__init__()265 self.dropout = nn.Dropout(p=dropout)266 position = torch.arange(max_len).unsqueeze(1)267 div_term = torch.exp(268 torch.arange(0, d_model, 2) * (-math.log(10000.0) / d_model)269 )270 pe = torch.zeros(1, max_len, d_model)271 pe[0, :, 0::2] = torch.sin(position * div_term)272 pe[0, :, 1::2] = torch.cos(position * div_term)273 self.register_buffer("pe", pe)274 275 def forward(self, x):276 x = x + self.pe[:, : x.size(1)]277 return self.dropout(x)278 279 280class VersatileAttention(Attention):281 def __init__(282 self,283 attention_mode=None,284 cross_frame_attention_mode=None,285 temporal_position_encoding=False,286 temporal_position_encoding_max_len=24,287 *args,288 **kwargs,289 ):290 super().__init__(*args, **kwargs)291 assert attention_mode == "Temporal"292 293 self.attention_mode = attention_mode294 self.is_cross_attention = kwargs["cross_attention_dim"] is not None295 296 self.pos_encoder = (297 PositionalEncoding(298 kwargs["query_dim"],299 dropout=0.0,300 max_len=temporal_position_encoding_max_len,301 )302 if (temporal_position_encoding and attention_mode == "Temporal")303 else None304 )305 306 def extra_repr(self):307 return f"(Module Info) Attention_Mode: {self.attention_mode}, Is_Cross_Attention: {self.is_cross_attention}"308 309 def set_use_memory_efficient_attention_xformers(310 self,311 use_memory_efficient_attention_xformers: bool,312 attention_op: Optional[Callable] = None,313 ):314 if use_memory_efficient_attention_xformers:315 if not is_xformers_available():316 raise ModuleNotFoundError(317 (318 "Refer to https://github.com/facebookresearch/xformers for more information on how to install"319 " xformers"320 ),321 name="xformers",322 )323 elif not torch.cuda.is_available():324 raise ValueError(325 "torch.cuda.is_available() should be True but is False. xformers' memory efficient attention is"326 " only available for GPU "327 )328 else:329 try:330 # Make sure we can run the memory efficient attention331 _ = xformers.ops.memory_efficient_attention(332 torch.randn((1, 2, 40), device="cuda"),333 torch.randn((1, 2, 40), device="cuda"),334 torch.randn((1, 2, 40), device="cuda"),335 )336 except Exception as e:337 raise e338 339 # XFormersAttnProcessor corrupts video generation and work with Pytorch 1.13.340 # Pytorch 2.0.1 AttnProcessor works the same as XFormersAttnProcessor in Pytorch 1.13.341 # You don't need XFormersAttnProcessor here.342 # processor = XFormersAttnProcessor(343 # attention_op=attention_op,344 # )345 processor = AttnProcessor()346 else:347 processor = AttnProcessor()348 349 self.set_processor(processor)350 351 def forward(352 self,353 hidden_states,354 encoder_hidden_states=None,355 attention_mask=None,356 video_length=None,357 **cross_attention_kwargs,358 ):359 if self.attention_mode == "Temporal":360 d = hidden_states.shape[1] # d means HxW361 hidden_states = rearrange(362 hidden_states, "(b f) d c -> (b d) f c", f=video_length363 )364 365 if self.pos_encoder is not None:366 hidden_states = self.pos_encoder(hidden_states)367 368 encoder_hidden_states = (369 repeat(encoder_hidden_states, "b n c -> (b d) n c", d=d)370 if encoder_hidden_states is not None371 else encoder_hidden_states372 )373 374 else:375 raise NotImplementedError376 377 hidden_states = self.processor(378 self,379 hidden_states,380 encoder_hidden_states=encoder_hidden_states,381 attention_mask=attention_mask,382 **cross_attention_kwargs,383 )384 385 if self.attention_mode == "Temporal":386 hidden_states = rearrange(hidden_states, "(b d) f c -> (b f) d c", d=d)387 388 return hidden_states389