feat: support custom onnx file and update model signatures

server/clip_server/executors/clip_onnx.py

@@ -23,6 +23,7 @@ def __init__(
         num_worker_preprocess: int = 4,
         minibatch_size: int = 32,
         traversal_paths: str = '@r',
+        model_path: Optional[str] = None,
         **kwargs,
     ):
         super().__init__(**kwargs)
@@ -33,7 +34,7 @@ def __init__(
         self._preprocess_tensor = clip._transform_ndarray(clip.MODEL_SIZE[name])
         self._pool = ThreadPool(processes=num_worker_preprocess)
 
-        self._model = CLIPOnnxModel(name)
+        self._model = CLIPOnnxModel(name, model_path)
 
         import torch
 

server/clip_server/model/clip_onnx.py

@@ -17,15 +17,30 @@
 
 
 class CLIPOnnxModel:
-    def __init__(self, name: str = None):
+    def __init__(self, name: str = None, model_path: str = None):
         if name in _MODELS:
-            cache_dir = os.path.expanduser(f'~/.cache/clip/{name.replace("/", "-")}')
-            self._textual_path = _download(
-                _S3_BUCKET + _MODELS[name][0], cache_dir, with_resume=True
-            )
-            self._visual_path = _download(
-                _S3_BUCKET + _MODELS[name][1], cache_dir, with_resume=True
-            )
+            if model_path is not None:
+                if (
+                    os.path.isdir(model_path)
+                    and os.path.exists(os.path.join(model_path, 'textual.onnx'))
+                    and os.path.exists(os.path.join(model_path, 'visual.onnx'))
+                ):
+                    self._textual_path = os.path.join(model_path, 'textual.onnx')
+                    self._visual_path = os.path.join(model_path, 'visual.onnx')
+                else:
+                    raise RuntimeError(
+                        f'Invalid model path: {model_path}.\nAn valid path should contain \'textual.onnx\' and \'visual.onnx\' files.'
+                    )
+            else:
+                cache_dir = os.path.expanduser(
+                    f'~/.cache/clip/{name.replace("/", "-")}'
+                )
+                self._textual_path = _download(
+                    _S3_BUCKET + _MODELS[name][0], cache_dir, with_resume=True
+                )
+                self._visual_path = _download(
+                    _S3_BUCKET + _MODELS[name][1], cache_dir, with_resume=True
+                )
         else:
             raise RuntimeError(
                 f'Model {name} not found; available models = {available_models()}'

tests/conftest.py

@@ -16,16 +16,26 @@ def random_port():
     return random_port
 
 
-@pytest.fixture(scope='session', params=['onnx', 'torch', 'hg'])
+@pytest.fixture(scope='session', params=['onnx', 'torch', 'hg', 'onnx_custom'])
 def make_flow(port_generator, request):
-    if request.param == 'onnx':
-        from clip_server.executors.clip_onnx import CLIPEncoder
-    elif request.param == 'torch':
-        from clip_server.executors.clip_torch import CLIPEncoder
+    if request.param != 'onnx_custom':
+        if request.param == 'onnx':
+            from clip_server.executors.clip_onnx import CLIPEncoder
+        elif request.param == 'torch':
+            from clip_server.executors.clip_torch import CLIPEncoder
+        else:
+            from clip_server.executors.clip_hg import CLIPEncoder
+
+        f = Flow(port=port_generator()).add(name=request.param, uses=CLIPEncoder)
     else:
-        from clip_server.executors.clip_hg import CLIPEncoder
+        import os
+        from clip_server.executors.clip_onnx import CLIPEncoder
 
-    f = Flow(port=port_generator()).add(name=request.param, uses=CLIPEncoder)
+        f = Flow(port=port_generator()).add(
+            name=request.param,
+            uses=CLIPEncoder,
+            uses_with={'model_path': os.path.expanduser('~/.cache/clip/ViT-B-32')},
+        )
     with f:
         yield f
 

tests/test_server.py

@@ -3,12 +3,12 @@
 import pytest
 from clip_server.model.clip import _transform_ndarray, _transform_blob, _download
 from docarray import Document
+from jina import Flow
 import numpy as np
 
 
 def test_server_download(tmpdir):
     _download('https://docarray.jina.ai/_static/favicon.png', tmpdir, with_resume=False)
-
     target_path = os.path.join(tmpdir, 'favicon.png')
     file_size = os.path.getsize(target_path)
     assert file_size > 0
@@ -25,6 +25,40 @@ def test_server_download(tmpdir):
     assert not os.path.exists(part_path)
 
 
+def test_make_onnx_flow_custom_path_wrong_name(port_generator):
+    from clip_server.executors.clip_onnx import CLIPEncoder
+    import os
+
+    f = Flow(port=port_generator()).add(
+        name='onnx',
+        uses=CLIPEncoder,
+        uses_with={
+            'name': 'ABC',
+            'model_path': os.path.expanduser('~/.cache/clip/ViT-B-32'),
+        },
+    )
+    with pytest.raises(Exception) as info:
+        with f:
+            f.post('/', Document(text='Hello world'))
+
+
+def test_make_onnx_flow_custom_path_wrong_path(port_generator):
+    from clip_server.executors.clip_onnx import CLIPEncoder
+    import os
+
+    f = Flow(port=port_generator()).add(
+        name='onnx',
+        uses=CLIPEncoder,
+        uses_with={
+            'name': 'ViT-B/32',
+            'model_path': 'ABC',
+        },
+    )
+    with pytest.raises(Exception) as info:
+        with f:
+            f.post('/', Document(text='Hello world'))
+
+
 @pytest.mark.parametrize(
     'image_uri',
     [