Update utils/model.py

utils/model.py

@@ -1,3 +1,5 @@
+import spaces
+
 import os
 import json
 
@@ -9,6 +11,7 @@ from model import FastSpeech2, ScheduledOptim
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+@spaces.GPU(duration=10)
 def get_model(args, configs, device, train=False):
     (preprocess_config, model_config, train_config) = configs
 
@@ -34,6 +37,7 @@ def get_model(args, configs, device, train=False):
     model.requires_grad_ = False
     return model
 
+@spaces.GPU(duration=10)
 def get_model_infer(ckpt_path, configs, device):
     (preprocess_config, model_config, train_config) = configs
 
@@ -49,7 +53,7 @@ def get_param_num(model):
     num_param = sum(param.numel() for param in model.parameters())
     return num_param
 
-
+@spaces.GPU(duration=10)
 def get_vocoder(config, device):
     name = config["vocoder"]["model"]
     speaker = config["vocoder"]["speaker"]
@@ -81,7 +85,7 @@ def get_vocoder(config, device):
 
     return vocoder
 
-
+@spaces.GPU(duration=10)
 def vocoder_infer(mels, vocoder, model_config, preprocess_config, lengths=None):
     name = model_config["vocoder"]["model"]
     with torch.no_grad():