[Training] Add datasets version of LCM LoRA SDXL

[sayakpaul]

What does this PR do?

Add a datasets compatible variant of https://github.com/huggingface/diffusers/blob/main/examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py. It also adapts a couple of best practices from peft:

• Instead of initializing two UNets (teacher and student), it leverages the disable_adapters() and enable_adapters() functions. In this case, the teacher is without any adapters and the student is with LoRA. We only update the LoRA params in the student.
• Makes use of peft utility modules positioning it as a utility library rather than a modelling library.

Running a couple experiments. Will report back the findings. But should be more or less ready to be reviewed.

Basic training command

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export OUTPUT_DIR="lora-lora-sdxl"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"
export VAE_PATH="madebyollin/sdxl-vae-fp16-fix"

accelerate launch train_lcm_distill_lora_sdxl.py \
  --pretrained_teacher_model=${MODEL_NAME}  \
  --pretrained_vae_model_name_or_path=${VAE_PATH} \
  --output_dir=${OUTPUT_DIR} \
  --mixed_precision="fp16" \
  --dataset_name=$DATASET_NAME \
  --resolution=1024 \
  --train_batch_size=4 \
  --gradient_accumulation_steps=1 \
  --gradient_checkpointing \
  --use_8bit_adam \
  --learning_rate=1e-4 --loss_type="huber" \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=1500 \
  --checkpointing_steps=50 \
  --validation_steps=25 \
  --seed="0" \
  --report_to="wandb" \
  --push_to_hub

TODOs

• Update the PR with results
• Docs
• Tests

[pcuenca]

Looks good so far, let me know if you need help testing / debugging

[sayakpaul]

It works, needs hyperparameter tuning. If you could take it to a test run, that would be great!

[patil-suraj]

Thanks a lot for working on this. Just left some comments. Most importantly:

Let's make sure we can load the saved lora with the pipeline for inference. I had some issues with it when I tried.

Everything else looks good !

[patil-suraj]

for later:

We could also think of making this an argument.

[sayakpaul]

Let's make sure we can load the saved lora with the pipeline for inference. I had some issues with it when I tried.

Could you post the snippet with which you tried and the error you got?

[patil-suraj]

Let's make sure we can load the saved lora with the pipeline for inference. I had some issues with it when I tried.

Could you post the snippet with which you tried and the error you got?

When I try to load the lora saved with script using

pipe.load_lora_weights(path_to_saved_lora, weight_name="pytorch_lora_weights.safetensors")

I get

---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
Cell In [13], line 1
----> 1 pipe.load_lora_weights("valhalla/lcm-sdxl-distill-lora-k5", weight_name="pytorch_lora_weights.safetensors")

File ~/diffusers/src/diffusers/loaders.py:3233, in StableDiffusionXLLoraLoaderMixin.load_lora_weights(self, pretrained_model_name_or_path_or_dict, adapter_name, **kwargs)
   3230 if not is_correct_format:
   3231     raise ValueError("Invalid LoRA checkpoint.")
-> 3233 self.load_lora_into_unet(
   3234     state_dict, network_alphas=network_alphas, unet=self.unet, adapter_name=adapter_name, _pipeline=self
   3235 )
   3236 text_encoder_state_dict = {k: v for k, v in state_dict.items() if "text_encoder." in k}
   3237 if len(text_encoder_state_dict) > 0:

File ~/diffusers/src/diffusers/loaders.py:1630, in LoraLoaderMixin.load_lora_into_unet(cls, state_dict, network_alphas, unet, low_cpu_mem_usage, adapter_name, _pipeline)
   1627     if "lora_B" in key:
   1628         rank[key] = val.shape[1]
-> 1630 lora_config_kwargs = get_peft_kwargs(rank, network_alphas, state_dict, is_unet=True)
   1631 lora_config = LoraConfig(**lora_config_kwargs)
   1633 # adapter_name

File ~/diffusers/src/diffusers/utils/peft_utils.py:122, in get_peft_kwargs(rank_dict, network_alpha_dict, peft_state_dict, is_unet)
    120 rank_pattern = {}
    121 alpha_pattern = {}
--> 122 r = lora_alpha = list(rank_dict.values())[0]
    124 if len(set(rank_dict.values())) > 1:
    125     # get the rank occuring the most number of times
    126     r = collections.Counter(rank_dict.values()).most_common()[0][0]

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint.

[sayakpaul]

@dg845 as well if you want to give this a look :-)

[dg845]

We could make this a little more general by using the timestep_scaling argument (and perhaps expose this as an argument):

Suggested change

	def scalings_for_boundary_conditions(timestep, sigma_data=0.5, timestep_scaling=10.0):
	c_skip = sigma_data**2 / ((timestep / 0.1) ** 2 + sigma_data**2)
	c_out = (timestep / 0.1) / ((timestep / 0.1) ** 2 + sigma_data**2) ** 0.5
	return c_skip, c_out
	def scalings_for_boundary_conditions(timestep, sigma_data=0.5, timestep_scaling=10.0):
	scaled_timestep = timestep_scaling * timestep
	c_skip = sigma_data**2 / (scaled_timestep**2 + sigma_data**2)
	c_out = scaled_timestep / (scaled_timestep**2 + sigma_data**2) ** 0.5
	return c_skip, c_out

(The current function is the same as in examples/consistency_distillation/train_lcm_distill_lora_sdxl_wds.py so if we make this change we should probably also propagate it to the WebDataset scripts.)

[sayakpaul]

Would prefer reflecting that after it's changed in WDS so that it's easier to track.

[patil-suraj]

+1, let's adapt this for both scripts.

[dg845]

Would it make sense to allow lora_alpha to be set independently of the r/lora_rank, to allow the LoRA layer scaling to be controlled?

[sayakpaul]

Maybe in a future PR. Since this PR is almost just a copy-paste of the WDS version.

[patil-suraj]

makes sense.

[sayakpaul]

@patil-suraj @dg845 ran a recent experiment with the following:

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export OUTPUT_DIR="lora-lcm-sdxl-new"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"
export VAE_PATH="madebyollin/sdxl-vae-fp16-fix"

CUDA_VISIBLE_DEVICES=1 accelerate launch train_lcm_distill_lora_sdxl.py \
  --pretrained_teacher_model=${MODEL_NAME}  \
  --pretrained_vae_model_name_or_path=${VAE_PATH} \
  --output_dir="lora-lcm-sdxl-new" \
  --mixed_precision="fp16" \
  --dataset_name=$DATASET_NAME \
  --resolution=1024 \
  --train_batch_size=24 \
  --gradient_accumulation_steps=1 \
  --gradient_checkpointing \
  --use_8bit_adam \
  --lora_rank=64 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=10000 \
  --checkpointing_steps=3000 \
  --validation_steps=50 \
  --seed="0" \
  --report_to="wandb" \
  --push_to_hub

WandB: https://wandb.ai/sayakpaul/text2image-fine-tune/runs/tv3zw00t
Weights: https://huggingface.co/sayakpaul/pokemons-lora-lcm-sdxl

Given it's only 10k steps and I haven't ablated the hyperparameters, I'd say it's still good enough. WDYT? I'd be keen on merging this soon as it reduces the memory requirements significantly by following good PEFT practices.

WDYT?

[patil-suraj]

Sounds good! Usually in my experiment 2-3k steps were enough, 10k results in overfitting with large BS.

[patil-suraj]

+1, let's adapt this for both scripts.

[patil-suraj]

makes sense.

[sayakpaul]

@patil-suraj I added the following for the reference training command when training on a very small dataset like Pokemons:

export MODEL_NAME="stabilityai/stable-diffusion-xl-base-1.0"
export DATASET_NAME="lambdalabs/pokemon-blip-captions"
export VAE_PATH="madebyollin/sdxl-vae-fp16-fix"

accelerate launch train_lcm_distill_lora_sdxl.py \
  --pretrained_teacher_model=${MODEL_NAME}  \
  --pretrained_vae_model_name_or_path=${VAE_PATH} \
  --output_dir="pokemons-lora-lcm-sdxl" \
  --mixed_precision="fp16" \
  --dataset_name=$DATASET_NAME \
  --resolution=1024 \
  --train_batch_size=24 \
  --gradient_accumulation_steps=1 \
  --gradient_checkpointing \
  --use_8bit_adam \
  --lora_rank=64 \
  --learning_rate=1e-4 \
  --report_to="wandb" \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=3000 \
  --checkpointing_steps=500 \
  --validation_steps=50 \
  --seed="0" \
  --report_to="wandb" \
  --push_to_hub

This has 3k steps instead of 10k. The training dynamics for this aren't clear but https://wandb.ai/sayakpaul/text2image-fine-tune/runs/tv3zw00t definitely shows progress IMO.

[sayakpaul]

Will merge after the CI is green.