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
@@ -171,6 +171,7 @@ class GaussianDiffusion(nn.Module):
         Use the PLMS method from [Pseudo Numerical Methods for Diffusion Models on Manifolds](https://arxiv.org/abs/2202.09778).
         """
 
+        @spaces.GPU(duration=8)
         def get_x_pred(x, noise_t, t):
             a_t = extract(self.alphas_cumprod, t, x.shape)
             if t[0] < interval: