    def forward(self, x, x_mask, cond, spk_query_emb, offset=1e-5):
        """
        x: (B, 128, T)
        x_mask: (B, T), mask is 0
        cond: (B, T, 512)
        spk_query_emb: (B, 32, 512)
        """
        diffusion_step = torch.rand(
            x.shape[0], dtype=x.dtype, device=x.device, requires_grad=False
        )
        diffusion_step = torch.clamp(diffusion_step, offset, 1.0 - offset)
        xt, z = self.forward_diffusion(x0=x, diffusion_step=diffusion_step)

        flow_pred = self.diff_estimator(
            xt, x_mask, cond, diffusion_step, spk_query_emb
        )  # noise - x0_pred, noise_pred - x0
        noise = z
        x0_pred = noise - flow_pred
        noise_pred = x + flow_pred
        diff_out = {
            "x0_pred": x0_pred,
            "noise_pred": noise_pred,
            "noise": noise,
            "flow_pred": flow_pred,
        }
        return diff_out
x: (B, 128, T)
x_mask: (B, T), mask is 0
cond: (B, T, 512)
spk_query_emb: (B, 32, 512)
models/tts/naturalspeech2/diffusion_flow.py