fix: openclip compatable

jina-ai · Jul 26, 2022 · 0ba7fad · 0ba7fad
1 parent ab23235
commit 0ba7fad
Show file tree

Hide file tree

Showing 2 changed files with 4 additions and 4 deletions.
diff --git a/server/clip_server/model/openclip_model.py b/server/clip_server/model/openclip_model.py
@@ -15,9 +15,8 @@
 from open_clip.model import (
     CLIP,
     convert_weights_to_fp16,
-    build_model_from_openai_state_dict,
 )
-from open_clip.factory import _MODEL_CONFIGS, load_checkpoint, load_openai_model
+from open_clip.factory import _MODEL_CONFIGS, load_state_dict, load_openai_model
 
 if TYPE_CHECKING:
     import torch
@@ -48,7 +47,8 @@ def __init__(self, name: str, device: str = 'cpu', jit: bool = False, **kwargs):
 
             self._model = CLIP(**model_cfg)
 
-            load_checkpoint(self._model, model_path)
+            state_dict = load_state_dict(model_path)
+            self._model.load_state_dict(state_dict, strict=True)
 
             if str(device) == 'cuda':
                 convert_weights_to_fp16(self._model)

diff --git a/server/setup.py b/server/setup.py
@@ -49,7 +49,7 @@
         'torchvision',
         'jina>=3.6.0',
         'prometheus-client',
-        'open_clip_torch',
+        'open_clip_torch>=1.3.0',
     ],
     extras_require={
         'onnx': [