Make ONNX runtime optional (#215)

Author: juanmc2005

requirements.txt

@@ -18,4 +18,3 @@ optuna>=2.10
 websocket-server>=0.6.4
 websocket-client>=0.58.0
 rich>=12.5.1
-onnxruntime-gpu>=1.16.1

setup.cfg

@@ -40,7 +40,6 @@ install_requires=
     websocket-server>=0.6.4
     websocket-client>=0.58.0
     rich>=12.5.1
-    onnxruntime-gpu>=1.16.1
 
 [options.packages.find]
 where=src

src/diart/models.py

@@ -5,7 +5,6 @@
 from typing import Optional, Text, Union, Callable, List
 
 import numpy as np
-import onnxruntime
 import torch
 import torch.nn as nn
 from requests import HTTPError
@@ -15,9 +14,16 @@
     from pyannote.audio.pipelines.speaker_verification import PretrainedSpeakerEmbedding
     from pyannote.audio.utils.powerset import Powerset
 
-    _has_pyannote = True
+    IS_PYANNOTE_AVAILABLE = True
 except ImportError:
-    _has_pyannote = False
+    IS_PYANNOTE_AVAILABLE = False
+
+try:
+    import onnxruntime as ort
+
+    IS_ONNX_AVAILABLE = True
+except ImportError:
+    IS_ONNX_AVAILABLE = False
 
 
 class PowersetAdapter(nn.Module):
@@ -88,11 +94,9 @@ def execution_provider(self) -> str:
         return f"{device}ExecutionProvider"
 
     def recreate_session(self):
-        options = onnxruntime.SessionOptions()
-        options.graph_optimization_level = (
-            onnxruntime.GraphOptimizationLevel.ORT_ENABLE_ALL
-        )
-        self.session = onnxruntime.InferenceSession(
+        options = ort.SessionOptions()
+        options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
+        self.session = ort.InferenceSession(
             self.path,
             sess_options=options,
             providers=[self.execution_provider],
@@ -168,7 +172,7 @@ def from_pyannote(
         -------
         wrapper: SegmentationModel
         """
-        assert _has_pyannote, "No pyannote.audio installation found"
+        assert IS_PYANNOTE_AVAILABLE, "No pyannote.audio installation found"
         return SegmentationModel(PyannoteLoader(model, use_hf_token))
 
     @staticmethod
@@ -177,6 +181,7 @@ def from_onnx(
         input_name: str = "waveform",
         output_name: str = "segmentation",
     ) -> "SegmentationModel":
+        assert IS_ONNX_AVAILABLE, "No ONNX installation found"
         return SegmentationModel(ONNXLoader(model_path, [input_name], output_name))
 
     @staticmethod
@@ -224,7 +229,7 @@ def from_pyannote(
         -------
         wrapper: EmbeddingModel
         """
-        assert _has_pyannote, "No pyannote.audio installation found"
+        assert IS_PYANNOTE_AVAILABLE, "No pyannote.audio installation found"
         loader = PyannoteLoader(model, use_hf_token)
         return EmbeddingModel(loader)
 
@@ -234,6 +239,7 @@ def from_onnx(
         input_names: List[str] | None = None,
         output_name: str = "embedding",
     ) -> "EmbeddingModel":
+        assert IS_ONNX_AVAILABLE, "No ONNX installation found"
         input_names = input_names or ["waveform", "weights"]
         loader = ONNXLoader(model_path, input_names, output_name)
         return EmbeddingModel(loader)