multimodalart/EchoMimic-zero
8
1# Adapted from https://github.com/guoyww/AnimateDiff/blob/main/animatediff/models/unet_blocks.py2 3from collections import OrderedDict4from dataclasses import dataclass5from os import PathLike6from pathlib import Path7from typing import Dict, List, Optional, Tuple, Union8 9import torch10import torch.nn as nn11import torch.utils.checkpoint12from diffusers.configuration_utils import ConfigMixin, register_to_config13from diffusers.models.attention_processor import AttentionProcessor14from diffusers.models.embeddings import TimestepEmbedding, Timesteps15from diffusers.models.modeling_utils import ModelMixin16from diffusers.utils import SAFETENSORS_WEIGHTS_NAME, WEIGHTS_NAME, BaseOutput, logging17from safetensors.torch import load_file18 19from .resnet import InflatedConv3d, InflatedGroupNorm20from .unet_3d_blocks import UNetMidBlock3DCrossAttn, get_down_block, get_up_block21 22logger = logging.get_logger(__name__) # pylint: disable=invalid-name23 24 25@dataclass26class UNet3DConditionOutput(BaseOutput):27 sample: torch.FloatTensor28 29 30class UNet3DConditionModel(ModelMixin, ConfigMixin):31 _supports_gradient_checkpointing = True32 33 @register_to_config34 def __init__(35 self,36 sample_size: Optional[int] = None,37 in_channels: int = 4,38 out_channels: int = 4,39 center_input_sample: bool = False,40 flip_sin_to_cos: bool = True,41 freq_shift: int = 0,42 down_block_types: Tuple[str] = (43 "CrossAttnDownBlock3D",44 "CrossAttnDownBlock3D",45 "CrossAttnDownBlock3D",46 "DownBlock3D",47 ),48 mid_block_type: str = "UNetMidBlock3DCrossAttn",49 up_block_types: Tuple[str] = (50 "UpBlock3D",51 "CrossAttnUpBlock3D",52 "CrossAttnUpBlock3D",53 "CrossAttnUpBlock3D",54 ),55 only_cross_attention: Union[bool, Tuple[bool]] = False,56 block_out_channels: Tuple[int] = (320, 640, 1280, 1280),57 layers_per_block: int = 2,58 downsample_padding: int = 1,59 mid_block_scale_factor: float = 1,60 act_fn: str = "silu",61 norm_num_groups: int = 32,62 norm_eps: float = 1e-5,63 cross_attention_dim: int = 1280,64 attention_head_dim: Union[int, Tuple[int]] = 8,65 dual_cross_attention: bool = False,66 use_linear_projection: bool = False,67 class_embed_type: Optional[str] = None,68 num_class_embeds: Optional[int] = None,69 upcast_attention: bool = False,70 resnet_time_scale_shift: str = "default",71 use_inflated_groupnorm=False,72 # Additional73 use_motion_module=False,74 motion_module_resolutions=(1, 2, 4, 8),75 motion_module_mid_block=False,76 motion_module_decoder_only=False,77 motion_module_type=None,78 motion_module_kwargs={},79 unet_use_cross_frame_attention=None,80 unet_use_temporal_attention=None,81 ):82 super().__init__()83 84 self.sample_size = sample_size85 time_embed_dim = block_out_channels[0] * 486 87 # input88 self.conv_in = InflatedConv3d(89 in_channels, block_out_channels[0], kernel_size=3, padding=(1, 1)90 )91 92 # time93 self.time_proj = Timesteps(block_out_channels[0], flip_sin_to_cos, freq_shift)94 timestep_input_dim = block_out_channels[0]95 96 self.time_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)97 98 # class embedding99 if class_embed_type is None and num_class_embeds is not None:100 self.class_embedding = nn.Embedding(num_class_embeds, time_embed_dim)101 elif class_embed_type == "timestep":102 self.class_embedding = TimestepEmbedding(timestep_input_dim, time_embed_dim)103 elif class_embed_type == "identity":104 self.class_embedding = nn.Identity(time_embed_dim, time_embed_dim)105 else:106 self.class_embedding = None107 108 self.down_blocks = nn.ModuleList([])109 self.mid_block = None110 self.up_blocks = nn.ModuleList([])111 112 if isinstance(only_cross_attention, bool):113 only_cross_attention = [only_cross_attention] * len(down_block_types)114 115 if isinstance(attention_head_dim, int):116 attention_head_dim = (attention_head_dim,) * len(down_block_types)117 118 # down119 output_channel = block_out_channels[0]120 for i, down_block_type in enumerate(down_block_types):121 res = 2**i122 input_channel = output_channel123 output_channel = block_out_channels[i]124 is_final_block = i == len(block_out_channels) - 1125 126 down_block = get_down_block(127 down_block_type,128 num_layers=layers_per_block,129 in_channels=input_channel,130 out_channels=output_channel,131 temb_channels=time_embed_dim,132 add_downsample=not is_final_block,133 resnet_eps=norm_eps,134 resnet_act_fn=act_fn,135 resnet_groups=norm_num_groups,136 cross_attention_dim=cross_attention_dim,137 attn_num_head_channels=attention_head_dim[i],138 downsample_padding=downsample_padding,139 dual_cross_attention=dual_cross_attention,140 use_linear_projection=use_linear_projection,141 only_cross_attention=only_cross_attention[i],142 upcast_attention=upcast_attention,143 resnet_time_scale_shift=resnet_time_scale_shift,144 unet_use_cross_frame_attention=unet_use_cross_frame_attention,145 unet_use_temporal_attention=unet_use_temporal_attention,146 use_inflated_groupnorm=use_inflated_groupnorm,147 use_motion_module=use_motion_module148 and (res in motion_module_resolutions)149 and (not motion_module_decoder_only),150 motion_module_type=motion_module_type,151 motion_module_kwargs=motion_module_kwargs,152 )153 self.down_blocks.append(down_block)154 155 # mid156 if mid_block_type == "UNetMidBlock3DCrossAttn":157 self.mid_block = UNetMidBlock3DCrossAttn(158 in_channels=block_out_channels[-1],159 temb_channels=time_embed_dim,160 resnet_eps=norm_eps,161 resnet_act_fn=act_fn,162 output_scale_factor=mid_block_scale_factor,163 resnet_time_scale_shift=resnet_time_scale_shift,164 cross_attention_dim=cross_attention_dim,165 attn_num_head_channels=attention_head_dim[-1],166 resnet_groups=norm_num_groups,167 dual_cross_attention=dual_cross_attention,168 use_linear_projection=use_linear_projection,169 upcast_attention=upcast_attention,170 unet_use_cross_frame_attention=unet_use_cross_frame_attention,171 unet_use_temporal_attention=unet_use_temporal_attention,172 use_inflated_groupnorm=use_inflated_groupnorm,173 use_motion_module=use_motion_module and motion_module_mid_block,174 motion_module_type=motion_module_type,175 motion_module_kwargs=motion_module_kwargs,176 )177 else:178 raise ValueError(f"unknown mid_block_type : {mid_block_type}")179 180 # count how many layers upsample the videos181 self.num_upsamplers = 0182 183 # up184 reversed_block_out_channels = list(reversed(block_out_channels))185 reversed_attention_head_dim = list(reversed(attention_head_dim))186 only_cross_attention = list(reversed(only_cross_attention))187 output_channel = reversed_block_out_channels[0]188 for i, up_block_type in enumerate(up_block_types):189 res = 2 ** (3 - i)190 is_final_block = i == len(block_out_channels) - 1191 192 prev_output_channel = output_channel193 output_channel = reversed_block_out_channels[i]194 input_channel = reversed_block_out_channels[195 min(i + 1, len(block_out_channels) - 1)196 ]197 198 # add upsample block for all BUT final layer199 if not is_final_block:200 add_upsample = True201 self.num_upsamplers += 1202 else:203 add_upsample = False204 205 up_block = get_up_block(206 up_block_type,207 num_layers=layers_per_block + 1,208 in_channels=input_channel,209 out_channels=output_channel,210 prev_output_channel=prev_output_channel,211 temb_channels=time_embed_dim,212 add_upsample=add_upsample,213 resnet_eps=norm_eps,214 resnet_act_fn=act_fn,215 resnet_groups=norm_num_groups,216 cross_attention_dim=cross_attention_dim,217 attn_num_head_channels=reversed_attention_head_dim[i],218 dual_cross_attention=dual_cross_attention,219 use_linear_projection=use_linear_projection,220 only_cross_attention=only_cross_attention[i],221 upcast_attention=upcast_attention,222 resnet_time_scale_shift=resnet_time_scale_shift,223 unet_use_cross_frame_attention=unet_use_cross_frame_attention,224 unet_use_temporal_attention=unet_use_temporal_attention,225 use_inflated_groupnorm=use_inflated_groupnorm,226 use_motion_module=use_motion_module227 and (res in motion_module_resolutions),228 motion_module_type=motion_module_type,229 motion_module_kwargs=motion_module_kwargs,230 )231 self.up_blocks.append(up_block)232 prev_output_channel = output_channel233 234 # out235 if use_inflated_groupnorm:236 self.conv_norm_out = InflatedGroupNorm(237 num_channels=block_out_channels[0],238 num_groups=norm_num_groups,239 eps=norm_eps,240 )241 else:242 self.conv_norm_out = nn.GroupNorm(243 num_channels=block_out_channels[0],244 num_groups=norm_num_groups,245 eps=norm_eps,246 )247 self.conv_act = nn.SiLU()248 self.conv_out = InflatedConv3d(249 block_out_channels[0], out_channels, kernel_size=3, padding=1250 )251 252 @property253 # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.attn_processors254 def attn_processors(self) -> Dict[str, AttentionProcessor]:255 r"""256 Returns:257 `dict` of attention processors: A dictionary containing all attention processors used in the model with258 indexed by its weight name.259 """260 # set recursively261 processors = {}262 263 def fn_recursive_add_processors(264 name: str,265 module: torch.nn.Module,266 processors: Dict[str, AttentionProcessor],267 ):268 if hasattr(module, "set_processor"):269 processors[f"{name}.processor"] = module.processor270 271 for sub_name, child in module.named_children():272 if "temporal_transformer" not in sub_name:273 fn_recursive_add_processors(f"{name}.{sub_name}", child, processors)274 275 return processors276 277 for name, module in self.named_children():278 if "temporal_transformer" not in name:279 fn_recursive_add_processors(name, module, processors)280 281 return processors282 283 def set_attention_slice(self, slice_size):284 r"""285 Enable sliced attention computation.286 287 When this option is enabled, the attention module will split the input tensor in slices, to compute attention288 in several steps. This is useful to save some memory in exchange for a small speed decrease.289 290 Args:291 slice_size (`str` or `int` or `list(int)`, *optional*, defaults to `"auto"`):292 When `"auto"`, halves the input to the attention heads, so attention will be computed in two steps. If293 `"max"`, maxium amount of memory will be saved by running only one slice at a time. If a number is294 provided, uses as many slices as `attention_head_dim // slice_size`. In this case, `attention_head_dim`295 must be a multiple of `slice_size`.296 """297 sliceable_head_dims = []298 299 def fn_recursive_retrieve_slicable_dims(module: torch.nn.Module):300 if hasattr(module, "set_attention_slice"):301 sliceable_head_dims.append(module.sliceable_head_dim)302 303 for child in module.children():304 fn_recursive_retrieve_slicable_dims(child)305 306 # retrieve number of attention layers307 for module in self.children():308 fn_recursive_retrieve_slicable_dims(module)309 310 num_slicable_layers = len(sliceable_head_dims)311 312 if slice_size == "auto":313 # half the attention head size is usually a good trade-off between314 # speed and memory315 slice_size = [dim // 2 for dim in sliceable_head_dims]316 elif slice_size == "max":317 # make smallest slice possible318 slice_size = num_slicable_layers * [1]319 320 slice_size = (321 num_slicable_layers * [slice_size]322 if not isinstance(slice_size, list)323 else slice_size324 )325 326 if len(slice_size) != len(sliceable_head_dims):327 raise ValueError(328 f"You have provided {len(slice_size)}, but {self.config} has {len(sliceable_head_dims)} different"329 f" attention layers. Make sure to match `len(slice_size)` to be {len(sliceable_head_dims)}."330 )331 332 for i in range(len(slice_size)):333 size = slice_size[i]334 dim = sliceable_head_dims[i]335 if size is not None and size > dim:336 raise ValueError(f"size {size} has to be smaller or equal to {dim}.")337 338 # Recursively walk through all the children.339 # Any children which exposes the set_attention_slice method340 # gets the message341 def fn_recursive_set_attention_slice(342 module: torch.nn.Module, slice_size: List[int]343 ):344 if hasattr(module, "set_attention_slice"):345 module.set_attention_slice(slice_size.pop())346 347 for child in module.children():348 fn_recursive_set_attention_slice(child, slice_size)349 350 reversed_slice_size = list(reversed(slice_size))351 for module in self.children():352 fn_recursive_set_attention_slice(module, reversed_slice_size)353 354 def _set_gradient_checkpointing(self, module, value=False):355 if hasattr(module, "gradient_checkpointing"):356 module.gradient_checkpointing = value357 358 # Copied from diffusers.models.unet_2d_condition.UNet2DConditionModel.set_attn_processor359 def set_attn_processor(360 self, processor: Union[AttentionProcessor, Dict[str, AttentionProcessor]]361 ):362 r"""363 Sets the attention processor to use to compute attention.364 365 Parameters:366 processor (`dict` of `AttentionProcessor` or only `AttentionProcessor`):367 The instantiated processor class or a dictionary of processor classes that will be set as the processor368 for **all** `Attention` layers.369 370 If `processor` is a dict, the key needs to define the path to the corresponding cross attention371 processor. This is strongly recommended when setting trainable attention processors.372 373 """374 count = len(self.attn_processors.keys())375 376 if isinstance(processor, dict) and len(processor) != count:377 raise ValueError(378 f"A dict of processors was passed, but the number of processors {len(processor)} does not match the"379 f" number of attention layers: {count}. Please make sure to pass {count} processor classes."380 )381 382 def fn_recursive_attn_processor(name: str, module: torch.nn.Module, processor):383 if hasattr(module, "set_processor"):384 if not isinstance(processor, dict):385 module.set_processor(processor)386 else:387 module.set_processor(processor.pop(f"{name}.processor"))388 389 for sub_name, child in module.named_children():390 if "temporal_transformer" not in sub_name:391 fn_recursive_attn_processor(f"{name}.{sub_name}", child, processor)392 393 for name, module in self.named_children():394 if "temporal_transformer" not in name:395 fn_recursive_attn_processor(name, module, processor)396 397 def forward(398 self,399 sample: torch.FloatTensor,400 timestep: Union[torch.Tensor, float, int],401 encoder_hidden_states: torch.Tensor,402 class_labels: Optional[torch.Tensor] = None,403 pose_cond_fea: Optional[torch.Tensor] = None,404 attention_mask: Optional[torch.Tensor] = None,405 down_block_additional_residuals: Optional[Tuple[torch.Tensor]] = None,406 mid_block_additional_residual: Optional[torch.Tensor] = None,407 return_dict: bool = True,408 ) -> Union[UNet3DConditionOutput, Tuple]:409 r"""410 Args:411 sample (`torch.FloatTensor`): (batch, channel, height, width) noisy inputs tensor412 timestep (`torch.FloatTensor` or `float` or `int`): (batch) timesteps413 encoder_hidden_states (`torch.FloatTensor`): (batch, sequence_length, feature_dim) encoder hidden states414 return_dict (`bool`, *optional*, defaults to `True`):415 Whether or not to return a [`models.unet_2d_condition.UNet2DConditionOutput`] instead of a plain tuple.416 417 Returns:418 [`~models.unet_2d_condition.UNet2DConditionOutput`] or `tuple`:419 [`~models.unet_2d_condition.UNet2DConditionOutput`] if `return_dict` is True, otherwise a `tuple`. When420 returning a tuple, the first element is the sample tensor.421 """422 # By default samples have to be AT least a multiple of the overall upsampling factor.423 # The overall upsampling factor is equal to 2 ** (# num of upsampling layears).424 # However, the upsampling interpolation output size can be forced to fit any upsampling size425 # on the fly if necessary.426 default_overall_up_factor = 2**self.num_upsamplers427 428 # upsample size should be forwarded when sample is not a multiple of `default_overall_up_factor`429 forward_upsample_size = False430 upsample_size = None431 432 if any(s % default_overall_up_factor != 0 for s in sample.shape[-2:]):433 logger.info("Forward upsample size to force interpolation output size.")434 forward_upsample_size = True435 436 # prepare attention_mask437 if attention_mask is not None:438 attention_mask = (1 - attention_mask.to(sample.dtype)) * -10000.0439 attention_mask = attention_mask.unsqueeze(1)440 441 # center input if necessary442 if self.config.center_input_sample:443 sample = 2 * sample - 1.0444 445 # time446 timesteps = timestep447 if not torch.is_tensor(timesteps):448 # This would be a good case for the `match` statement (Python 3.10+)449 is_mps = sample.device.type == "mps"450 if isinstance(timestep, float):451 dtype = torch.float32 if is_mps else torch.float64452 else:453 dtype = torch.int32 if is_mps else torch.int64454 timesteps = torch.tensor([timesteps], dtype=dtype, device=sample.device)455 elif len(timesteps.shape) == 0:456 timesteps = timesteps[None].to(sample.device)457 458 # broadcast to batch dimension in a way that's compatible with ONNX/Core ML459 timesteps = timesteps.expand(sample.shape[0])460 461 t_emb = self.time_proj(timesteps)462 463 # timesteps does not contain any weights and will always return f32 tensors464 # but time_embedding might actually be running in fp16. so we need to cast here.465 # there might be better ways to encapsulate this.466 t_emb = t_emb.to(dtype=self.dtype)467 emb = self.time_embedding(t_emb)468 469 if self.class_embedding is not None:470 if class_labels is None:471 raise ValueError(472 "class_labels should be provided when num_class_embeds > 0"473 )474 475 if self.config.class_embed_type == "timestep":476 class_labels = self.time_proj(class_labels)477 478 class_emb = self.class_embedding(class_labels).to(dtype=self.dtype)479 emb = emb + class_emb480 481 # pre-process482 sample = self.conv_in(sample)483 if pose_cond_fea is not None:484 sample = sample + pose_cond_fea485 486 # down487 down_block_res_samples = (sample,)488 for downsample_block in self.down_blocks:489 if (490 hasattr(downsample_block, "has_cross_attention")491 and downsample_block.has_cross_attention492 ):493 sample, res_samples = downsample_block(494 hidden_states=sample,495 temb=emb,496 encoder_hidden_states=encoder_hidden_states,497 attention_mask=attention_mask,498 )499 else:500 sample, res_samples = downsample_block(501 hidden_states=sample,502 temb=emb,503 encoder_hidden_states=encoder_hidden_states,504 )505 506 down_block_res_samples += res_samples507 508 if down_block_additional_residuals is not None:509 new_down_block_res_samples = ()510 511 for down_block_res_sample, down_block_additional_residual in zip(512 down_block_res_samples, down_block_additional_residuals513 ):514 down_block_res_sample = (515 down_block_res_sample + down_block_additional_residual516 )517 new_down_block_res_samples += (down_block_res_sample,)518 519 down_block_res_samples = new_down_block_res_samples520 521 # mid522 sample = self.mid_block(523 sample,524 emb,525 encoder_hidden_states=encoder_hidden_states,526 attention_mask=attention_mask,527 )528 529 if mid_block_additional_residual is not None:530 sample = sample + mid_block_additional_residual531 532 # up533 for i, upsample_block in enumerate(self.up_blocks):534 is_final_block = i == len(self.up_blocks) - 1535 536 res_samples = down_block_res_samples[-len(upsample_block.resnets) :]537 down_block_res_samples = down_block_res_samples[538 : -len(upsample_block.resnets)539 ]540 541 # if we have not reached the final block and need to forward the542 # upsample size, we do it here543 if not is_final_block and forward_upsample_size:544 upsample_size = down_block_res_samples[-1].shape[2:]545 546 if (547 hasattr(upsample_block, "has_cross_attention")548 and upsample_block.has_cross_attention549 ):550 sample = upsample_block(551 hidden_states=sample,552 temb=emb,553 res_hidden_states_tuple=res_samples,554 encoder_hidden_states=encoder_hidden_states,555 upsample_size=upsample_size,556 attention_mask=attention_mask,557 )558 else:559 sample = upsample_block(560 hidden_states=sample,561 temb=emb,562 res_hidden_states_tuple=res_samples,563 upsample_size=upsample_size,564 encoder_hidden_states=encoder_hidden_states,565 )566 567 # post-process568 sample = self.conv_norm_out(sample)569 sample = self.conv_act(sample)570 sample = self.conv_out(sample)571 572 if not return_dict:573 return (sample,)574 575 return UNet3DConditionOutput(sample=sample)576 577 @classmethod578 def from_pretrained_2d(579 cls,580 pretrained_model_path: PathLike,581 motion_module_path: PathLike,582 subfolder=None,583 unet_additional_kwargs=None,584 mm_zero_proj_out=False,585 ):586 pretrained_model_path = Path(pretrained_model_path)587 motion_module_path = Path(motion_module_path)588 if subfolder is not None:589 pretrained_model_path = pretrained_model_path.joinpath(subfolder)590 logger.info(591 f"loaded temporal unet's pretrained weights from {pretrained_model_path} ..."592 )593 594 config_file = pretrained_model_path / "config.json"595 if not (config_file.exists() and config_file.is_file()):596 raise RuntimeError(f"{config_file} does not exist or is not a file")597 598 unet_config = cls.load_config(config_file)599 unet_config["_class_name"] = cls.__name__600 unet_config["down_block_types"] = [601 "CrossAttnDownBlock3D",602 "CrossAttnDownBlock3D",603 "CrossAttnDownBlock3D",604 "DownBlock3D",605 ]606 unet_config["up_block_types"] = [607 "UpBlock3D",608 "CrossAttnUpBlock3D",609 "CrossAttnUpBlock3D",610 "CrossAttnUpBlock3D",611 ]612 unet_config["mid_block_type"] = "UNetMidBlock3DCrossAttn"613 614 model = cls.from_config(unet_config, **unet_additional_kwargs)615 # load the vanilla weights616 if pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME).exists():617 logger.debug(618 f"loading safeTensors weights from {pretrained_model_path} ..."619 )620 state_dict = load_file(621 pretrained_model_path.joinpath(SAFETENSORS_WEIGHTS_NAME), device="cpu"622 )623 624 elif pretrained_model_path.joinpath(WEIGHTS_NAME).exists():625 logger.debug(f"loading weights from {pretrained_model_path} ...")626 state_dict = torch.load(627 pretrained_model_path.joinpath(WEIGHTS_NAME),628 map_location="cpu",629 weights_only=True,630 )631 else:632 raise FileNotFoundError(f"no weights file found in {pretrained_model_path}")633 634 # load the motion module weights635 if motion_module_path.exists() and motion_module_path.is_file():636 if motion_module_path.suffix.lower() in [".pth", ".pt", ".ckpt"]:637 logger.info(f"Load motion module params from {motion_module_path}")638 motion_state_dict = torch.load(639 motion_module_path, map_location="cpu", weights_only=True640 )641 elif motion_module_path.suffix.lower() == ".safetensors":642 motion_state_dict = load_file(motion_module_path, device="cpu")643 else:644 raise RuntimeError(645 f"unknown file format for motion module weights: {motion_module_path.suffix}"646 )647 if mm_zero_proj_out:648 logger.info(f"Zero initialize proj_out layers in motion module...")649 new_motion_state_dict = OrderedDict()650 for k in motion_state_dict:651 if "proj_out" in k:652 continue653 new_motion_state_dict[k] = motion_state_dict[k]654 motion_state_dict = new_motion_state_dict655 656 # merge the state dicts657 state_dict.update(motion_state_dict)658 659 # load the weights into the model660 m, u = model.load_state_dict(state_dict, strict=False)661 logger.debug(f"### missing keys: {len(m)}; \n### unexpected keys: {len(u)};")662 663 params = [664 p.numel() if "temporal" in n else 0 for n, p in model.named_parameters()665 ]666 logger.info(f"Loaded {sum(params) / 1e6}M-parameter motion module")667 668 return model669