Stable Diffusion XL with Dreambooth and Lora: how to structure dataset with a custom captioning locally

[SylwiaNowakowska]

Hi!

I am using https://github.com/huggingface/diffusers/blob/main/examples/dreambooth/train_dreambooth_lora_sdxl.py

I would like to fine-tune on a custom dataset containing images and corresponding prompts.
Here is how I launch the script:

accelerate launch train_dreambooth_lora_sdxl_0_25.py
--pretrained_model_name_or_path='/home/stable-diffusion-xl-base-1.0'
--pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix"
--instance_data_dir='/home/Dataset/data'
--image_column='file_name'
--caption_column="text"
--repeats=1
--instance_prompt="In the style of MAMM"
--validation_prompt="Yellow flower, in the style of MAMM."
--num_validation_images=4
--validation_epochs=1
--output_dir='/home/TEST_OUTPUT'
--seed=42
--resolution=1024
--crops_coords_top_left_h=0
--crops_coords_top_left_w=0
--center_crop
--train_text_encoder
--train_batch_size=1
--sample_batch_size=1
--max_train_steps=20
--checkpointing_steps=50
--checkpoints_total_limit=5
--gradient_accumulation_steps=5
--gradient_checkpointing
--learning_rate=2e-06
--text_encoder_lr=5e-6
--lr_scheduler="constant"
--snr_gamma=5.0
--lr_warmup_steps=0
--lr_num_cycles=1
--lr_power=1.0
--dataloader_num_workers=0
--optimizer="AdamW"
--adam_beta1=0.9
--adam_beta2=0.999
--adam_weight_decay=1e-04
--adam_weight_decay_text_encoder=1e-03
--adam_epsilon=1e-08
--max_grad_norm=1.0
--report_to=wandb
--mixed_precision="fp16"
--prior_generation_precision="fp16"
--local_rank=-1
--use_8bit_adam
--rank=4

The question is how should my folder structure look like?

Currently, the Dataset folder has a following structure

data/metadata.jsonl
data/images/01.png
data/images/02.png

With metadata.jsonl file structured like this:
{"file_name": "images/01.png", "text": "In the style of MAMM, a drawing of a pink flower"}
{"file_name": "images/02.png", "text": "In the style of MAMM, a drawing of a yellow flower"}

However I get then
PIL.UnidentifiedImageError: cannot identify image file 'home/Dataset/data/metadata.jsonl'

The same error happens with a following structure:

data/metadata.jsonl
data/01.png
data/02.png

With metadata.jsonl file structured like this:
{"file_name": "01.png", "text": "In the style of MAMM, a drawing of a pink flower"}
{"file_name": "02.png", "text": "In the style of MAMM, a drawing of a yellow flower"}

I would appreciate your help.

[a-r-r-o-w]

Hi @SylwiaNowakowska, I believe you're following the advanced sdxl lora blog for training. If not, the guide explains some great techniques for achieving good results and I'd highly recommend you checking it out.

Could you try running with --dataset_name=/home/Dataset/data instead of --instance_data_dir? The latter is for when you're not using multi-caption datasets. You'd also need to pip install datasets in order to use the dataset_name argument.

You might also need --caption_column="text" or whatever the attribute containing the prompts is called.

[a-r-r-o-w]

Could you try this?

accelerate launch train_dreambooth_lora_sdxl_0_25.py
--pretrained_model_name_or_path='/home/stable-diffusion-xl-base-1.0'
--pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix"
--dataset_name="/home/Dataset/data"
--caption_column="text"
--repeats=1
--instance_prompt="In the style of MAMM"
--validation_prompt="Yellow flower, in the style of MAMM."
--num_validation_images=4
--validation_epochs=1
--output_dir='/home/TEST_OUTPUT'
--seed=42
--resolution=1024
--crops_coords_top_left_h=0
--crops_coords_top_left_w=0
--center_crop
--train_text_encoder
--train_batch_size=1
--sample_batch_size=1
--max_train_steps=20
--checkpointing_steps=50
--checkpoints_total_limit=5
--gradient_accumulation_steps=5
--gradient_checkpointing
--learning_rate=2e-06
--text_encoder_lr=5e-6
--lr_scheduler="constant"
--snr_gamma=5.0
--lr_warmup_steps=0
--lr_num_cycles=1
--lr_power=1.0
--dataloader_num_workers=0
--optimizer="AdamW"
--adam_beta1=0.9
--adam_beta2=0.999
--adam_weight_decay=1e-04
--adam_weight_decay_text_encoder=1e-03
--adam_epsilon=1e-08
--max_grad_norm=1.0
--report_to=wandb
--mixed_precision="fp16"
--prior_generation_precision="fp16"
--local_rank=-1
--use_8bit_adam
--rank=4

Hopefully this should start training without issues.

[SylwiaNowakowska]

Dear @a-r-r-o-w,

Thank you so much for your fast reply!

I run the script as you suggested, the result is the same:
The image files are found:
Resolving data files: 100%|███████████████| 765/765 [00:00<00:00, 418937.53it/s]
And then gives the error:
ValueError: --caption_column value 'text' not found in dataset columns. Dataset columns are: image

FYI: I have diffusers 0.25.0

I would appreciate your feedback.

[SylwiaNowakowska]

I actually manged to make the custom dataset working by setting in the script:

!accelerate launch train_dreambooth_lora_sdxl_0_25.py
--pretrained_model_name_or_path='/media/brayz/brayz_storage/stable-diffusion-xl-base-1.0'
--pretrained_vae_model_name_or_path="madebyollin/sdxl-vae-fp16-fix"
--cache_dir='/media/brayz/brayz_storage/MAMMO_DIFFUSION_Project/cache_dir'
--dataset_name='/home/brayz/Mammo_Diffusion/TOY_DATASET_HEALTHY/data'
--image_column="image"
--caption_column="text"
--repeats=1
...

Setting the folder structure to:
data/metadata.jsonl
data/train/01.png
data/train/02.png

Structuring metadata.json file as follows:
{"file_name": "train/01.png", "text": "In the style of MAMM, a drawing of a pink flower"}
{"file_name": "train/02.png", "text": "In the style of MAMM, a drawing of a yellow flower"}

So, one folder had to be named "train".
The --image_column had to be "image" but in metadata.json it had to be called "file_name".

Thank you again for your great help!

[Valdiolus]

For now this doesn't work on sdxl - tried to do like @SylwiaNowakowska and the error is:

ImportError: You are trying to load your data using the datasets library. If you wish to train using custom captions please install the datasets library: pip install datasets. If you wish to load a local folder containing images only, specify --instance_data_dir instead.

pip install datasets did already.

[Valdiolus]

Found out that name dataset was in my PATH, removed and it works.

[sayakpaul]

Cc: @linoytsaban too in case you have anything more to add.

[linoytsaban]

sorry missed this, @a-r-r-o-w had it all covered from what I see! Thanks a lot! 😁
Attaching this doc about creating image dataset with the datasets library in case it adds some more clarity :)