Add support for Debiased Estimation loss to dreambooth (#2194)

README.md

@@ -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)
@@ -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)
 

kohya_gui/dreambooth_gui.py

@@ -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())
@@ -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())
@@ -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())
@@ -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,
@@ -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(

test/config/dreambooth-AdamW8bit.json

@@ -1,8 +1,14 @@
 {
+  "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,
@@ -10,46 +16,66 @@
   "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,
@@ -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,
@@ -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,

test/config/locon-AdamW8bit.json

@@ -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,
@@ -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,
@@ -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,