Update difpoint/src/models/predictor.py

difpoint/src/models/predictor.py

@@ -10,7 +10,8 @@ import torch
 from torch.cuda import nvtx
 from collections import OrderedDict
 import platform
-import spaces
+
+import spaces 
 
 try:
     import tensorrt as trt
@@ -40,15 +41,15 @@ class TensorRTPredictor:
     """
     Implements inference for the EfficientDet TensorRT engine.
     """
-
+    @spaces.GPU
     def __init__(self, **kwargs):
         """
         :param engine_path: The path to the serialized engine to load from disk.
         """
         if platform.system().lower() == 'linux':
-            ctypes.CDLL("./downloaded_repo/pretrained_weights/liveportrait_onnx/libgrid_sample_3d_plugin.so", mode=ctypes.RTLD_GLOBAL)
+            ctypes.CDLL("./difpoint/checkpoints/liveportrait_onnx/libgrid_sample_3d_plugin.so", mode=ctypes.RTLD_GLOBAL)
         else:
-            ctypes.CDLL("./downloaded_repo/pretrained_weights/liveportrait_onnx/grid_sample_3d_plugin.dll", mode=ctypes.RTLD_GLOBAL)
+            ctypes.CDLL("./difpoint/checkpoints/liveportrait_onnx/grid_sample_3d_plugin.dll", mode=ctypes.RTLD_GLOBAL)
         # Load TRT engine
         self.logger = trt.Logger(trt.Logger.VERBOSE)
         trt.init_libnvinfer_plugins(self.logger, "")
@@ -172,7 +173,6 @@ class TensorRTPredictor:
         del self.outputs
         del self.tensors
 
-
 class OnnxRuntimePredictor:
     """
     OnnxRuntime Prediction
@@ -190,7 +190,10 @@ 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
-        self.onnx_model = onnxruntime.InferenceSession(model_path, providers=providers, sess_options=opts)
+        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()