Fix dpmpp_2s_ancestral_cfg_pp: save temp[0] before second model call

[amatiramisu]

Was using dpmpp_2s_ancestral_cfg_pp for a bit and noticed that it was way too stable at high CFG values. Looked into the code and found this:

Lines 1325-1327 in comfy/k_diffusion/sampling.py:

x_2 = (sigma_fn(s) / sigma_fn(t)) * (x + (denoised - temp[0])) - (-h * r).expm1() * denoised
denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
x = (sigma_fn(t_next) / sigma_fn(t)) * (x + (denoised - temp[0])) - (-h).expm1() * denoised_2

Both lines use denoised from the first model call, so both should use the matching temp[0] from that call. But the second model call triggers post_cfg_function and overwrites temp[0], so line 1327 ends up using mismatched values.

Fix - save temp[0] before the second model call:

uncond_denoised = temp[0]
x_2 = (sigma_fn(s) / sigma_fn(t)) * (x + (denoised - uncond_denoised)) - (-h * r).expm1() * denoised
denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
x = (sigma_fn(t_next) / sigma_fn(t)) * (x + (denoised - uncond_denoised)) - (-h).expm1() * denoised_2

[amatiramisu]

Looking into this further, I am wondering if there was a specific reason to not follow the original CFG++ paper implementation on this. Left is the current approach (including my initial fix), and right would be the one following the CFG++ paper.
I did not upscale it or change anything else about it, so the differences are entirely current approach vs reference implementation. Difficult to say if one is better than the other, personally I say that reference has slightly better eyes and the ears come out a bit better. If more nature = better is pretty subjective though.
|
If aligning with the paper is desired, the change would be:

uncond_denoised = temp[0]
x_2 = (sigma_fn(s) / sigma_fn(t)) * x - (-h * r).expm1() * uncond_denoised
denoised_2 = model(x_2, sigma_fn(s) * s_in, **extra_args)
uncond_denoised_2 = temp[0]
x = denoised_2 - torch.exp(-h) * uncond_denoised_2 + (sigma_fn(t_next) / sigma_fn(t)) * x

Happy to add this to the PR or open a separate one if preferred.