Fabrice-TIERCELIN commited on
Commit
e098237
·
verified ·
1 Parent(s): 190aa8e

Update diffusers_helper/models/hunyuan_video_packed.py

Browse files
diffusers_helper/models/hunyuan_video_packed.py CHANGED
@@ -362,7 +362,7 @@ class HunyuanVideoIndividualTokenRefiner(nn.Module):
362
  batch_size = attention_mask.shape[0]
363
  seq_len = attention_mask.shape[1]
364
  attention_mask = attention_mask.to(hidden_states.device).bool()
365
- self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).repeat(1, 1, seq_len, 1)
366
  self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
367
  self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
368
  self_attn_mask[:, :, :, 0] = True
 
362
  batch_size = attention_mask.shape[0]
363
  seq_len = attention_mask.shape[1]
364
  attention_mask = attention_mask.to(hidden_states.device).bool()
365
+ self_attn_mask_1 = attention_mask.view(batch_size, 1, 1, seq_len).expand(-1, -1, seq_len, -1)
366
  self_attn_mask_2 = self_attn_mask_1.transpose(2, 3)
367
  self_attn_mask = (self_attn_mask_1 & self_attn_mask_2).bool()
368
  self_attn_mask[:, :, :, 0] = True