feat: support not yet supported onnx models

jina-ai · Sep 22, 2022 · f48850f · f48850f
1 parent c690c24
commit f48850f
Showing 1 changed file with 26 additions and 2 deletions.
diff --git a/server/clip_server/model/clip_onnx.py b/server/clip_server/model/clip_onnx.py
@@ -218,10 +218,34 @@ def start_sessions(
     ):
         import onnxruntime as ort
 
-        self._visual_session = ort.InferenceSession(self._visual_path, **kwargs)
+        if self._visual_path.endswith('.zip'):
+            import zipfile
+            import tempfile
+
+            with zipfile.ZipFile(
+                self._visual_path, 'r'
+            ) as zip_ref, tempfile.TemporaryDirectory() as tmp_dir:
+                zip_ref.extractall(tmp_dir)
+                self._visual_session = ort.InferenceSession(
+                    tmp_dir + '/visual.onnx', **kwargs
+                )
+        else:
+            self._visual_session = ort.InferenceSession(self._visual_path, **kwargs)
         self._visual_session.disable_fallback()
 
-        self._textual_session = ort.InferenceSession(self._textual_path, **kwargs)
+        if self._textual_path.endswith('.zip'):
+            import zipfile
+            import tempfile
+
+            with zipfile.ZipFile(
+                self._textual_path, 'r'
+            ) as zip_ref, tempfile.TemporaryDirectory() as tmp_dir:
+                zip_ref.extractall(tmp_dir)
+                self._textual_session = ort.InferenceSession(
+                    tmp_dir + '/textual.onnx', **kwargs
+                )
+        else:
+            self._textual_session = ort.InferenceSession(self._textual_path, **kwargs)
         self._textual_session.disable_fallback()
 
     def encode_image(self, image_input: Dict):