Update face_parsing/parse_mask.py

face_parsing/parse_mask.py

@@ -7,9 +7,11 @@ from PIL import Image
 from tqdm import tqdm
 import torch.nn.functional as F
 import torchvision.transforms as transforms
+import spaces
 
 from . model import BiSeNet
 
+@spaces.GPU(enable_queue=True)
 class SoftErosion(nn.Module):
     def __init__(self, kernel_size=15, threshold=0.6, iterations=1):
         super(SoftErosion, self).__init__()
@@ -59,7 +61,8 @@ transform = transforms.Compose([
 
 
 
-def init_parsing_model(model_path, device="cpu"):
+@spaces.GPU(enable_queue=True)
+def init_parsing_model(model_path, device="cuda"):
     net = BiSeNet(19)
     net.to(device)
     net.load_state_dict(torch.load(model_path))