Text encoder still not working correctly with LoRa Dreambooth training script

[JohnnyRacer]

Hello, I am getting much better results using the --train_text_encoder flag with the Dreambooth script. However, the actual outputed LoRa .pt files from models trained with train_text_encoder gives very bad results after using monkeypatch to generate images. I suspect that the text encoder's weights are still not saved properly. I tried to save the pipeline directly after each epoch from within the training script, but loading it using diffusers gives me strange errors about torch not being able to parse the linear layers. Does anyone have similar experiences with training the text encoder or have any idea why this is happening?

Images sampled from within the training loop (train_text_encoder enabled) :

Images sampled after model was monkeypatch with the trained LoRa weights (train_text_encoder enabled) :

The images doesn't seem to correlate with the samples generated while training and has very little cohesiveness with the training images used.

[cloneofsimo]

Did you also patch the text encoder as well?

[cloneofsimo]

Have a look at: https://github.com/cloneofsimo/lora/blob/master/scripts/run_with_text_lora_also.ipynb

[JohnnyRacer]

Did you also patch the text encoder as well?

@cloneofsimo

Seems like the repo was corrupted when I cloned it. I recloned and retrained the model and the text encoder files are there.

[JohnnyRacer]

@cloneofsimo
Got the text encoder patching to work , but it seems like it overfits very easily and gives quite bad results at times compared to training Dreambooth without LoRa. And even stranger is that there are a lot more artifacts after using the text encoder patch like watermarks from the training images that are showing up that had not appeared in the training like the examples I had show earlier like the examples below. Any clue why this is occuring?

Update :

Using tune_lora_scale(pipe.text_encoder, 0.7) to lower the strength of the text encoder weights seem to help a little, but also reduces image quality overall.

[cloneofsimo]

Nice to see that it works. I've found that when it overfits, tuning down both unet and text encoder helps, but with different scale. Also try learning it with prior preservation. Seems like you are dreamboothing with faces, and guide like https://github.com/nitrosocke/dreambooth-training-guide from @nitrosocke might help. (Also says to crop out any artifacts)

[cloneofsimo]

https://github.com/cloneofsimo/lora#what-happens-to-text-encoder-lora-and-unet-lora Have a look at this part as well

[brian6091]

I've posted some images of the effect of tuning both the unet and the text encoder (with prior preservation) in the discussions section:

#37

[JohnnyRacer]

I've posted some images of the effect of tuning both the unet and the text encoder (with prior preservation) in the discussions section:

#37

Thanks for the amazing guide @brian6091 ! Have you experimented with using a class dataset to act as a regularization similar to what was mentioned in the original Dreambooth paper? I tried to use it but it gives me strange errors about accelerate.launch() not being instantiated. I think it could give a noticeable increase image quality compared to using pipeline itself to automatically generate class images.

[JohnnyRacer]

Nice to see that it works. I've found that when it overfits, tuning down both unet and text encoder helps, but with different scale. Also try learning it with prior preservation. Seems like you are dreamboothing with faces, and guide like https://github.com/nitrosocke/dreambooth-training-guide from @nitrosocke might help. (Also says to crop out any artifacts)

From my experience so far it seems like training with LoRA requires much cleaner and specific datasets to work (i.e. aligned faces) , trying to train on more diverse datasets (i.e. pictures of a person in different poses) often leads to bad results.

[brian6091]

I've posted some images of the effect of tuning both the unet and the text encoder (with prior preservation) in the discussions section:
#37

Thanks for the amazing guide @brian6091 ! Have you experimented with using a class dataset to act as a regularization similar to what was mentioned in the original Dreambooth paper? I tried to use it but it gives me strange errors about accelerate.launch() not being instantiated. I think it could give a noticeable increase image quality compared to using pipeline itself to automatically generate class images.

I used a class dataset for the training I showed in the discussion post. It was the same for both Dreambooth and LoRA. I used the SD pipeline to generate the class images. I didn't get time to show the effects of both Dreambooth and LoRA on the class, but I may add some to a future discussion.

I was under the impression that the Dreambooth method was to use SD to generate class images. What did you do to get the error? If you post the error, maybe we could help debug it.

[JohnnyRacer]

@brian6091

I am getting this error if I pass in a directory with --class_data_dir :

ValueError: Please make sure to properly initialize your accelerator via 'accelerator = Accelerator()' before using any functionality from the 'accelerate'  library.

Without adding this argument it trains fine. I've tried using both accelerate launch as well as just using python to start the script but it gives the same errors. This is on the latest commit to the repo from today. Any ideas on why this is happening?

[brian6091]

@JohnnyRacer
Are you running in Colab? Which script or notebook are you using. I think I've seen this kind of error when there is a malformed input.

[JohnnyRacer]

@brian6091

No I'm running it locally. I am using the train_lora_dreambooth.py script from this repo with no modifications. The ./aradesh_train_col is where I have my training images and ./aradesh_reg_col is where the regularization images are located.

The command below works

 accelerate launch train_lora_dreambooth.py \
  --pretrained_model_name_or_path="../sd-v1-5-diffusers" \
  --instance_data_dir="./aradesh_train_col" \
  --output_dir="./aradesh-sd1.5-textenc" \
  --instance_prompt="sks aradesh" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=600 \
  --use_8bit_adam \
  --train_text_encoder

The command below doesn't work and gives the accelerate launch error that I mentioned earlier.

 accelerate launch train_lora_dreambooth.py \
  --pretrained_model_name_or_path="../sd-v1-5-diffusers" \
  --instance_data_dir="./aradesh_train_col" \
  --class_data_dir="./aradesh_reg_col" \
  --output_dir="./aradesh-sd1.5-textenc" \
  --instance_prompt="sks aradesh" \
  --resolution=512 \
  --train_batch_size=1 \
  --gradient_accumulation_steps=1 \
  --learning_rate=1e-4 \
  --lr_scheduler="constant" \
  --lr_warmup_steps=0 \
  --max_train_steps=600 \
  --use_8bit_adam \
  --train_text_encoder

[brian6091]

@JohnnyRacer

I think adding this flag would work:

--with_prior_preservation

I would also add this:

--prior_loss_weight=1.0

since the default weight is oddly huge

[JohnnyRacer]

@brian6091

Thanks for the help, I got it to work but also had to include the --class_prompt argument as well. Did you use the same prompt for the --class_prompt as the --instance_prompt when you were training your models?