levihsu commited on
Commit
005ec87
1 Parent(s): ce7bb03

Update preprocess/humanparsing/run_parsing.py

Browse files
preprocess/humanparsing/run_parsing.py CHANGED
@@ -11,12 +11,12 @@ import torch
11
 
12
  class Parsing:
13
  def __init__(self, gpu_id: int):
14
- self.gpu_id = gpu_id
15
- torch.cuda.set_device(gpu_id)
16
  session_options = ort.SessionOptions()
17
  session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
18
  session_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
19
- session_options.add_session_config_entry('gpu_id', str(gpu_id))
20
  self.session = ort.InferenceSession(os.path.join(Path(__file__).absolute().parents[2].absolute(), 'checkpoints/humanparsing/parsing_atr.onnx'),
21
  sess_options=session_options, providers=['CPUExecutionProvider'])
22
  self.lip_session = ort.InferenceSession(os.path.join(Path(__file__).absolute().parents[2].absolute(), 'checkpoints/humanparsing/parsing_lip.onnx'),
@@ -24,6 +24,6 @@ class Parsing:
24
 
25
 
26
  def __call__(self, input_image):
27
- torch.cuda.set_device(self.gpu_id)
28
  parsed_image, face_mask = onnx_inference(self.session, self.lip_session, input_image)
29
  return parsed_image, face_mask
 
11
 
12
  class Parsing:
13
  def __init__(self, gpu_id: int):
14
+ # self.gpu_id = gpu_id
15
+ # torch.cuda.set_device(gpu_id)
16
  session_options = ort.SessionOptions()
17
  session_options.graph_optimization_level = ort.GraphOptimizationLevel.ORT_ENABLE_ALL
18
  session_options.execution_mode = ort.ExecutionMode.ORT_SEQUENTIAL
19
+ # session_options.add_session_config_entry('gpu_id', str(gpu_id))
20
  self.session = ort.InferenceSession(os.path.join(Path(__file__).absolute().parents[2].absolute(), 'checkpoints/humanparsing/parsing_atr.onnx'),
21
  sess_options=session_options, providers=['CPUExecutionProvider'])
22
  self.lip_session = ort.InferenceSession(os.path.join(Path(__file__).absolute().parents[2].absolute(), 'checkpoints/humanparsing/parsing_lip.onnx'),
 
24
 
25
 
26
  def __call__(self, input_image):
27
+ # torch.cuda.set_device(self.gpu_id)
28
  parsed_image, face_mask = onnx_inference(self.session, self.lip_session, input_image)
29
  return parsed_image, face_mask