# Copyright 2023 The HuggingFace Team. All rights reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # Adapted from https://github.com/huggingface/diffusers/blob/v0.16.1/src/diffusers/models/unet_3d_blocks.py from typing import Any, Dict, List, Literal, Optional, Tuple, Union import logging import torch from torch import nn from diffusers.utils import is_torch_version from diffusers.models.transformer_2d import ( Transformer2DModel as DiffusersTransformer2DModel, ) from diffusers.models.resnet import Downsample2D, ResnetBlock2D, Upsample2D from ..data.data_util import batch_adain_conditioned_tensor from .resnet import TemporalConvLayer from .temporal_transformer import TransformerTemporalModel from .transformer_2d import Transformer2DModel from .attention_processor import ReferEmbFuseAttention logger = logging.getLogger(__name__) # 注： # (1) 原代码的`use_linear_projection`默认值均为True，与2D-SD模型不符，load时报错。因此均改为False # (2) 原代码调用`Transformer2DModel`的输入参数顺序为n_channels // attn_num_head_channels, attn_num_head_channels, # 与2D-SD模型不符。因此把顺序交换 # (3) 增加了temporal attention用的frame embedding输入 # note: # 1. The default value of `use_linear_projection` in the original code is True, which is inconsistent with the 2D-SD model and causes an error when loading. Therefore, it is changed to False. # 2. The original code calls `Transformer2DModel` with the input parameter order of n_channels // attn_num_head_channels, attn_num_head_channels, which is inconsistent with the 2D-SD model. Therefore, the order is reversed. # 3. Added the frame embedding input used by the temporal attention def get_down_block( down_block_type, num_layers, in_channels, out_channels, temb_channels, femb_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", temporal_transformer: Union[nn.Module, None] = TransformerTemporalModel, temporal_conv_block: Union[nn.Module, None] = TemporalConvLayer, need_spatial_position_emb: bool = False, need_t2i_ip_adapter: bool = False, ip_adapter_cross_attn: bool = False, need_t2i_facein: bool = False, need_t2i_ip_adapter_face: bool = False, need_adain_temporal_cond: bool = False, resnet_2d_skip_time_act: bool = False, need_refer_emb: bool = False, ): if (isinstance(down_block_type, str) and down_block_type == "DownBlock3D") or ( isinstance(down_block_type, nn.Module) and down_block_type.__name__ == "DownBlock3D" ): return DownBlock3D( num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, femb_channels=femb_channels, add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, downsample_padding=downsample_padding, resnet_time_scale_shift=resnet_time_scale_shift, temporal_conv_block=temporal_conv_block, need_adain_temporal_cond=need_adain_temporal_cond, resnet_2d_skip_time_act=resnet_2d_skip_time_act, need_refer_emb=need_refer_emb, attn_num_head_channels=attn_num_head_channels, ) elif ( isinstance(down_block_type, str) and down_block_type == "CrossAttnDownBlock3D" ) or ( isinstance(down_block_type, nn.Module) and down_block_type.__name__ == "CrossAttnDownBlock3D" ): if cross_attention_dim is None: raise ValueError( "cross_attention_dim must be specified for CrossAttnDownBlock3D" ) return CrossAttnDownBlock3D( num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, femb_channels=femb_channels, add_downsample=add_downsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, downsample_padding=downsample_padding, cross_attention_dim=cross_attention_dim, attn_num_head_channels=attn_num_head_channels, dual_cross_attention=dual_cross_attention, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, resnet_time_scale_shift=resnet_time_scale_shift, temporal_conv_block=temporal_conv_block, temporal_transformer=temporal_transformer, need_spatial_position_emb=need_spatial_position_emb, need_t2i_ip_adapter=need_t2i_ip_adapter, ip_adapter_cross_attn=ip_adapter_cross_attn, need_t2i_facein=need_t2i_facein, need_t2i_ip_adapter_face=need_t2i_ip_adapter_face, need_adain_temporal_cond=need_adain_temporal_cond, resnet_2d_skip_time_act=resnet_2d_skip_time_act, need_refer_emb=need_refer_emb, ) raise ValueError(f"{down_block_type} does not exist.") def get_up_block( up_block_type, num_layers, in_channels, out_channels, prev_output_channel, temb_channels, femb_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", temporal_conv_block: Union[nn.Module, None] = TemporalConvLayer, temporal_transformer: Union[nn.Module, None] = TransformerTemporalModel, need_spatial_position_emb: bool = False, need_t2i_ip_adapter: bool = False, ip_adapter_cross_attn: bool = False, need_t2i_facein: bool = False, need_t2i_ip_adapter_face: bool = False, need_adain_temporal_cond: bool = False, resnet_2d_skip_time_act: bool = False, ): if (isinstance(up_block_type, str) and up_block_type == "UpBlock3D") or ( isinstance(up_block_type, nn.Module) and up_block_type.__name__ == "UpBlock3D" ): return UpBlock3D( num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel, temb_channels=temb_channels, femb_channels=femb_channels, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, resnet_time_scale_shift=resnet_time_scale_shift, temporal_conv_block=temporal_conv_block, need_adain_temporal_cond=need_adain_temporal_cond, resnet_2d_skip_time_act=resnet_2d_skip_time_act, ) elif (isinstance(up_block_type, str) and up_block_type == "CrossAttnUpBlock3D") or ( isinstance(up_block_type, nn.Module) and up_block_type.__name__ == "CrossAttnUpBlock3D" ): if cross_attention_dim is None: raise ValueError( "cross_attention_dim must be specified for CrossAttnUpBlock3D" ) return CrossAttnUpBlock3D( num_layers=num_layers, in_channels=in_channels, out_channels=out_channels, prev_output_channel=prev_output_channel, temb_channels=temb_channels, femb_channels=femb_channels, add_upsample=add_upsample, resnet_eps=resnet_eps, resnet_act_fn=resnet_act_fn, resnet_groups=resnet_groups, cross_attention_dim=cross_attention_dim, attn_num_head_channels=attn_num_head_channels, dual_cross_attention=dual_cross_attention, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, resnet_time_scale_shift=resnet_time_scale_shift, temporal_conv_block=temporal_conv_block, temporal_transformer=temporal_transformer, need_spatial_position_emb=need_spatial_position_emb, need_t2i_ip_adapter=need_t2i_ip_adapter, ip_adapter_cross_attn=ip_adapter_cross_attn, need_t2i_facein=need_t2i_facein, need_t2i_ip_adapter_face=need_t2i_ip_adapter_face, need_adain_temporal_cond=need_adain_temporal_cond, resnet_2d_skip_time_act=resnet_2d_skip_time_act, ) raise ValueError(f"{up_block_type} does not exist.") class UNetMidBlock3DCrossAttn(nn.Module): print_idx = 0 def __init__( self, in_channels: int, temb_channels: int, femb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, 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, temporal_conv_block: Union[nn.Module, None] = TemporalConvLayer, temporal_transformer: Union[nn.Module, None] = TransformerTemporalModel, need_spatial_position_emb: bool = False, need_t2i_ip_adapter: bool = False, ip_adapter_cross_attn: bool = False, need_t2i_facein: bool = False, need_t2i_ip_adapter_face: bool = False, need_adain_temporal_cond: bool = False, resnet_2d_skip_time_act: bool = False, ): super().__init__() self.has_cross_attention = True self.attn_num_head_channels = attn_num_head_channels resnet_groups = ( resnet_groups if resnet_groups is not None else min(in_channels // 4, 32) ) # there is always at least one resnet resnets = [ ResnetBlock2D( in_channels=in_channels, out_channels=in_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=resnet_2d_skip_time_act, ) ] if temporal_conv_block is not None: temp_convs = [ temporal_conv_block( in_channels, in_channels, dropout=0.1, femb_channels=femb_channels, ) ] else: temp_convs = [None] attentions = [] temp_attentions = [] for _ in range(num_layers): attentions.append( Transformer2DModel( attn_num_head_channels, in_channels // attn_num_head_channels, in_channels=in_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, use_linear_projection=use_linear_projection, upcast_attention=upcast_attention, cross_attn_temporal_cond=need_t2i_ip_adapter, ip_adapter_cross_attn=ip_adapter_cross_attn, need_t2i_facein=need_t2i_facein, need_t2i_ip_adapter_face=need_t2i_ip_adapter_face, ) ) if temporal_transformer is not None: temp_attention = temporal_transformer( attn_num_head_channels, in_channels // attn_num_head_channels, in_channels=in_channels, num_layers=1, femb_channels=femb_channels, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, need_spatial_position_emb=need_spatial_position_emb, ) else: temp_attention = None temp_attentions.append(temp_attention) resnets.append( ResnetBlock2D( in_channels=in_channels, out_channels=in_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=resnet_2d_skip_time_act, ) ) if temporal_conv_block is not None: temp_convs.append( temporal_conv_block( in_channels, in_channels, dropout=0.1, femb_channels=femb_channels, ) ) else: temp_convs.append(None) self.resnets = nn.ModuleList(resnets) self.temp_convs = nn.ModuleList(temp_convs) self.attentions = nn.ModuleList(attentions) self.temp_attentions = nn.ModuleList(temp_attentions) self.need_adain_temporal_cond = need_adain_temporal_cond def forward( self, hidden_states, temb=None, femb=None, encoder_hidden_states=None, attention_mask=None, num_frames=1, cross_attention_kwargs=None, sample_index: torch.LongTensor = None, vision_conditon_frames_sample_index: torch.LongTensor = None, spatial_position_emb: torch.Tensor = None, refer_self_attn_emb: List[torch.Tensor] = None, refer_self_attn_emb_mode: Literal["read", "write"] = "read", ): hidden_states = self.resnets[0](hidden_states, temb) if self.temp_convs[0] is not None: hidden_states = self.temp_convs[0]( hidden_states, femb=femb, num_frames=num_frames, sample_index=sample_index, vision_conditon_frames_sample_index=vision_conditon_frames_sample_index, ) for attn, temp_attn, resnet, temp_conv in zip( self.attentions, self.temp_attentions, self.resnets[1:], self.temp_convs[1:] ): hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, self_attn_block_embs=refer_self_attn_emb, self_attn_block_embs_mode=refer_self_attn_emb_mode, ).sample if temp_attn is not None: hidden_states = temp_attn( hidden_states, femb=femb, num_frames=num_frames, cross_attention_kwargs=cross_attention_kwargs, encoder_hidden_states=encoder_hidden_states, sample_index=sample_index, vision_conditon_frames_sample_index=vision_conditon_frames_sample_index, spatial_position_emb=spatial_position_emb, ).sample hidden_states = resnet(hidden_states, temb) if temp_conv is not None: hidden_states = temp_conv( hidden_states, femb=femb, num_frames=num_frames, sample_index=sample_index, vision_conditon_frames_sample_index=vision_conditon_frames_sample_index, ) if ( self.need_adain_temporal_cond and num_frames > 1 and sample_index is not None ): if self.print_idx == 0: logger.debug(f"adain to vision_condition") hidden_states = batch_adain_conditioned_tensor( hidden_states, num_frames=num_frames, need_style_fidelity=False, src_index=sample_index, dst_index=vision_conditon_frames_sample_index, ) self.print_idx += 1 return hidden_states class CrossAttnDownBlock3D(nn.Module): print_idx = 0 def __init__( self, in_channels: int, out_channels: int, temb_channels: int, femb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, 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, temporal_conv_block: Union[nn.Module, None] = TemporalConvLayer, temporal_transformer: Union[nn.Module, None] = TransformerTemporalModel, need_spatial_position_emb: bool = False, need_t2i_ip_adapter: bool = False, ip_adapter_cross_attn: bool = False, need_t2i_facein: bool = False, need_t2i_ip_adapter_face: bool = False, need_adain_temporal_cond: bool = False, resnet_2d_skip_time_act: bool = False, need_refer_emb: bool = False, ): super().__init__() resnets = [] attentions = [] temp_attentions = [] temp_convs = [] self.has_cross_attention = True self.attn_num_head_channels = attn_num_head_channels self.need_refer_emb = need_refer_emb if need_refer_emb: refer_emb_attns = [] for i in range(num_layers): in_channels = in_channels if i == 0 else out_channels resnets.append( ResnetBlock2D( in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=resnet_2d_skip_time_act, ) ) if temporal_conv_block is not None: temp_convs.append( temporal_conv_block( out_channels, out_channels, dropout=0.1, femb_channels=femb_channels, ) ) else: temp_convs.append(None) attentions.append( Transformer2DModel( attn_num_head_channels, out_channels // attn_num_head_channels, in_channels=out_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, cross_attn_temporal_cond=need_t2i_ip_adapter, ip_adapter_cross_attn=ip_adapter_cross_attn, need_t2i_facein=need_t2i_facein, need_t2i_ip_adapter_face=need_t2i_ip_adapter_face, ) ) if temporal_transformer is not None: temp_attention = temporal_transformer( attn_num_head_channels, out_channels // attn_num_head_channels, in_channels=out_channels, num_layers=1, femb_channels=femb_channels, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, need_spatial_position_emb=need_spatial_position_emb, ) else: temp_attention = None temp_attentions.append(temp_attention) if need_refer_emb: refer_emb_attns.append( ReferEmbFuseAttention( query_dim=out_channels, heads=attn_num_head_channels, dim_head=out_channels // attn_num_head_channels, dropout=0, bias=False, cross_attention_dim=None, upcast_attention=False, ) ) self.resnets = nn.ModuleList(resnets) self.temp_convs = nn.ModuleList(temp_convs) self.attentions = nn.ModuleList(attentions) self.temp_attentions = nn.ModuleList(temp_attentions) if add_downsample: self.downsamplers = nn.ModuleList( [ Downsample2D( out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op", ) ] ) if need_refer_emb: refer_emb_attns.append( ReferEmbFuseAttention( query_dim=out_channels, heads=attn_num_head_channels, dim_head=out_channels // attn_num_head_channels, dropout=0, bias=False, cross_attention_dim=None, upcast_attention=False, ) ) else: self.downsamplers = None self.gradient_checkpointing = False self.need_adain_temporal_cond = need_adain_temporal_cond if need_refer_emb: self.refer_emb_attns = nn.ModuleList(refer_emb_attns) logger.debug(f"cross attn downblock 3d need_refer_emb, {self.need_refer_emb}") def forward( self, hidden_states: torch.FloatTensor, temb: Optional[torch.FloatTensor] = None, femb: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, attention_mask: Optional[torch.FloatTensor] = None, num_frames: int = 1, cross_attention_kwargs: Optional[Dict[str, Any]] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, sample_index: torch.LongTensor = None, vision_conditon_frames_sample_index: torch.LongTensor = None, spatial_position_emb: torch.Tensor = None, refer_embs: Optional[List[torch.Tensor]] = None, refer_self_attn_emb: List[torch.Tensor] = None, refer_self_attn_emb_mode: Literal["read", "write"] = "read", ): # TODO(Patrick, William) - attention mask is not used output_states = () for i_downblock, (resnet, temp_conv, attn, temp_attn) in enumerate( zip(self.resnets, self.temp_convs, self.attentions, self.temp_attentions) ): # print("crossattndownblock3d, attn,", type(attn), cross_attention_kwargs) if self.training and self.gradient_checkpointing: if self.print_idx == 0: logger.debug( f"self.training and self.gradient_checkpointing={self.training and self.gradient_checkpointing}" ) def create_custom_forward(module, return_dict=None): def custom_forward(*inputs): if return_dict is not None: return module(*inputs, return_dict=return_dict) else: return module(*inputs) return custom_forward ckpt_kwargs: Dict[str, Any] = ( {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} ) hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs, ) if self.print_idx == 0: logger.debug(f"unet3d after resnet {hidden_states.mean()}") if temp_conv is not None: hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(temp_conv), hidden_states, num_frames, sample_index, vision_conditon_frames_sample_index, femb, **ckpt_kwargs, ) hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(attn, return_dict=False), hidden_states, encoder_hidden_states, None, # timestep None, # added_cond_kwargs None, # class_labels cross_attention_kwargs, attention_mask, encoder_attention_mask, refer_self_attn_emb, refer_self_attn_emb_mode, **ckpt_kwargs, )[0] if temp_attn is not None: hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(temp_attn, return_dict=False), hidden_states, femb, # None, # encoder_hidden_states, encoder_hidden_states, None, # timestep None, # class_labels num_frames, cross_attention_kwargs, sample_index, vision_conditon_frames_sample_index, spatial_position_emb, **ckpt_kwargs, )[0] else: hidden_states = resnet(hidden_states, temb) if self.print_idx == 0: logger.debug(f"unet3d after resnet {hidden_states.mean()}") if temp_conv is not None: hidden_states = temp_conv( hidden_states, femb=femb, num_frames=num_frames, sample_index=sample_index, vision_conditon_frames_sample_index=vision_conditon_frames_sample_index, ) hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, self_attn_block_embs=refer_self_attn_emb, self_attn_block_embs_mode=refer_self_attn_emb_mode, ).sample if temp_attn is not None: hidden_states = temp_attn( hidden_states, femb=femb, num_frames=num_frames, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, sample_index=sample_index, vision_conditon_frames_sample_index=vision_conditon_frames_sample_index, spatial_position_emb=spatial_position_emb, ).sample if ( self.need_adain_temporal_cond and num_frames > 1 and sample_index is not None ): if self.print_idx == 0: logger.debug(f"adain to vision_condition") hidden_states = batch_adain_conditioned_tensor( hidden_states, num_frames=num_frames, need_style_fidelity=False, src_index=sample_index, dst_index=vision_conditon_frames_sample_index, ) # 使用 attn 的方式 来融合 down_block_refer_emb if self.print_idx == 0: logger.debug( f"downblock, {i_downblock}, self.need_refer_emb={self.need_refer_emb}" ) if self.need_refer_emb and refer_embs is not None: if self.print_idx == 0: logger.debug( f"{i_downblock}, self.refer_emb_attns {refer_embs[i_downblock].shape}" ) hidden_states = self.refer_emb_attns[i_downblock]( hidden_states, refer_embs[i_downblock], num_frames=num_frames ) else: if self.print_idx == 0: logger.debug(f"crossattndownblock refer_emb_attns, no this step") output_states += (hidden_states,) if self.downsamplers is not None: for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states) if ( self.need_adain_temporal_cond and num_frames > 1 and sample_index is not None ): if self.print_idx == 0: logger.debug(f"adain to vision_condition") hidden_states = batch_adain_conditioned_tensor( hidden_states, num_frames=num_frames, need_style_fidelity=False, src_index=sample_index, dst_index=vision_conditon_frames_sample_index, ) # 使用 attn 的方式 来融合 down_block_refer_emb # TODO: adain和 refer_emb的顺序 # TODO：adain 首帧特征还是refer_emb的 if self.need_refer_emb and refer_embs is not None: i_downblock += 1 hidden_states = self.refer_emb_attns[i_downblock]( hidden_states, refer_embs[i_downblock], num_frames=num_frames ) output_states += (hidden_states,) self.print_idx += 1 return hidden_states, output_states class DownBlock3D(nn.Module): print_idx = 0 def __init__( self, in_channels: int, out_channels: int, temb_channels: int, femb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, 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, temporal_conv_block: Union[nn.Module, None] = TemporalConvLayer, need_adain_temporal_cond: bool = False, resnet_2d_skip_time_act: bool = False, need_refer_emb: bool = False, attn_num_head_channels: int = 1, ): super().__init__() resnets = [] temp_convs = [] self.need_refer_emb = need_refer_emb if need_refer_emb: refer_emb_attns = [] self.attn_num_head_channels = attn_num_head_channels for i in range(num_layers): in_channels = in_channels if i == 0 else out_channels resnets.append( ResnetBlock2D( in_channels=in_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=resnet_2d_skip_time_act, ) ) if temporal_conv_block is not None: temp_convs.append( temporal_conv_block( out_channels, out_channels, dropout=0.1, femb_channels=femb_channels, ) ) else: temp_convs.append(None) if need_refer_emb: refer_emb_attns.append( ReferEmbFuseAttention( query_dim=out_channels, heads=attn_num_head_channels, dim_head=out_channels // attn_num_head_channels, dropout=0, bias=False, cross_attention_dim=None, upcast_attention=False, ) ) self.resnets = nn.ModuleList(resnets) self.temp_convs = nn.ModuleList(temp_convs) if add_downsample: self.downsamplers = nn.ModuleList( [ Downsample2D( out_channels, use_conv=True, out_channels=out_channels, padding=downsample_padding, name="op", ) ] ) if need_refer_emb: refer_emb_attns.append( ReferEmbFuseAttention( query_dim=out_channels, heads=attn_num_head_channels, dim_head=out_channels // attn_num_head_channels, dropout=0, bias=False, cross_attention_dim=None, upcast_attention=False, ) ) else: self.downsamplers = None self.gradient_checkpointing = False self.need_adain_temporal_cond = need_adain_temporal_cond if need_refer_emb: self.refer_emb_attns = nn.ModuleList(refer_emb_attns) def forward( self, hidden_states, temb=None, num_frames=1, sample_index: torch.LongTensor = None, vision_conditon_frames_sample_index: torch.LongTensor = None, spatial_position_emb: torch.Tensor = None, femb=None, refer_embs: Optional[Tuple[torch.Tensor]] = None, refer_self_attn_emb: List[torch.Tensor] = None, refer_self_attn_emb_mode: Literal["read", "write"] = "read", ): output_states = () for i_downblock, (resnet, temp_conv) in enumerate( zip(self.resnets, self.temp_convs) ): if self.training and self.gradient_checkpointing: def create_custom_forward(module): def custom_forward(*inputs): return module(*inputs) return custom_forward ckpt_kwargs: Dict[str, Any] = ( {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} ) hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs, ) if temp_conv is not None: hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(temp_conv), hidden_states, num_frames, sample_index, vision_conditon_frames_sample_index, femb, **ckpt_kwargs, ) else: hidden_states = resnet(hidden_states, temb) if temp_conv is not None: hidden_states = temp_conv( hidden_states, femb=femb, num_frames=num_frames, sample_index=sample_index, vision_conditon_frames_sample_index=vision_conditon_frames_sample_index, ) if ( self.need_adain_temporal_cond and num_frames > 1 and sample_index is not None ): if self.print_idx == 0: logger.debug(f"adain to vision_condition") hidden_states = batch_adain_conditioned_tensor( hidden_states, num_frames=num_frames, need_style_fidelity=False, src_index=sample_index, dst_index=vision_conditon_frames_sample_index, ) if self.need_refer_emb and refer_embs is not None: hidden_states = self.refer_emb_attns[i_downblock]( hidden_states, refer_embs[i_downblock], num_frames=num_frames ) output_states += (hidden_states,) if self.downsamplers is not None: for downsampler in self.downsamplers: hidden_states = downsampler(hidden_states) if ( self.need_adain_temporal_cond and num_frames > 1 and sample_index is not None ): if self.print_idx == 0: logger.debug(f"adain to vision_condition") hidden_states = batch_adain_conditioned_tensor( hidden_states, num_frames=num_frames, need_style_fidelity=False, src_index=sample_index, dst_index=vision_conditon_frames_sample_index, ) if self.need_refer_emb and refer_embs is not None: i_downblock += 1 hidden_states = self.refer_emb_attns[i_downblock]( hidden_states, refer_embs[i_downblock], num_frames=num_frames ) output_states += (hidden_states,) self.print_idx += 1 return hidden_states, output_states class CrossAttnUpBlock3D(nn.Module): print_idx = 0 def __init__( self, in_channels: int, out_channels: int, prev_output_channel: int, temb_channels: int, femb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, 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, temporal_conv_block: Union[nn.Module, None] = TemporalConvLayer, temporal_transformer: Union[nn.Module, None] = TransformerTemporalModel, need_spatial_position_emb: bool = False, need_t2i_ip_adapter: bool = False, ip_adapter_cross_attn: bool = False, need_t2i_facein: bool = False, need_t2i_ip_adapter_face: bool = False, need_adain_temporal_cond: bool = False, resnet_2d_skip_time_act: bool = False, ): super().__init__() resnets = [] temp_convs = [] attentions = [] temp_attentions = [] self.has_cross_attention = True self.attn_num_head_channels = attn_num_head_channels for i in range(num_layers): res_skip_channels = in_channels if (i == num_layers - 1) else out_channels resnet_in_channels = prev_output_channel if i == 0 else out_channels resnets.append( ResnetBlock2D( in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=resnet_2d_skip_time_act, ) ) if temporal_conv_block is not None: temp_convs.append( temporal_conv_block( out_channels, out_channels, dropout=0.1, femb_channels=femb_channels, ) ) else: temp_convs.append(None) attentions.append( Transformer2DModel( attn_num_head_channels, out_channels // attn_num_head_channels, in_channels=out_channels, num_layers=1, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, use_linear_projection=use_linear_projection, only_cross_attention=only_cross_attention, upcast_attention=upcast_attention, cross_attn_temporal_cond=need_t2i_ip_adapter, ip_adapter_cross_attn=ip_adapter_cross_attn, need_t2i_facein=need_t2i_facein, need_t2i_ip_adapter_face=need_t2i_ip_adapter_face, ) ) if temporal_transformer is not None: temp_attention = temporal_transformer( attn_num_head_channels, out_channels // attn_num_head_channels, in_channels=out_channels, num_layers=1, femb_channels=femb_channels, cross_attention_dim=cross_attention_dim, norm_num_groups=resnet_groups, need_spatial_position_emb=need_spatial_position_emb, ) else: temp_attention = None temp_attentions.append(temp_attention) self.resnets = nn.ModuleList(resnets) self.temp_convs = nn.ModuleList(temp_convs) self.attentions = nn.ModuleList(attentions) self.temp_attentions = nn.ModuleList(temp_attentions) if add_upsample: self.upsamplers = nn.ModuleList( [Upsample2D(out_channels, use_conv=True, out_channels=out_channels)] ) else: self.upsamplers = None self.gradient_checkpointing = False self.need_adain_temporal_cond = need_adain_temporal_cond def forward( self, hidden_states: torch.FloatTensor, res_hidden_states_tuple: Tuple[torch.FloatTensor, ...], temb: Optional[torch.FloatTensor] = None, femb: Optional[torch.FloatTensor] = None, encoder_hidden_states: Optional[torch.FloatTensor] = None, num_frames: int = 1, cross_attention_kwargs: Optional[Dict[str, Any]] = None, upsample_size: Optional[int] = None, attention_mask: Optional[torch.FloatTensor] = None, encoder_attention_mask: Optional[torch.FloatTensor] = None, sample_index: torch.LongTensor = None, vision_conditon_frames_sample_index: torch.LongTensor = None, spatial_position_emb: torch.Tensor = None, refer_self_attn_emb: List[torch.Tensor] = None, refer_self_attn_emb_mode: Literal["read", "write"] = "read", ): for resnet, temp_conv, attn, temp_attn in zip( self.resnets, self.temp_convs, self.attentions, self.temp_attentions ): # pop res hidden states res_hidden_states = res_hidden_states_tuple[-1] res_hidden_states_tuple = res_hidden_states_tuple[:-1] hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) if self.training and self.gradient_checkpointing: def create_custom_forward(module, return_dict=None): def custom_forward(*inputs): if return_dict is not None: return module(*inputs, return_dict=return_dict) else: return module(*inputs) return custom_forward ckpt_kwargs: Dict[str, Any] = ( {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} ) hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs, ) if temp_conv is not None: hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(temp_conv), hidden_states, num_frames, sample_index, vision_conditon_frames_sample_index, femb, **ckpt_kwargs, ) hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(attn, return_dict=False), hidden_states, encoder_hidden_states, None, # timestep None, # added_cond_kwargs None, # class_labels cross_attention_kwargs, attention_mask, encoder_attention_mask, refer_self_attn_emb, refer_self_attn_emb_mode, **ckpt_kwargs, )[0] if temp_attn is not None: hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(temp_attn, return_dict=False), hidden_states, femb, # None, # encoder_hidden_states, encoder_hidden_states, None, # timestep None, # class_labels num_frames, cross_attention_kwargs, sample_index, vision_conditon_frames_sample_index, spatial_position_emb, **ckpt_kwargs, )[0] else: hidden_states = resnet(hidden_states, temb) if temp_conv is not None: hidden_states = temp_conv( hidden_states, num_frames=num_frames, femb=femb, sample_index=sample_index, vision_conditon_frames_sample_index=vision_conditon_frames_sample_index, ) hidden_states = attn( hidden_states, encoder_hidden_states=encoder_hidden_states, cross_attention_kwargs=cross_attention_kwargs, self_attn_block_embs=refer_self_attn_emb, self_attn_block_embs_mode=refer_self_attn_emb_mode, ).sample if temp_attn is not None: hidden_states = temp_attn( hidden_states, femb=femb, num_frames=num_frames, cross_attention_kwargs=cross_attention_kwargs, encoder_hidden_states=encoder_hidden_states, sample_index=sample_index, vision_conditon_frames_sample_index=vision_conditon_frames_sample_index, spatial_position_emb=spatial_position_emb, ).sample if ( self.need_adain_temporal_cond and num_frames > 1 and sample_index is not None ): if self.print_idx == 0: logger.debug(f"adain to vision_condition") hidden_states = batch_adain_conditioned_tensor( hidden_states, num_frames=num_frames, need_style_fidelity=False, src_index=sample_index, dst_index=vision_conditon_frames_sample_index, ) if self.upsamplers is not None: for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states, upsample_size) if ( self.need_adain_temporal_cond and num_frames > 1 and sample_index is not None ): if self.print_idx == 0: logger.debug(f"adain to vision_condition") hidden_states = batch_adain_conditioned_tensor( hidden_states, num_frames=num_frames, need_style_fidelity=False, src_index=sample_index, dst_index=vision_conditon_frames_sample_index, ) self.print_idx += 1 return hidden_states class UpBlock3D(nn.Module): print_idx = 0 def __init__( self, in_channels: int, prev_output_channel: int, out_channels: int, temb_channels: int, femb_channels: int, dropout: float = 0.0, num_layers: int = 1, resnet_eps: float = 1e-6, 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, temporal_conv_block: Union[nn.Module, None] = TemporalConvLayer, need_adain_temporal_cond: bool = False, resnet_2d_skip_time_act: bool = False, ): super().__init__() resnets = [] temp_convs = [] for i in range(num_layers): res_skip_channels = in_channels if (i == num_layers - 1) else out_channels resnet_in_channels = prev_output_channel if i == 0 else out_channels resnets.append( ResnetBlock2D( in_channels=resnet_in_channels + res_skip_channels, out_channels=out_channels, temb_channels=temb_channels, eps=resnet_eps, groups=resnet_groups, dropout=dropout, time_embedding_norm=resnet_time_scale_shift, non_linearity=resnet_act_fn, output_scale_factor=output_scale_factor, pre_norm=resnet_pre_norm, skip_time_act=resnet_2d_skip_time_act, ) ) if temporal_conv_block is not None: temp_convs.append( temporal_conv_block( out_channels, out_channels, dropout=0.1, femb_channels=femb_channels, ) ) else: temp_convs.append(None) self.resnets = nn.ModuleList(resnets) self.temp_convs = nn.ModuleList(temp_convs) if add_upsample: self.upsamplers = nn.ModuleList( [Upsample2D(out_channels, use_conv=True, out_channels=out_channels)] ) else: self.upsamplers = None self.gradient_checkpointing = False self.need_adain_temporal_cond = need_adain_temporal_cond def forward( self, hidden_states, res_hidden_states_tuple, temb=None, upsample_size=None, num_frames=1, sample_index: torch.LongTensor = None, vision_conditon_frames_sample_index: torch.LongTensor = None, spatial_position_emb: torch.Tensor = None, femb=None, refer_self_attn_emb: List[torch.Tensor] = None, refer_self_attn_emb_mode: Literal["read", "write"] = "read", ): for resnet, temp_conv in zip(self.resnets, self.temp_convs): # pop res hidden states res_hidden_states = res_hidden_states_tuple[-1] res_hidden_states_tuple = res_hidden_states_tuple[:-1] hidden_states = torch.cat([hidden_states, res_hidden_states], dim=1) if self.training and self.gradient_checkpointing: def create_custom_forward(module): def custom_forward(*inputs): return module(*inputs) return custom_forward ckpt_kwargs: Dict[str, Any] = ( {"use_reentrant": False} if is_torch_version(">=", "1.11.0") else {} ) hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(resnet), hidden_states, temb, **ckpt_kwargs, ) if temp_conv is not None: hidden_states = torch.utils.checkpoint.checkpoint( create_custom_forward(temp_conv), hidden_states, num_frames, sample_index, vision_conditon_frames_sample_index, femb, **ckpt_kwargs, ) else: hidden_states = resnet(hidden_states, temb) if temp_conv is not None: hidden_states = temp_conv( hidden_states, num_frames=num_frames, femb=femb, sample_index=sample_index, vision_conditon_frames_sample_index=vision_conditon_frames_sample_index, ) if ( self.need_adain_temporal_cond and num_frames > 1 and sample_index is not None ): if self.print_idx == 0: logger.debug(f"adain to vision_condition") hidden_states = batch_adain_conditioned_tensor( hidden_states, num_frames=num_frames, need_style_fidelity=False, src_index=sample_index, dst_index=vision_conditon_frames_sample_index, ) if self.upsamplers is not None: for upsampler in self.upsamplers: hidden_states = upsampler(hidden_states, upsample_size) if ( self.need_adain_temporal_cond and num_frames > 1 and sample_index is not None ): if self.print_idx == 0: logger.debug(f"adain to vision_condition") hidden_states = batch_adain_conditioned_tensor( hidden_states, num_frames=num_frames, need_style_fidelity=False, src_index=sample_index, dst_index=vision_conditon_frames_sample_index, ) self.print_idx += 1 return hidden_states