Update usr/diff/shallow_diffusion_tts.py

usr/diff/shallow_diffusion_tts.py

@@ -16,6 +16,7 @@ from modules.diffsinger_midi.fs2 import FastSpeech2MIDI
 from utils.hparams import hparams
 
 
+
 def exists(x):
     return x is not None
 
@@ -157,7 +158,7 @@ class GaussianDiffusion(nn.Module):
 
     @torch.no_grad()
     def p_sample(self, x, t, cond, clip_denoised=True, repeat_noise=False):
-        b, *_, device = *x.shape, "cuda"
+        b, *_, device = *x.shape, x.device
         model_mean, _, model_log_variance = self.p_mean_variance(x=x, t=t, cond=cond, clip_denoised=clip_denoised)
         noise = noise_like(x.shape, device, repeat_noise)
         # no noise when t == 0
@@ -172,7 +173,10 @@ class GaussianDiffusion(nn.Module):
 
         def get_x_pred(x, noise_t, t):
             a_t = extract(self.alphas_cumprod, t, x.shape)
-            a_prev = extract(self.alphas_cumprod, torch.max(t-interval, torch.zeros_like(t)), x.shape)
+            if t[0] < interval:
+                a_prev = torch.ones_like(a_t)
+            else:
+                a_prev = extract(self.alphas_cumprod, torch.max(t-interval, torch.zeros_like(t)), x.shape)
             a_t_sq, a_prev_sq = a_t.sqrt(), a_prev.sqrt()
 
             x_delta = (a_prev - a_t) * ((1 / (a_t_sq * (a_t_sq + a_prev_sq))) * x - 1 / (a_t_sq * (((1 - a_prev) * a_t).sqrt() + ((1 - a_t) * a_prev).sqrt())) * noise_t)
@@ -228,7 +232,7 @@ class GaussianDiffusion(nn.Module):
 
     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, "cuda"
+        b, *_, device = *txt_tokens.shape, txt_tokens.device
         ret = self.fs2(txt_tokens, mel2ph, spk_embed, ref_mels, f0, uv, energy,
                        skip_decoder=(not infer), infer=infer, **kwargs)
         cond = ret['decoder_inp'].transpose(1, 2)
@@ -287,7 +291,7 @@ class GaussianDiffusion(nn.Module):
 class OfflineGaussianDiffusion(GaussianDiffusion):
     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, "cuda"
+        b, *_, device = *txt_tokens.shape, txt_tokens.device
 
         ret = self.fs2(txt_tokens, mel2ph, spk_embed, ref_mels, f0, uv, energy,
                        skip_decoder=True, infer=True, **kwargs)
@@ -316,4 +320,4 @@ class OfflineGaussianDiffusion(GaussianDiffusion):
                 x = self.p_sample(x, torch.full((b,), i, device=device, dtype=torch.long), cond)
             x = x[:, 0].transpose(1, 2)
             ret['mel_out'] = self.denorm_spec(x)
-        return ret
+        return ret