Update face_parsing/swap.py

face_parsing/swap.py

@@ -4,6 +4,7 @@ import torch.nn.functional as F
 import torchvision.transforms as transforms
 import cv2
 import numpy as np
+import spaces
 
 from .model import BiSeNet
 
@@ -61,7 +62,9 @@ class SoftErosion(nn.Module):
 
 device = "cpu"
 
-def init_parser(pth_path, mode="cpu"):
+
+@spaces.GPU(enable_queue=True)
+def init_parser(pth_path, mode="cuda"):
     global device
     device = mode
     n_classes = 19
@@ -75,6 +78,7 @@ def init_parser(pth_path, mode="cpu"):
     return net
 
 
+@spaces.GPU(enable_queue=True)
 def image_to_parsing(img, net):
     img = cv2.resize(img, (512, 512))
     img = img[:,:,::-1]
@@ -99,6 +103,7 @@ def get_mask(parsing, classes):
     return res
 
 
+@spaces.GPU(enable_queue=True)
 def swap_regions(source, target, net, smooth_mask, includes=[1,2,3,4,5,10,11,12,13], blur=10):
     parsing = image_to_parsing(source, net)
 
@@ -125,6 +130,7 @@ def swap_regions(source, target, net, smooth_mask, includes=[1,2,3,4,5,10,11,12,
 
     return result
 
+@spaces.GPU(enable_queue=True)
 def mask_regions_to_list(values):
     out_ids = []
     for value in values: