    def forward_diffusion(self, x0, diffusion_step):
        """
        x0: (B, 128, T)
        time_step: (B,)
        """
        time_step = diffusion_step.unsqueeze(-1).unsqueeze(-1)
        z = torch.randn(x0.shape, dtype=x0.dtype, device=x0.device, requires_grad=False)
        xt = (1 - time_step) * x0 + time_step * z
        return xt, z
x0: (B, 128, T)
time_step: (B,)
models/tts/naturalspeech2/diffusion_flow.py