Update usr/diff/shallow_diffusion_tts.py

usr/diff/shallow_diffusion_tts.py

@@ -157,7 +157,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, x.device
+        b, *_, device = *x.shape, "cuda"
         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,10 +172,7 @@ class GaussianDiffusion(nn.Module):
 
         def get_x_pred(x, noise_t, t):
             a_t = extract(self.alphas_cumprod, 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_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)
@@ -231,7 +228,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, txt_tokens.device
+        b, *_, device = *txt_tokens.shape, "cuda"
         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)
@@ -290,7 +287,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, txt_tokens.device
+        b, *_, device = *txt_tokens.shape, "cuda"
 
         ret = self.fs2(txt_tokens, mel2ph, spk_embed, ref_mels, f0, uv, energy,
                        skip_decoder=True, infer=True, **kwargs)
@@ -319,4 +316,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