Fix DDIMParallelScheduler batch path ignoring final_alpha_cumprod

[Osamaali313]

Problem

DDIMParallelScheduler implements ParaDiGMS, whose whole premise is that the parallel path produces the same result as the sequential DDIM path. For the terminal step (prev_timestep < 0), the scalar methods use self.final_alpha_cumprod:

# _get_variance (line 274) and step (line 448)
alpha_prod_t_prev = self.alphas_cumprod[prev_timestep] if prev_timestep >= 0 else self.final_alpha_cumprod

but the batched methods hardcode 1.0:

# _batch_get_variance (line 285) and batch_step_no_noise (line 580)
alpha_prod_t_prev[prev_t < 0] = torch.tensor(1.0)

self.final_alpha_cumprod is set at init to:

self.final_alpha_cumprod = torch.tensor(1.0) if set_alpha_to_one else self.alphas_cumprod[0]

So it equals 1.0 only when set_alpha_to_one=True. With set_alpha_to_one=False — the default in Stable Diffusion 1.x scheduler configs — it is alphas_cumprod[0] (~0.999). On the final denoising step the batched path therefore diverges from the scalar path, breaking the parallel==sequential guarantee.

Fix

Use self.final_alpha_cumprod in both batched methods, matching the scalar siblings:

alpha_prod_t_prev[prev_t < 0] = self.final_alpha_cumprod

final_alpha_cumprod is already moved to the correct device before batch_step_no_noise uses it (line 577), and the scalar siblings reference it the same way.

Reproduction

Porting the verbatim terminal-step arithmetic of both paths, with set_alpha_to_one=False (final_alpha_cumprod ≈ 0.999):

max|sequential - parallel(current, 1.0)|          = 6.76e-2
max|sequential - parallel(fixed, final_alpha)|    = 0.00e+00

The current code diverges by ~6.8e-2 on the final step; the fix makes the parallel path bit-exactly equal to the sequential path.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.