Update difpoint/src/models/predictor.py

difpoint/src/models/predictor.py

@@ -190,13 +190,21 @@ class OnnxRuntimePredictor:
         # opts.inter_op_num_threads = kwargs.get("num_threads", 4)
         # opts.intra_op_num_threads = kwargs.get("num_threads", 4)
         # opts.log_severity_level = 3
-        def recreate_session():
-            return onnxruntime.InferenceSession(model_path, providers=providers, sess_options=opts)
-
-        self.onnx_model = recreate_session()
-        self.inputs = self.onnx_model.get_inputs()
-        self.outputs = self.onnx_model.get_outputs()
+        #self.onnx_model = onnxruntime.InferenceSession(model_path, providers=providers, sess_options=opts)
+        #self.inputs = self.onnx_model.get_inputs()
+        #self.outputs = self.onnx_model.get_outputs()
+        self.onnx_model = None
 
+    def _load_model(self):
+        """Lazy initialization of the ONNX model (only when needed)."""
+        if self.onnx_model is None:
+            providers = ['CUDAExecutionProvider', 'CoreMLExecutionProvider', 'CPUExecutionProvider']
+            print(f"OnnxRuntime use {providers}")
+            opts = onnxruntime.SessionOptions()
+            self.onnx_model = onnxruntime.InferenceSession(self.model_path, providers=providers, sess_options=opts)
+            self.inputs = self.onnx_model.get_inputs()
+            self.outputs = self.onnx_model.get_outputs()
+            
     def input_spec(self):
         """
         Get the specs for the input tensor of the network. Useful to prepare memory allocations.
@@ -222,6 +230,7 @@ class OnnxRuntimePredictor:
         return specs
 
     def predict(self, *data):
+        self._load_model()
         input_feeds = {}
         for i in range(len(data)):
             if self.inputs[i].type == 'tensor(float16)':