Update face_swapper.py

face_swapper.py

@@ -20,6 +20,7 @@ arcface_dst = np.array(
     dtype=np.float32)
 
 
+@spaces.GPU(enable_queue=True)
 def estimate_norm(lmk, image_size=112, mode='arcface'):
     assert lmk.shape == (5, 2)
     assert image_size % 112 == 0 or image_size % 128 == 0
@@ -37,14 +38,16 @@ def estimate_norm(lmk, image_size=112, mode='arcface'):
     return M
 
 
+@spaces.GPU(enable_queue=True)
 def norm_crop2(img, landmark, image_size=112, mode='arcface'):
     M = estimate_norm(landmark, image_size, mode)
     warped = cv2.warpAffine(img, M, (image_size, image_size), borderValue=0.0)
     return warped, M
 
 
-@spaces.GPU(enable_queue=True)
 class Inswapper():
+    
+    @spaces.GPU(enable_queue=True)
     def __init__(self, model_file=None, batch_size=32, providers=['CUDAExecutionProvider,CPUExecutionProvider']):
         self.model_file = model_file
         self.batch_size = batch_size
@@ -56,13 +59,15 @@ class Inswapper():
         self.session_options = onnxruntime.SessionOptions()
         self.session = onnxruntime.InferenceSession(self.model_file, sess_options=self.session_options, providers=providers)
 
+    @spaces.GPU(enable_queue=True)
     def forward(self, imgs, latents):
         preds = []
         for img, latent in zip(imgs, latents):
             img = img / 255
             pred = self.session.run(['output'], {'target': img, 'source': latent})[0]
             preds.append(pred)
-
+    
+    @spaces.GPU(enable_queue=True)
     def get(self, imgs, target_faces, source_faces):
         imgs = list(imgs)
 
@@ -79,7 +84,8 @@ class Inswapper():
             matrs[idx] = matrix
 
         return (preds, matrs)
-
+    
+    @spaces.GPU(enable_queue=True)
     def prepare_data(self, img, target_face, source_face):
         if isinstance(img, str):
             img = cv2.imread(img)
@@ -94,6 +100,7 @@ class Inswapper():
 
         return (matrix, blob, latent)
 
+    @spaces.GPU(enable_queue=True)
     def batch_forward(self, img_list, target_f_list, source_f_list):
         num_samples = len(img_list)
         num_batches = (num_samples + self.batch_size - 1) // self.batch_size
@@ -111,6 +118,7 @@ class Inswapper():
             yield batch_pred, batch_matr
 
 
+@spaces.GPU(enable_queue=True)
 def paste_to_whole(foreground, background, matrix, mask=None, crop_mask=(0,0,0,0), blur_amount=0.1, erode_amount = 0.15, blend_method='linear'):
     inv_matrix = cv2.invertAffineTransform(matrix)
     fg_shape = foreground.shape[:2]