Add safeglobals to allow metaclip models to load with weights_only=True, add bigG metaclip while working with them

Author: rwightman

src/open_clip/factory.py

@@ -26,6 +26,20 @@
 _MODEL_CONFIGS = {}  # directory (model_name: config) of model architecture configs
 
 
+try:
+    import _codecs
+    import numpy as np
+    # add safe globals, known to be needed for metaclip weights
+    torch.serialization.add_safe_globals([
+        _codecs.encode,  # now in pytorch main but some pytorch versions w/ weights_only flag don't have it
+        np.core.multiarray.scalar,
+        np.dtype,
+        np.dtypes.Float64DType,
+    ])
+except Exception:
+    pass
+
+
 def _natural_key(string_):
     return [int(s) if s.isdigit() else s for s in re.split(r'(\d+)', string_.lower())]
 

src/open_clip/model_configs/ViT-bigG-14-quickgelu.json

@@ -0,0 +1,19 @@
+{
+    "embed_dim": 1280,
+    "quick_gelu": true,
+    "vision_cfg": {
+        "image_size": 224,
+        "layers": 48,
+        "width": 1664,
+        "head_width": 104,
+        "mlp_ratio": 4.9231,
+        "patch_size": 14
+    },
+    "text_cfg": {
+        "context_length": 77,
+        "vocab_size": 49408,
+        "width": 1280,
+        "heads": 20,
+        "layers": 32
+    }
+}

src/open_clip/pretrained.py

@@ -280,6 +280,10 @@ def _mccfg(url='', hf_hub='', **kwargs):
     laion2b_s39b_b160k=_pcfg(hf_hub='laion/CLIP-ViT-bigG-14-laion2B-39B-b160k/'),
 )
 
+_VITbigG14_quickgelu = dict(
+    metaclip_fullcc=_pcfg(url='https://dl.fbaipublicfiles.com/MMPT/metaclip/G14_fullcc2.5b.pt'),
+)
+
 _robertaViTB32 = dict(
     laion2b_s12b_b32k=_pcfg(hf_hub='laion/CLIP-ViT-B-32-roberta-base-laion2B-s12B-b32k/'),
 )
@@ -356,6 +360,7 @@ def _mccfg(url='', hf_hub='', **kwargs):
     "ViT-H-14-378-quickgelu": _VITH14_378_quickgelu,
     "ViT-g-14": _VITg14,
     "ViT-bigG-14": _VITbigG14,
+    "ViT-bigG-14-quickgelu": _VITbigG14_quickgelu,
 
     "roberta-ViT-B-32": _robertaViTB32,
     "xlm-roberta-base-ViT-B-32": _xlmRobertaBaseViTB32,