Training own LCM-Lora

[AeroDEmi]

I have a dataset that I have trained using the text_to_image_lora_sdxl.py from the examples. Now, I want to train an LCM-Lora using the same dataset, but the results are pretty weird. I'm using this script: /examples/consistency_distillation/train_lcm_distill_lora_sdxl.py.

I have tried loading my Lora + the LCM Lora and setting the adapter weights, but the quality is not the best.

This is how I'm loading the weights:

vae = AutoencoderKL.from_pretrained(
    'madebyollin/sdxl-vae-fp16-fix', 
    torch_dtype=torch.float16, 
    use_safetensors=True
)   
pipe = StableDiffusionXLPipeline.from_pretrained(
       'stabilityai/stable-diffusion-xl-base-1.0',
       vae=vae,
       torch_dtype=torch.float16,
       use_safetensors=True,
       variant="fp16",
       )
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to('cuda')

pipe.load_lora_weights('./pokemons-lora-lcm-sdxl/checkpoint-100/', weight_name='pytorch_lora_weights.safetensors', adapter_name='icon')

images = pipe(
    prompt= "Bulbasur",
    negative_prompt= "unprofessional, words, text, photo, photograph, blurry, frame, outline frame, color background",
    height=512,
    width=512,
    num_images_per_prompt=4,
    guidance_scale=1.0,
    num_inference_steps=8,
    generator=torch.Generator('cuda').manual_seed(0)
)

[tolgacangoz]

Could you tell us what your exact training commands were for train_text_to_image_lora_sdxl.py and train_lcm_distill_lora_sdxl.py, also inference code for train_text_to_image_lora_sdxl.py-trained pipeline was?

[AeroDEmi]

This is the normal Lora:

!accelerate launch --config_file "1gpu.yaml" diffusers/examples/text_to_image/train_text_to_image_lora_sdxl.py \
  --pretrained_model_name_or_path="stabilityai/stable-diffusion-xl-base-1.0"  \
  --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \
  --train_data_dir="training/pokemon" \
  --output_dir="pokemons-lora-sdxl-text-encoder" \
  --mixed_precision="fp16" \
  --resolution=512 \
  --train_batch_size=4 \
  --gradient_accumulation_steps=1 \
  --gradient_checkpointing \
  --rank=4 \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=1000 \
  --checkpointing_steps=100 \
  --train_text_encoder \
  --seed="0"

Inference:

vae = AutoencoderKL.from_pretrained(
    'madebyollin/sdxl-vae-fp16-fix', 
    torch_dtype=torch.float16, 
    use_safetensors=True
)

pipe = StableDiffusionXLPipeline.from_pretrained(
       'stabilityai/stable-diffusion-xl-base-1.0',
       vae=vae,
       torch_dtype=torch.float16,
       use_safetensors=True,
       variant="fp16",
       )
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to('cuda')
pipe.load_lora_weights('./pokemons-lora-sdxl-text-encoder/checkpoint-1000/', weight_name='pytorch_lora_weights.safetensors', adapter_name='icon')

images = pipe(
    prompt= "a pokemon of a golden sea turtle with wings in a white background",
    negative_prompt= "unprofessional, words, text, photo, photograph, blurry, frame, outline frame, color background",
    height=512,
    width=512,
    num_images_per_prompt=4,
    guidance_scale=3,
    num_inference_steps=20,
    generator=torch.Generator('cuda').manual_seed(0)
)

Results:

This is from the LCM LORA:

accelerate launch --config_file "1gpu.yaml" diffusers/examples/consistency_distillation/train_lcm_distill_lora_sdxl.py \
  --pretrained_teacher_model="stabilityai/stable-diffusion-xl-base-1.0"  \
  --pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix" \
  --train_data_dir="training/pokemon" \
  --output_dir="pokemons-lora-lcm-sdxl" \
  --mixed_precision="fp16" \
  --resolution=512 \
  --train_batch_size=4 \
  --gradient_accumulation_steps=1 \
  --gradient_checkpointing \
  --lora_rank=4 \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=1000 \
  --checkpointing_steps=100 \
  --seed="0"

This is the inference:

vae = AutoencoderKL.from_pretrained(
    'madebyollin/sdxl-vae-fp16-fix', 
    torch_dtype=torch.float16, 
    use_safetensors=True
)   
unet = UNet2DConditionModel.from_pretrained("latent-consistency/lcm-sdxl", torch_dtype=torch.float16, variant="fp16")
pipe = StableDiffusionXLPipeline.from_pretrained(
       'stabilityai/stable-diffusion-xl-base-1.0',
       vae=vae,
       unet=unet,
       torch_dtype=torch.float16,
       use_safetensors=True,
       variant="fp16",
       )
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to('cuda')

images = pipe(
    prompt= "a pokemon of a golden sea turtle with wings in a white background",
    negative_prompt= "unprofessional, words, text, photo, photograph, blurry, frame, outline frame, color background",
    height=512,
    width=512,
    num_images_per_prompt=1,
    guidance_scale=8.0,
    num_inference_steps=8,
    generator=torch.Generator('cuda').manual_seed(0)
)

Results:

[tolgacangoz]

In the docs page, they say that they generated 768 x 768 images; so could you try this resolution?

[AeroDEmi]

I have trained again with --resolution=768. Seems like it helped but I'm wondering if I should use the other training script to compare them...
This is the inference:

vae = AutoencoderKL.from_pretrained(
    'madebyollin/sdxl-vae-fp16-fix', 
    torch_dtype=torch.float16, 
    use_safetensors=True
)   
unet = UNet2DConditionModel.from_pretrained("latent-consistency/lcm-sdxl", torch_dtype=torch.float16, variant="fp16")
pipe = StableDiffusionXLPipeline.from_pretrained(
       'stabilityai/stable-diffusion-xl-base-1.0',
       vae=vae,
       unet=unet,
       torch_dtype=torch.float16,
       use_safetensors=True,
       variant="fp16",
       )
pipe.scheduler = LCMScheduler.from_config(pipe.scheduler.config)
pipe.to('cuda')

pipe.load_lora_weights('./pokemons-lora-lcm-sdxl/checkpoint-1000/', weight_name='pytorch_lora_weights.safetensors', adapter_name='pkm')
images = pipe(
    prompt= "a pokemon of a golden sea turtle with wings in a white background",
    negative_prompt= "unprofessional, words, text, photo, photograph, blurry, frame, outline frame, color background",
    height=768,
    width=768,
    num_images_per_prompt=4,
    guidance_scale=3.0,
    num_inference_steps=8,
    generator=torch.Generator('cuda').manual_seed(0)
)

[tolgacangoz]

SD-XL mostly prefers 1024x1024. Could you try with this resolution as well for both?

[sayakpaul]

As pointed out by @StandardAI, resolution is really crucial for SDXL models to shine well. They fail terribly on smaller resolutions, actually. If using smaller resolutions is a must for you, then maybe try out other models that are based on SD and not SDXL?