[Fix] try import expand mask from transformers (#2064)

* try import expand mask from transformers * fix lint * fix lint * fix lint
open-mmlab · Nov 10, 2023 · 2a95e50 · 2a95e50
1 parent c2f8f3a
commit 2a95e50
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/mmagic/models/editors/fastcomposer/fastcomposer_util.py b/mmagic/models/editors/fastcomposer/fastcomposer_util.py
@@ -12,7 +12,13 @@
 from transformers import (CLIPModel, CLIPPreTrainedModel, CLIPTextModel,
                           CLIPVisionConfig, CLIPVisionModel)
 from transformers.modeling_outputs import BaseModelOutputWithPooling
-from transformers.models.clip.modeling_clip import _expand_mask
+
+from mmagic.utils import try_import
+
+_expand_mask = try_import('transformers.models.clip.modeling_clip')
+if _expand_mask is None:
+    _expand_mask = try_import(
+        'ransformers.models.clip.modeling_clip._prepare_4d_attention_mask')
 
 
 class FastComposerModel(nn.Module):