Getting x(t) given x(t-1) using scheduler

[darshats]

Hi,
I'm exploring adding the Repaint algo inside the sdxl pipeline. For this, at any given timestep t, we need to noise the latent again (forward pass) and denoise it back, a set number of times (jump steps). We cant use the add_noise since it does the cumulative result assuming x(0) input. To just add noise for a specific timestep, I expect that the single step noising would work something like this:

def get_xt_from_xt_minus_1(xt_minus_1, timestep):
            alpha_t = self.scheduler.alphas[int(timestep.item())].to(xt_minus_1.device)
            sigma_t = torch.sqrt(self.scheduler.betas[int(timestep.item())]).to(xt_minus_1.device)

            mean = torch.sqrt(alpha_t) * xt_minus_1
            xt = mean + sigma_t * randn_tensor(xt_minus_1.shape, device=xt_minus_1.device, dtype=xt_minus_1.dtype)
            return xt

This isnt giving the right results. Am I missing the scale factor somewhere, or something else? Any insight greatly appreciated!
Thanks,
Darshat