Update usr/diff/shallow_diffusion_tts.py

usr/diff/shallow_diffusion_tts.py

@@ -15,7 +15,7 @@ from modules.fastspeech.fs2 import FastSpeech2
 from modules.diffsinger_midi.fs2 import FastSpeech2MIDI
 from utils.hparams import hparams
 
-
+import spaces
 
 def exists(x):
     return x is not None
@@ -226,7 +226,8 @@ class GaussianDiffusion(nn.Module):
             raise NotImplementedError()
 
         return loss
-
+        
+    @spaces.GPU(duration=180)
     def forward(self, txt_tokens, mel2ph=None, spk_embed=None,
                 ref_mels=None, f0=None, uv=None, energy=None, infer=False, **kwargs):
         b, *_, device = *txt_tokens.shape, txt_tokens.device