Skip to content

Commit

Permalink
Add support for Debiased Estimation loss to dreambooth
Browse files Browse the repository at this point in the history
  • Loading branch information
bmaltais committed Apr 2, 2024
1 parent b3a78fe commit 5664ae1
Show file tree
Hide file tree
Showing 4 changed files with 61 additions and 5 deletions.
5 changes: 5 additions & 0 deletions README.md
Original file line number Diff line number Diff line change
Expand Up @@ -36,6 +36,9 @@ The GUI allows you to set the training parameters and generate and run the requi
- [Troubleshooting](#troubleshooting)
- [Page File Limit](#page-file-limit)
- [No module called tkinter](#no-module-called-tkinter)
- [LORA Training on TESLA V100 - GPU Utilization Issue](#lora-training-on-tesla-v100---gpu-utilization-issue)
- [Issue Summary](#issue-summary)
- [Potential Solutions](#potential-solutions)
- [SDXL training](#sdxl-training)
- [Masked loss](#masked-loss)
- [Change History](#change-history)
Expand Down Expand Up @@ -434,6 +437,8 @@ ControlNet dataset is used to specify the mask. The mask images should be the RG
- The [English version of the dataset settings documentation](./docs/config_README-en.md) is added. PR [#1175](https://github.com/kohya-ss/sd-scripts/pull/1175) Thanks to darkstorm2150!

- Add GUI support for the new parameters listed above.
- Move accelerate launch parameters to new `Accelerate launch` accordion above `Model` accordion.
- Add support for `Debiased Estimation loss` to Dreambooth settings.

### 2024/03/21 (v23.0.15)

Expand Down
5 changes: 5 additions & 0 deletions kohya_gui/dreambooth_gui.py
Original file line number Diff line number Diff line change
Expand Up @@ -151,6 +151,7 @@ def save_configuration(
scale_v_pred_loss_like_noise_pred,
min_timestep,
max_timestep,
debiased_estimation_loss,
):
# Get list of function parameters and values
parameters = list(locals().items())
Expand Down Expand Up @@ -289,6 +290,7 @@ def open_configuration(
scale_v_pred_loss_like_noise_pred,
min_timestep,
max_timestep,
debiased_estimation_loss,
):
# Get list of function parameters and values
parameters = list(locals().items())
Expand Down Expand Up @@ -422,6 +424,7 @@ def train_model(
scale_v_pred_loss_like_noise_pred,
min_timestep,
max_timestep,
debiased_estimation_loss,
):
# Get list of function parameters and values
parameters = list(locals().items())
Expand Down Expand Up @@ -580,6 +583,7 @@ def train_model(
"clip_skip": clip_skip,
"color_aug": color_aug,
"dataset_config": dataset_config,
"debiased_estimation_loss": debiased_estimation_loss,
"enable_bucket": enable_bucket,
"epoch": epoch,
"flip_aug": flip_aug,
Expand Down Expand Up @@ -906,6 +910,7 @@ def dreambooth_tab(
advanced_training.scale_v_pred_loss_like_noise_pred,
advanced_training.min_timestep,
advanced_training.max_timestep,
advanced_training.debiased_estimation_loss,
]

configuration.button_open_config.click(
Expand Down
51 changes: 47 additions & 4 deletions test/config/dreambooth-AdamW8bit.json
Original file line number Diff line number Diff line change
@@ -1,55 +1,81 @@
{
"LoRA_type": "Kohya LoCon",
"LyCORIS_preset": "full",
"adaptive_noise_scale": 0,
"additional_parameters": "",
"block_alphas": "",
"block_dims": "",
"block_lr_zero_threshold": "",
"bucket_no_upscale": true,
"bucket_reso_steps": 64,
"bypass_mode": false,
"cache_latents": true,
"cache_latents_to_disk": false,
"caption_dropout_every_n_epochs": 0.0,
"caption_dropout_rate": 0.05,
"caption_extension": "",
"clip_skip": 2,
"color_aug": false,
"constrain": 0.0,
"conv_alpha": 8,
"conv_block_alphas": "",
"conv_block_dims": "",
"conv_dim": 16,
"dataset_config": "",
"debiased_estimation_loss": true,
"decompose_both": false,
"dim_from_weights": false,
"dora_wd": false,
"down_lr_weight": "",
"enable_bucket": true,
"epoch": 1,
"factor": -1,
"flip_aug": false,
"fp8_base": false,
"full_bf16": false,
"full_fp16": false,
"gpu_ids": "",
"gradient_accumulation_steps": 1,
"gradient_checkpointing": false,
"ip_noise_gamma": 0.1,
"ip_noise_gamma_random_strength": true,
"keep_tokens": "0",
"learning_rate": 5e-05,
"learning_rate_te": 1e-05,
"learning_rate_te1": 1e-05,
"learning_rate_te2": 1e-05,
"log_tracker_config": "",
"log_tracker_name": "",
"logging_dir": "./test/logs",
"lora_network_weights": "",
"lr_scheduler": "constant",
"lr_scheduler_args": "",
"lr_scheduler_num_cycles": "",
"lr_scheduler_power": "",
"lr_warmup": 0,
"main_process_port": 0,
"masked_loss": false,
"max_bucket_reso": 2048,
"max_data_loader_n_workers": "0",
"max_grad_norm": 1,
"max_resolution": "512,512",
"max_timestep": 1000,
"max_token_length": "75",
"max_train_epochs": "",
"max_train_steps": "",
"mem_eff_attn": false,
"mid_lr_weight": "",
"min_bucket_reso": 256,
"min_snr_gamma": 0,
"min_timestep": 0,
"mixed_precision": "bf16",
"model_list": "runwayml/stable-diffusion-v1-5",
"module_dropout": 0.1,
"multi_gpu": false,
"multires_noise_discount": 0,
"multires_noise_iterations": 0,
"no_token_padding": false,
"network_alpha": 8,
"network_dim": 16,
"network_dropout": 0.1,
"noise_offset": 0.05,
"noise_offset_random_strength": true,
"noise_offset_type": "Original",
"num_cpu_threads_per_process": 2,
"num_machines": 1,
Expand All @@ -62,7 +88,10 @@
"pretrained_model_name_or_path": "runwayml/stable-diffusion-v1-5",
"prior_loss_weight": 1.0,
"random_crop": false,
"rank_dropout": 0.1,
"rank_dropout_scale": false,
"reg_data_dir": "",
"rescaled": false,
"resume": "",
"sample_every_n_epochs": 0,
"sample_every_n_steps": 25,
Expand All @@ -75,13 +104,27 @@
"save_model_as": "safetensors",
"save_precision": "fp16",
"save_state": false,
"save_state_on_train_end": false,
"scale_v_pred_loss_like_noise_pred": false,
"scale_weight_norms": 1,
"sdxl": false,
"sdxl_cache_text_encoder_outputs": false,
"sdxl_no_half_vae": true,
"seed": "1234",
"shuffle_caption": false,
"stop_text_encoder_training": 0,
"text_encoder_lr": 0.0,
"train_batch_size": 4,
"train_data_dir": "./test/img",
"train_norm": false,
"train_on_input": false,
"training_comment": "",
"unet_lr": 0.0,
"unit": 1,
"up_lr_weight": "",
"use_cp": false,
"use_scalar": false,
"use_tucker": false,
"use_wandb": false,
"v2": false,
"v_parameterization": false,
Expand Down
5 changes: 4 additions & 1 deletion test/config/locon-AdamW8bit.json
Original file line number Diff line number Diff line change
Expand Up @@ -22,7 +22,7 @@
"conv_block_dims": "",
"conv_dim": 16,
"dataset_config": "",
"debiased_estimation_loss": false,
"debiased_estimation_loss": true,
"decompose_both": false,
"dim_from_weights": false,
"dora_wd": false,
Expand Down Expand Up @@ -50,6 +50,8 @@
"lr_scheduler_num_cycles": "",
"lr_scheduler_power": "",
"lr_warmup": 0,
"main_process_port": 0,
"masked_loss": false,
"max_bucket_reso": 2048,
"max_data_loader_n_workers": "0",
"max_grad_norm": 1,
Expand Down Expand Up @@ -102,6 +104,7 @@
"save_model_as": "safetensors",
"save_precision": "fp16",
"save_state": false,
"save_state_on_train_end": false,
"scale_v_pred_loss_like_noise_pred": false,
"scale_weight_norms": 1,
"sdxl": false,
Expand Down

0 comments on commit 5664ae1

Please sign in to comment.