Acivate GPU Processing in the preprocess process

.gitattributes

@@ -36,4 +36,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.png filter=lfs diff=lfs merge=lfs -text
 *.so filter=lfs diff=lfs merge=lfs -text
 *.jpg filter=lfs diff=lfs merge=lfs -text
-*.jpeg filter=lfs diff=lfs merge=lfs -text
+*.jpeg filter=lfs diff=lfs merge=lfs -text

preprocess/humanparsing/run_parsing.py

@@ -46,8 +46,9 @@ def download_model_file(model_path, url):
 
 class Parsing:
     def __init__(self, gpu_id: int):
-        # self.gpu_id = gpu_id
-        # torch.cuda.set_device(gpu_id)
+        # zeroGPU
+        self.gpu_id = gpu_id
+        torch.cuda.set_device(gpu_id)
         
         # Model paths
         base_path = Path(__file__).absolute().parents[2].absolute()
@@ -74,6 +75,7 @@ class Parsing:
         
 
     def __call__(self, input_image):
-        # torch.cuda.set_device(self.gpu_id)
+        # zeroGPU
+        torch.cuda.set_device(self.gpu_id)
         parsed_image, face_mask = onnx_inference(self.session, self.lip_session, input_image)
         return parsed_image, face_mask

preprocess/openpose/run_openpose.py

@@ -27,12 +27,14 @@ import pdb
 
 class OpenPose:
     def __init__(self, gpu_id: int):
-        # self.gpu_id = gpu_id
-        # torch.cuda.set_device(gpu_id)
+        # zeroGPU
+        self.gpu_id = gpu_id
+        torch.cuda.set_device(gpu_id)
         self.preprocessor = OpenposeDetector()
 
     def __call__(self, input_image, resolution=384):
-        # torch.cuda.set_device(self.gpu_id)
+        # zeroGPU
+        torch.cuda.set_device(self.gpu_id)
         if isinstance(input_image, Image.Image):
             input_image = np.asarray(input_image)
         elif type(input_image) == str: