Clean Up Comments in LCM(-LoRA) Distillation Scripts.

[dg845]

What does this PR do?

This PR cleans up and improves the comments in the LCM and LCM-LoRA distillation scripts. Includes a few small bug fixes.

Before submitting

• This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
• Did you read the contributor guideline?
• Did you read our philosophy doc (important for complex PRs)?
• Was this discussed/approved via a GitHub issue or the forum? Please add a link to it if that's the case.
• Did you make sure to update the documentation with your changes? Here are the
  documentation guidelines, and
  here are tips on formatting docstrings.
• Did you write any new necessary tests?

Who can review?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

@patrickvonplaten
@patil-suraj
@sayakpaul

[sayakpaul]

Hmm this seems like a big change to me. I wonder how much of an effect it has on the results.

[dg845]

Hi @sayakpaul, this change should be a no-op when the teacher LDM checkpoint has self.scheduler.config.prediction_type == "epsilon" (as is the case for the runwayml/stable-diffusion-v1-5 and stabilityai/stable-diffusion-xl-base-1.0 checkpoints suggested in the README examples). The current logic of directly using the residual output of the teacher unet is correct when the prediction_type is epsilon:

diffusers/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py

Line 1289 in 4836cfa

pred_noise = cond_teacher_output + w * (cond_teacher_output - uncond_teacher_output)

but in general we need to convert the residual unet output to a predicted epsilon for use in DDIMSolver.ddim_step, which is done by the new predicted_source_noise function, which is based on DDIMScheduler.step:

diffusers/src/diffusers/schedulers/scheduling_ddim.py

Lines 412 to 422 in 4836cfa

	# 3. compute predicted original sample from predicted noise also called
	# "predicted x_0" of formula (12) from https://arxiv.org/pdf/2010.02502.pdf
	if self.config.prediction_type == "epsilon":
	pred_original_sample = (sample - beta_prod_t ** (0.5) * model_output) / alpha_prod_t ** (0.5)
	pred_epsilon = model_output
	elif self.config.prediction_type == "sample":
	pred_original_sample = model_output
	pred_epsilon = (sample - alpha_prod_t ** (0.5) * pred_original_sample) / beta_prod_t ** (0.5)
	elif self.config.prediction_type == "v_prediction":
	pred_original_sample = (alpha_prod_t**0.5) * sample - (beta_prod_t**0.5) * model_output
	pred_epsilon = (alpha_prod_t**0.5) * model_output + (beta_prod_t**0.5) * sample

Note that pred_epsilon is exactly equal to model_output when prediction_type == "epsilon".

[sayakpaul]

Thanks for the solid explanation!

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[patil-suraj]

Really nice and explanatory comments, great job!
Looks good to me, just left some nits.