# -*- encoding: utf-8 -*- ''' @File : cogvideo_model.py @Time : 2022/07/11 16:12:05 @Author : Wenyi Hong @Version : 1.0 @Contact : hwy22@mails.tsinghua.edu.cn ''' # here put the import lib import torch from SwissArmyTransformer.model.base_model import BaseModel, BaseMixin from SwissArmyTransformer.mpu.utils import split_tensor_along_last_dim from SwissArmyTransformer.model.transformer import unscaled_init_method from SwissArmyTransformer.mpu import ColumnParallelLinear, RowParallelLinear import torch.nn.functional as F from deepspeed.runtime.activation_checkpointing.checkpointing import get_cuda_rng_tracker import math class PositionEmbeddingMixin(BaseMixin): def __init__(self, additional_sequence_length, hidden_size, init_method_std=0.02, reinit_slice=slice(512, 912), ): super(PositionEmbeddingMixin, self).__init__() self.reinit_slice = reinit_slice self.position_embeddings = torch.nn.Embedding(additional_sequence_length, hidden_size) torch.nn.init.normal_(self.position_embeddings.weight, mean=0.0, std=init_method_std) def reinit(self, parent_model=None): old_weights = self.transformer.position_embeddings.weight.data[self.reinit_slice] old_len, hidden_size = old_weights.shape assert hidden_size == self.position_embeddings.weight.shape[-1] self.position_embeddings.weight.data.view(-1, old_len, hidden_size).copy_(old_weights) def window_partition(x, window_size): """ Args: x: (B, framenum, H, W, C) window_size (int): window size Returns: windows: (num_windows*B, frame_num, window_size, window_size, C) """ B, framenum, H, W, C = x.shape x = x.view(B, framenum, H // window_size, window_size, W // window_size, window_size, C) windows = x.permute(0, 2, 4, 1, 3, 5, 6).contiguous().view(-1, framenum, window_size, window_size, C) return windows def window_reverse(windows, window_size, H, W): """ Args: windows: (num_windows*B, frame_num, window_size, window_size, C) window_size (int): Window size H (int): Height of image W (int): Width of image Returns: x: (B, frame_num, H, W, C) """ B = int(windows.shape[0] / (H * W / window_size / window_size)) framenum = windows.shape[1] x = windows.view(B, H // window_size, W // window_size, framenum, window_size, window_size, -1) x = x.permute(0, 3, 1, 4, 2, 5, 6).contiguous().view(B, framenum, H, W, -1) return x class WindowAttentionMixin(BaseMixin): def __init__(self, num_layers, hidden_size, frame_resolution, window_size, shift_size, n_head, frame_num, init_method=unscaled_init_method(0.02), output_layer_init_method=unscaled_init_method(0.02), ): super(WindowAttentionMixin, self).__init__() self.num_layers = num_layers # replace attention in the LAST n layers self.query_key_value = torch.nn.ModuleList( [ColumnParallelLinear(hidden_size, 3*hidden_size,stride=3, gather_output=False,init_method=init_method) for layer_id in range(num_layers) ]) self.dense = torch.nn.ModuleList( [RowParallelLinear( hidden_size, hidden_size, input_is_parallel=True, init_method=output_layer_init_method, bias=True, module=self, name="dense", ) for layer_id in range(num_layers) ]) self.n_head = n_head self.window_size = window_size self.frame_resolution = frame_resolution self.frame_len = frame_resolution * frame_resolution assert frame_resolution % window_size == 0 assert 0 < shift_size < window_size nW = (self.frame_resolution // self.window_size) ** 2 ws_squre = self.window_size * self.window_size # odd non-shift, even shift img_mask = torch.zeros((1, 1, frame_resolution, frame_resolution, 1)) h_slices = (slice(0, -shift_size), slice(-shift_size, None)) w_slices = (slice(0, -shift_size), slice(-shift_size, None)) cnt = 0 for h in h_slices: for w in w_slices: img_mask[:, :, h, w, :] = cnt cnt += 1 mask_windows = window_partition(img_mask, self.window_size) # nW, 1, window_size, window_size, 1 mask_windows = mask_windows.view(-1, self.window_size * self.window_size) sub_attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2) #[nW, self.window_size * self.window_size, self.window_size * self.window_size] sub_attn_mask = sub_attn_mask.masked_fill(sub_attn_mask != 0, float(0.0)).masked_fill(sub_attn_mask == 0, float(1.00)) attn_mask = sub_attn_mask.repeat(1, frame_num, frame_num) self.attn_mask_sequential = attn_mask.clone().tril() self.causal_mask_sequential = torch.ones(1, ws_squre*frame_num, ws_squre*frame_num).tril() self.causal_mask_interp = torch.ones(1, ws_squre*frame_num, ws_squre*frame_num) self.attn_mask_interp = attn_mask.clone() # bi-dir for bi_idx in range(0, frame_num, 2): for uni_idx in range(1, frame_num, 2): self.attn_mask_interp[:, bi_idx*ws_squre:(bi_idx+1)*ws_squre, uni_idx*ws_squre:(uni_idx+1)*ws_squre] = 0 self.causal_mask_interp[:, bi_idx*ws_squre:(bi_idx+1)*ws_squre, uni_idx*ws_squre:(uni_idx+1)*ws_squre] = 0 # uni-dir for uni_idx in range(1, frame_num, 2): self.attn_mask_interp[:, ws_squre*uni_idx:ws_squre*(uni_idx+1), ws_squre*uni_idx:ws_squre*(uni_idx+1)].tril_() self.causal_mask_interp[:, ws_squre*uni_idx:ws_squre*(uni_idx+1), ws_squre*uni_idx:ws_squre*(uni_idx+1)].tril_() for uni_idx2 in range(uni_idx+2, frame_num, 2): self.attn_mask_interp[:, ws_squre*uni_idx:ws_squre*(uni_idx+1), ws_squre*uni_idx2:ws_squre*(uni_idx2+1)] = 0 self.causal_mask_interp[:, ws_squre*uni_idx:ws_squre*(uni_idx+1), ws_squre*uni_idx2:ws_squre*(uni_idx2+1)] = 0 # expand dim self.attn_mask_sequential = self.attn_mask_sequential[None, None, :, None] self.attn_mask_interp = self.attn_mask_interp[None, None, :, None] self.causal_mask_sequential = self.causal_mask_sequential[None, None, :, None] self.causal_mask_interp = self.causal_mask_interp[None, None, :, None] self.shift_sizes = [0, shift_size] # self.register_buffer("attn_mask", attn_mask) # self.register_buffer("causal_mask", causal_mask) self.mask_initialized = False self.attn_distribution = torch.nn.ParameterList([ torch.nn.Parameter(torch.zeros(hidden_size)) for _ in range(num_layers) ]) def reinit(self, *pre_mixins): start_layer = len(self.transformer.layers) - self.num_layers assert start_layer >= 0 for layer_id in range(self.num_layers): old_attention = self.transformer.layers[start_layer + layer_id].attention self.query_key_value[layer_id].weight.data.copy_(old_attention.query_key_value.weight.data) self.query_key_value[layer_id].bias.data.copy_(old_attention.query_key_value.bias.data) def attention_extra(self, frame_hidden_state, layer_id, attn_dropout, text_hidden_state=None, text_attn_mask=None, mode_sequential=True): # pb relax swin_pb_relax = True alpha = 16 # frame_hidden_state [batchsize, frame_num*frame_size, n_head*hiddensize_perhead] if not self.mask_initialized: self.attn_mask_sequential = self.attn_mask_sequential.to(device=frame_hidden_state.device, dtype=frame_hidden_state.dtype) self.causal_mask_sequential = self.causal_mask_sequential.to(device=frame_hidden_state.device, dtype=frame_hidden_state.dtype) self.attn_mask_interp = self.attn_mask_interp.to(device=frame_hidden_state.device, dtype=frame_hidden_state.dtype) self.causal_mask_interp = self.causal_mask_interp.to(device=frame_hidden_state.device, dtype=frame_hidden_state.dtype) self.mask_initialized = True b0, s1, h0 = frame_hidden_state.shape h = h0 // self.n_head frame_len = self.frame_resolution * self.frame_resolution frame_num = s1 // frame_len assert frame_num*frame_len == s1 wind_square = self.window_size * self.window_size nW = frame_len // wind_square bswin = b0 * nW causal_mask = self.causal_mask_sequential if mode_sequential else self.causal_mask_interp attn_mask = self.attn_mask_sequential if mode_sequential else self.attn_mask_interp if text_hidden_state is not None: s0 = text_hidden_state.shape[1] qkv_text = self.query_key_value[layer_id](text_hidden_state).reshape(b0, s0, 3, self.n_head, h).permute(2, 0, 3, 1, 4) #[3, b0, n_head, s0, h] q_text, k_text, v_text = qkv_text[0], qkv_text[1], qkv_text[2] # shift frame_hidden_state = frame_hidden_state.reshape(b0, frame_num, self.frame_resolution, self.frame_resolution, h0) if self.shift_sizes[layer_id%2] > 0: frame_hidden_state = torch.roll(frame_hidden_state, shifts=(-self.shift_sizes[layer_id%2], -self.shift_sizes[layer_id%2]), dims=(2,3)) # window partition frame_hidden_state = window_partition(frame_hidden_state, self.window_size).reshape(bswin, frame_num*wind_square, h0) qkv = self.query_key_value[layer_id](frame_hidden_state).reshape(bswin, frame_num*wind_square, 3, self.n_head, h)\ .permute(2, 0, 3, 1, 4) #[3, bswin, n_head, frame_num*wind_size*wind_size, h] q, k, v = qkv[0], qkv[1], qkv[2] # pb-relax if swin_pb_relax: attn = torch.matmul(q / (math.sqrt(h)*alpha), k.transpose(-1, -2)) else: attn = torch.matmul(q / math.sqrt(h), k.transpose(-1, -2)) if self.shift_sizes[layer_id%2] > 0: # attn = attn.view(bswin // nW, nW, self.n_head, frame_num*wind_square, frame_num*wind_square) + self.attn_mask.unsqueeze(1).unsqueeze(0) attn = torch.mul(attn.view(bswin // nW, nW, self.n_head, frame_num*wind_square, frame_num*wind_square), attn_mask)\ - 10000.0 * (1.0 - attn_mask) attn = attn.view(bswin, self.n_head, frame_num*wind_square, frame_num*wind_square) else: attn = torch.mul(attn.view(bswin // nW, nW, self.n_head, frame_num*wind_square, frame_num*wind_square), causal_mask)\ - 10000.0 * (1.0 - causal_mask) attn = attn.view(bswin, self.n_head, frame_num*wind_square, frame_num*wind_square) if swin_pb_relax: swin_pb_relax_const = torch.max(attn.reshape(bswin, self.n_head, -1), dim=-1, keepdim=True)[0].detach().unsqueeze(-1) attn = (attn - swin_pb_relax_const)*alpha if text_hidden_state is None: attn = F.softmax(attn, dim=-1) if attn_dropout is not None: with get_cuda_rng_tracker().fork(): attn = attn_dropout(attn) context_swin = torch.matmul(attn, v).permute(0, 2, 1, 3).reshape(bswin, frame_num, self.window_size, self.window_size, h0) else: assert text_attn_mask is not None text_attn_mask = text_attn_mask.unsqueeze(2).unsqueeze(2) # pb-relax if swin_pb_relax: attn_frame2text = torch.matmul(q.reshape(b0, -1, self.n_head, frame_num*wind_square, h) / (math.sqrt(h)*alpha), k_text.unsqueeze(1).transpose(-1, -2)) attn_frame2text = (attn_frame2text-swin_pb_relax_const.reshape(b0, -1, self.n_head, 1, 1))*alpha else: attn_frame2text = torch.matmul(q.reshape(b0, -1, self.n_head, frame_num*wind_square, h) / math.sqrt(h), k_text.unsqueeze(1).transpose(-1, -2)) attn_frame2text = torch.mul(text_attn_mask, attn_frame2text) - 10000.0 * (1.0 - text_attn_mask) attn_frame2text = attn_frame2text.reshape(bswin, self.n_head, frame_num*wind_square, s0) attn = torch.cat((attn, attn_frame2text), dim=-1) attn = F.softmax(attn, dim=-1) if attn_dropout is not None: with get_cuda_rng_tracker().fork(): attn = attn_dropout(attn) context_swin = (torch.matmul(attn[..., :-s0], v) + torch.matmul(attn[..., -s0:].reshape(b0, -1, self.n_head,frame_num*wind_square, s0), v_text.unsqueeze(1))\ .reshape(bswin, self.n_head, frame_num*wind_square, h))\ .permute(0, 2, 1, 3).reshape(bswin, frame_num, self.window_size, self.window_size, h0) context_swin = window_reverse(context_swin, self.window_size, self.frame_resolution, self.frame_resolution) # reverse cycle shift if self.shift_sizes[layer_id%2] > 0: context_swin = torch.roll(context_swin, shifts=(self.shift_sizes[layer_id%2], self.shift_sizes[layer_id%2]), dims=(2,3)) context_swin = context_swin.reshape(b0, s1, h0) return context_swin class FullAttentionMixin(BaseMixin): def __init__(self, num_layers, hidden_size, frame_resolution, n_head, frame_num, init_method=unscaled_init_method(0.02), output_layer_init_method=unscaled_init_method(0.02), ): super(FullAttentionMixin, self).__init__() self.num_layers = num_layers # replace attention in the LAST n layers self.query_key_value = torch.nn.ModuleList( [ColumnParallelLinear(hidden_size, 3*hidden_size,stride=3, gather_output=False,init_method=init_method) for layer_id in range(num_layers) ]) self.dense = torch.nn.ModuleList( [RowParallelLinear( hidden_size, hidden_size, input_is_parallel=True, init_method=output_layer_init_method, bias=True, module=self, name="dense",) for layer_id in range(num_layers) ]) self.n_head = n_head self.frame_resolution = frame_resolution self.frame_len = frame_resolution * frame_resolution self.causal_mask = torch.ones(1, 1, self.frame_len*frame_num, self.frame_len*frame_num).tril() self.mask_initialized = False self.attn_distribution = torch.nn.ParameterList([ torch.nn.Parameter(torch.zeros(hidden_size)) for _ in range(num_layers) ]) def reinit(self, *pre_mixins): start_layer = len(self.transformer.layers) - self.num_layers assert start_layer >= 0 for layer_id in range(self.num_layers): base_attention = self.transformer.layers[start_layer + layer_id].attention self.query_key_value[layer_id].weight.data.copy_(base_attention.query_key_value.weight.data) self.query_key_value[layer_id].bias.data.copy_(base_attention.query_key_value.bias.data) def attention_extra(self, frame_hidden_state, layer_id, attn_dropout, text_hidden_state=None, text_attn_mask=None, mode_sequential=False): # pb relax # frame_hidden_state [batchsize, frame_num*frame_size, n_head*hiddensize_perhead] assert mode_sequential == True # only swin_pb_relax = True alpha = 16 if not self.mask_initialized: self.causal_mask = self.causal_mask.to(device=frame_hidden_state.device, dtype=frame_hidden_state.dtype) self.mask_initialized = True b0, s1, h0 = frame_hidden_state.shape h = h0 // self.n_head frame_len = self.frame_resolution * self.frame_resolution frame_num = s1 // frame_len assert frame_num*frame_len == s1 qkv = self.query_key_value[layer_id](frame_hidden_state).reshape(b0, s1, 3, self.n_head, h)\ .permute(2, 0, 3, 1, 4) #[3, b0, n_head, s1, h] q, k, v = qkv[0], qkv[1], qkv[2] # frames-to-frames if swin_pb_relax: attn = torch.matmul(q / (math.sqrt(h)*alpha), k.transpose(-1, -2)) else: attn = torch.matmul(q / math.sqrt(h), k.transpose(-1, -2)) attn = torch.mul(attn, self.causal_mask) - 10000.0 * (1.0 - self.causal_mask) if swin_pb_relax: swin_pb_relax_const = torch.max(attn.reshape(b0, self.n_head, -1), dim=-1, keepdim=True)[0].detach().unsqueeze(-1) attn = (attn - swin_pb_relax_const)*alpha if text_hidden_state is None: attn = F.softmax(attn, dim=-1) if attn_dropout is not None: with get_cuda_rng_tracker().fork(): attn = attn_dropout(attn) context_swin = torch.matmul(attn, v).permute(0, 2, 1, 3).reshape(b0, s1, h0) else: # frame-to-text assert text_attn_mask is not None s0 = text_hidden_state.shape[1] qkv_text = self.query_key_value[layer_id](text_hidden_state).reshape(b0, s0, 3, self.n_head, h).permute(2, 0, 3, 1, 4) #[3, b0, n_head, s0, h] q_text, k_text, v_text = qkv_text[0], qkv_text[1], qkv_text[2] text_attn_mask = text_attn_mask.unsqueeze(2) if swin_pb_relax: attn_frame2text = torch.matmul(q.reshape(b0, self.n_head, s1, h) / (math.sqrt(h)*alpha), k_text.transpose(-1, -2)) attn_frame2text = (attn_frame2text-swin_pb_relax_const.reshape(b0, self.n_head, 1, 1))*alpha else: attn_frame2text = torch.matmul(q.reshape(b0, self.n_head, s1, h) / math.sqrt(h), k_text.transpose(-1, -2)) attn_frame2text = torch.mul(text_attn_mask, attn_frame2text) - 10000.0 * (1.0 - text_attn_mask) attn_frame2text = attn_frame2text.reshape(b0, self.n_head, s1, s0) attn = torch.cat((attn, attn_frame2text), dim=-1) attn = F.softmax(attn, dim=-1) if attn_dropout is not None: with get_cuda_rng_tracker().fork(): attn = attn_dropout(attn) context_frame = (torch.matmul(attn[..., :-s0], v) + torch.matmul(attn[..., -s0:].reshape(b0, self.n_head,s1, s0), v_text))\ .permute(0, 2, 1, 3).reshape(b0, s1, h0) return context_frame def attention_localframe_and_text(q0, k0, v0, attention_mask_totxt, attention_mask_local, n_head, text_len, frame_len, frame_num, attention_dropout=None, layer_id=0, **kwargs): b, s0, h0 = q0.shape s1 = s0 - text_len h = h0 // n_head assert q0.shape[1] == v0.shape[1] == k0.shape[1] == text_len+frame_len*frame_num # attention_mask_totxt [b, 1, 1, text_len] # attention_mask_local [1, 1, frame_num, frame_len, frame_len] # attention_mask: [1, 1, text_len+frame_len, text_len+frame_len] q0 = q0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3) v0 = v0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3) k0 = k0.reshape(b, s0, n_head, h).permute(0, 2, 1, 3) k0T = k0.transpose(-1, -2) # score: any2text score_any2text = torch.matmul(q0 / math.sqrt(q0.shape[-1]), k0T[..., :text_len]) score_any2text_part1 = torch.mul(score_any2text[..., :text_len, :], attention_mask_totxt) \ - 10000.0 * (1.0 - attention_mask_totxt) score_any2text_part2 = torch.mul(score_any2text[..., text_len:, :], attention_mask_totxt) - \ 10000.0 * (1.0 - attention_mask_totxt) # score: frame local q0_frame = q0[:, :, text_len:].reshape(b, n_head, frame_num, frame_len, h) v0_frame = v0[:, :, text_len:].reshape(b, n_head, frame_num, frame_len, h) k0T_frame = k0[:, :, text_len:].reshape(b, n_head, frame_num, frame_len, h).transpose(-1, -2) score_frame_local0 = torch.matmul(q0_frame / math.sqrt(q0_frame.shape[-1]), k0T_frame) score_frame_local0 = torch.mul(score_frame_local0, attention_mask_local) \ - 10000.0 * (1.0 - attention_mask_local) # context for frame score_frame_all = torch.cat((score_any2text_part2, score_frame_local0.view(b, n_head, s1, frame_len)), dim=-1) attention_probs_frame = F.softmax(score_frame_all, dim=-1) if attention_dropout is not None: with get_cuda_rng_tracker().fork(): attention_probs_frame = attention_dropout(attention_probs_frame) context_frame2text = torch.matmul(attention_probs_frame[..., :text_len], v0[..., :text_len, :]) # [b, n_head, s1, h] context_frame_local0 = torch.matmul(attention_probs_frame[..., text_len:text_len+frame_len].\ view(b, n_head, frame_num, frame_len, frame_len), v0_frame).view(b, n_head, s1, h) context_frame = (context_frame2text + context_frame_local0).transpose(1, 2).reshape(b, s1, h0) # context for text attention_probs_text = F.softmax(score_any2text_part1, dim=-1) if attention_dropout is not None: with get_cuda_rng_tracker().fork(): attention_probs_text = attention_dropout(attention_probs_text) context_text2text = torch.matmul(attention_probs_text, v0[..., :text_len, :]) context_text2text = context_text2text.transpose(1, 2).reshape(b, text_len, h0) return context_text2text, context_frame class CogVideoModel(BaseModel): def __init__(self, args, transformer=None, parallel_output=True): super().__init__(args, transformer=transformer, parallel_output=parallel_output) self.stage = args.cogvideo_stage # 1 or 2 self.mode_sequential = True if self.stage==1 else False self.layout = args.layout # [64, 64+400, 64+5*400] self.n_head = args.num_attention_heads frame_resolution = int(math.sqrt(self.layout[1]-self.layout[0])) frame_num = (args.layout[2]-args.layout[0])//(args.layout[1]-args.layout[0]) frame_len = self.layout[1]-self.layout[0] self.add_mixin('extra_position_embedding', PositionEmbeddingMixin( args.additional_seqlen, args.hidden_size )) if args.window_size == -1: # full attention assert self.stage == 1 self.add_mixin('attention_plus', FullAttentionMixin( num_layers=args.num_layers, hidden_size=args.hidden_size, frame_resolution=frame_resolution, n_head=args.num_attention_heads, frame_num=frame_num, )) else: self.add_mixin('attention_plus', WindowAttentionMixin( num_layers=args.num_layers, hidden_size=args.hidden_size, frame_resolution=frame_resolution, window_size=args.window_size, shift_size=args.window_size//2, n_head=args.num_attention_heads, frame_num=frame_num, )) # attention_mask_local self.attention_mask_local_sequential = torch.ones(1, 1, frame_num, frame_len, frame_len).tril().unsqueeze(0) self.attention_mask_local_interp = torch.ones(1, 1, frame_num, frame_len, frame_len) for idx in range(1, frame_num, 2): self.attention_mask_local_interp[:, :, idx:idx+1].tril_() self.attention_mask_local_interp = self.attention_mask_local_interp.unsqueeze(0) self.mask_initialized = False @classmethod def add_model_specific_args(cls, parser): group = parser.add_argument_group('CogVideoModel', 'CogVideo model configurations') group.add_argument("--layout", type=str, default='64, 464, 2064', help='text_len, textlen+frame_len, textlen+frame_len*frame_num') group.add_argument("--window-size", type=int, default=10, help="swin attention's window size in temperal channel, -1 represents full attention") group.add_argument("--additional-seqlen", type=int, default=2000) group.add_argument("--cogvideo-stage", type=int, default=1, choices=[1,2]) return parser def disable_untrainable_params(self): self.transformer.requires_grad_(False) def position_embedding_forward(self, position_ids, **kw_args): position = position_ids[..., :(64+400)] position_plus = position_ids[..., (64+400):] position_embeddings = torch.cat( ( self.transformer.position_embeddings(position), self.get_mixin('extra_position_embedding').position_embeddings(position_plus-(512+400)) ), dim=-2 ) return position_embeddings def attention_forward(self, hidden_states, mask, layer_id, **kw_args): # mask.shape=[bs, 1, 1, 64] if not self.mask_initialized: self.attention_mask_local_sequential = self.attention_mask_local_sequential.to(device=hidden_states.device, dtype=hidden_states.dtype) self.attention_mask_local_interp = self.attention_mask_local_interp.to(device=hidden_states.device, dtype=hidden_states.dtype) self.mask_initialized = True attn_module = self.transformer.layers[layer_id].attention hidden_size = hidden_states.shape[-1] bs = hidden_states.shape[0] # base model qkv mixed_raw_layer = attn_module.query_key_value(hidden_states) q0, k0, v0 = split_tensor_along_last_dim(mixed_raw_layer, 3) dropout_fn = self.transformer.layers[layer_id].attention.attention_dropout if self.training else None attention_mask_local = self.attention_mask_local_sequential if self.mode_sequential else self.attention_mask_local_interp context_text, context_frame_local_text = attention_localframe_and_text( q0, k0, v0, attention_mask_totxt=mask, attention_mask_local=attention_mask_local, n_head=attn_module.num_attention_heads_per_partition, text_len=self.layout[0], frame_len=self.layout[1]-self.layout[0], frame_num=(self.layout[2]-self.layout[0])//(self.layout[1]-self.layout[0]), attention_dropout=dropout_fn, layer_id=layer_id, ) context_frame_swin = self.get_mixin('attention_plus').attention_extra( hidden_states[:, self.layout[0]:], layer_id, dropout_fn, text_hidden_state=hidden_states[:, :self.layout[0]], text_attn_mask=mask[..., 0, :], mode_sequential=self.mode_sequential) attn_distrib = torch.sigmoid(self.get_mixin('attention_plus').attn_distribution[layer_id]) attn_distrib = attn_distrib.unsqueeze(0).unsqueeze(0) output_text = attn_module.dense(context_text) output_frame = torch.mul(attn_module.dense(context_frame_local_text), attn_distrib)\ +torch.mul(self.get_mixin('attention_plus').dense[layer_id](context_frame_swin), 1-attn_distrib) output = torch.cat((output_text, output_frame), dim=-2) return output