fix to _make_causal_mask and _expand_mask

external/arc2face/models.py

@@ -6,28 +6,27 @@ from transformers.modeling_outputs import BaseModelOutputWithPooling
 try:
     from transformers.models.clip.modeling_clip import _make_causal_mask, _expand_mask
 except ImportError:
-    from transformers.modeling_attn_mask_utils import _create_4d_causal_attention_mask
-
+    # transformers >=4.47 removed these internal helpers from modeling_clip.
+    # Reimplement them directly from the transformers 4.34 source so the mask
+    # format (additive, shape [bsz,1,tgt,src]) matches what CLIPEncoder expects.
     def _make_causal_mask(input_ids_shape, dtype, device, past_key_values_length=0):
-        batch_size, tgt_len = input_ids_shape
-        return _create_4d_causal_attention_mask(
-            input_shape=(batch_size, tgt_len),
-            dtype=dtype,
-            device=device,
-            past_key_values_length=past_key_values_length,
-        )
+        bsz, tgt_len = input_ids_shape
+        mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)
+        mask_cond = torch.arange(tgt_len, device=device)
+        mask.masked_fill_(mask_cond < (mask_cond + 1).view(tgt_len, 1), 0)
+        mask = mask.to(dtype)
+        if past_key_values_length > 0:
+            mask = torch.cat(
+                [torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1
+            )
+        return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
 
     def _expand_mask(mask, dtype, tgt_len=None):
-        batch_size, src_len = mask.shape
+        bsz, src_len = mask.shape
         tgt_len = tgt_len if tgt_len is not None else src_len
-
-        expanded_mask = mask[:, None, None, :].expand(batch_size, 1, tgt_len, src_len)
-        inverted_mask = 1.0 - expanded_mask.to(dtype)
-
-        return inverted_mask.masked_fill(
-            inverted_mask.to(torch.bool),
-            torch.finfo(dtype).min,
-        )
+        expanded = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
+        inverted = 1.0 - expanded
+        return inverted.masked_fill(inverted.to(torch.bool), torch.finfo(dtype).min)
 
 class CLIPTextModelWrapper(CLIPTextModel):
     # Adapted from https://github.com/huggingface/transformers/blob/v4.34.1/src/transformers/models/clip/modeling_clip.py#L812