Add support for masked_loss to the GUI

kohya_gui/class_advanced_training.py

@@ -222,8 +222,13 @@ class AdvancedTraining:
                 choices=["none", "sdpa", "xformers"],
                 value="xformers",
             )
-            self.color_aug = gr.Checkbox(label="Color augmentation", value=False)
-            self.flip_aug = gr.Checkbox(label="Flip augmentation", value=False)
+            self.color_aug = gr.Checkbox(label="Color augmentation", value=False, info="Enable weak color augmentation")
+            self.flip_aug = gr.Checkbox(label="Flip augmentation", value=False, info="Enable horizontal flip augmentation")
+            self.masked_loss = gr.Checkbox(
+                label="Masked loss",
+                value=False,
+                info="Apply mask for calculating loss. conditioning_data_dir is required for dataset"
+            )
             self.min_snr_gamma = gr.Slider(
                 label="Min SNR gamma",
                 value=0,

kohya_gui/common_gui.py

@@ -1237,6 +1237,10 @@ def run_cmd_advanced_training(**kwargs):
         gpu_ids = kwargs.get("gpu_ids")
         if not gpu_ids == "":
             run_cmd += f' --gpu_ids="{gpu_ids}"'
+            
+    if "masked_loss" in kwargs:
+        if kwargs.get("masked_loss"): # Test if the value is true as it could be false
+            run_cmd += " --masked_loss"
 
     if "max_data_loader_n_workers" in kwargs:
         max_data_loader_n_workers = kwargs.get("max_data_loader_n_workers")

kohya_gui/dreambooth_gui.py

@@ -94,6 +94,7 @@ def save_configuration(
     prior_loss_weight,
     color_aug,
     flip_aug,
+    masked_loss,
     clip_skip,
     vae,
     num_processes,
@@ -230,6 +231,7 @@ def open_configuration(
     prior_loss_weight,
     color_aug,
     flip_aug,
+    masked_loss,
     clip_skip,
     vae,
     num_processes,
@@ -361,6 +363,7 @@ def train_model(
     prior_loss_weight,
     color_aug,
     flip_aug,
+    masked_loss,
     clip_skip,
     vae,
     num_processes,
@@ -575,6 +578,7 @@ def train_model(
         "enable_bucket": enable_bucket,
         "epoch": epoch,
         "flip_aug": flip_aug,
+        "masked_loss": masked_loss,
         "full_bf16": full_bf16,
         "full_fp16": full_fp16,
         "gradient_accumulation_steps": gradient_accumulation_steps,
@@ -835,6 +839,7 @@ def dreambooth_tab(
             advanced_training.prior_loss_weight,
             advanced_training.color_aug,
             advanced_training.flip_aug,
+            advanced_training.masked_loss,
             advanced_training.clip_skip,
             advanced_training.vae,
             advanced_training.num_processes,

kohya_gui/finetune_gui.py

@@ -69,6 +69,7 @@ def save_configuration(
     max_bucket_reso,
     batch_size,
     flip_aug,
+    masked_loss,
     caption_metadata_filename,
     latent_metadata_filename,
     full_path,
@@ -213,6 +214,7 @@ def open_configuration(
     max_bucket_reso,
     batch_size,
     flip_aug,
+    masked_loss,
     caption_metadata_filename,
     latent_metadata_filename,
     full_path,
@@ -364,6 +366,7 @@ def train_model(
     max_bucket_reso,
     batch_size,
     flip_aug,
+    masked_loss,
     caption_metadata_filename,
     latent_metadata_filename,
     full_path,
@@ -609,6 +612,7 @@ def train_model(
         "dataset_repeats": dataset_repeats,
         "enable_bucket": True,
         "flip_aug": flip_aug,
+        "masked_loss": masked_loss,
         "full_bf16": full_bf16,
         "full_fp16": full_fp16,
         "gradient_accumulation_steps": gradient_accumulation_steps,
@@ -897,6 +901,7 @@ def finetune_tab(headless=False, config: dict = {}):
             max_bucket_reso,
             batch_size,
             advanced_training.flip_aug,
+            advanced_training.masked_loss,
             caption_metadata_filename,
             latent_metadata_filename,
             full_path,

kohya_gui/lora_gui.py

@@ -130,6 +130,7 @@ def save_configuration(
     dim_from_weights,
     color_aug,
     flip_aug,
+    masked_loss,
     clip_skip,
     num_processes,
     num_machines,
@@ -315,6 +316,7 @@ def open_configuration(
     dim_from_weights,
     color_aug,
     flip_aug,
+    masked_loss,
     clip_skip,
     num_processes,
     num_machines,
@@ -528,6 +530,7 @@ def train_model(
     dim_from_weights,
     color_aug,
     flip_aug,
+    masked_loss,
     clip_skip,
     num_processes,
     num_machines,
@@ -930,6 +933,7 @@ def train_model(
         "enable_bucket": enable_bucket,
         "epoch": epoch,
         "flip_aug": flip_aug,
+        "masked_loss": masked_loss,
         "fp8_base": fp8_base,
         "full_bf16": full_bf16,
         "full_fp16": full_fp16,
@@ -1963,6 +1967,7 @@ def lora_tab(
             dim_from_weights,
             advanced_training.color_aug,
             advanced_training.flip_aug,
+            advanced_training.masked_loss,
             advanced_training.clip_skip,
             advanced_training.num_processes,
             advanced_training.num_machines,

test/config/dataset-masked_loss.toml

@@ -0,0 +1,15 @@
+[[datasets]]
+resolution = 512
+batch_size = 4
+keep_tokens = 1
+enable_bucket = true
+min_bucket_reso = 64
+max_bucket_reso = 1024
+bucket_reso_steps = 32
+bucket_no_upscale = true
+
+  [[datasets.subsets]]
+  image_dir = '.\test\img\10_darius kawasaki person'
+  num_repeats = 10
+  caption_extension = '.txt'
+  conditioning_data_dir = '.\test\masked_loss'

test/config/dreambooth-AdamW8bit-masked_loss-toml.json

@@ -0,0 +1,100 @@
+{
+  "adaptive_noise_scale": 0,
+  "additional_parameters": "",
+  "bucket_no_upscale": true,
+  "bucket_reso_steps": 64,
+  "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,
+  "dataset_config": "D:/kohya_ss/test/config/dataset-masked_loss.toml",
+  "enable_bucket": true,
+  "epoch": 1,
+  "flip_aug": false,
+  "full_bf16": false,
+  "full_fp16": false,
+  "gpu_ids": "",
+  "gradient_accumulation_steps": 1,
+  "gradient_checkpointing": false,
+  "ip_noise_gamma": 0,
+  "ip_noise_gamma_random_strength": false,
+  "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",
+  "lr_scheduler": "constant",
+  "lr_scheduler_args": "",
+  "lr_scheduler_num_cycles": "",
+  "lr_scheduler_power": "",
+  "lr_warmup": 0,
+  "masked_loss": true,
+  "max_bucket_reso": 2048,
+  "max_data_loader_n_workers": "0",
+  "max_resolution": "512,512",
+  "max_timestep": 1000,
+  "max_token_length": "75",
+  "max_train_epochs": "",
+  "max_train_steps": "",
+  "mem_eff_attn": false,
+  "min_bucket_reso": 256,
+  "min_snr_gamma": 0,
+  "min_timestep": 0,
+  "mixed_precision": "bf16",
+  "model_list": "runwayml/stable-diffusion-v1-5",
+  "multi_gpu": false,
+  "multires_noise_discount": 0,
+  "multires_noise_iterations": 0,
+  "no_token_padding": false,
+  "noise_offset": 0.05,
+  "noise_offset_random_strength": false,
+  "noise_offset_type": "Original",
+  "num_cpu_threads_per_process": 2,
+  "num_machines": 1,
+  "num_processes": 1,
+  "optimizer": "AdamW8bit",
+  "optimizer_args": "",
+  "output_dir": "./test/output",
+  "output_name": "db-AdamW8bit-masked_loss-toml",
+  "persistent_data_loader_workers": false,
+  "pretrained_model_name_or_path": "runwayml/stable-diffusion-v1-5",
+  "prior_loss_weight": 1.0,
+  "random_crop": false,
+  "reg_data_dir": "",
+  "resume": "",
+  "sample_every_n_epochs": 0,
+  "sample_every_n_steps": 25,
+  "sample_prompts": "a painting of a gas mask , by darius kawasaki",
+  "sample_sampler": "euler_a",
+  "save_every_n_epochs": 1,
+  "save_every_n_steps": 0,
+  "save_last_n_steps": 0,
+  "save_last_n_steps_state": 0,
+  "save_model_as": "safetensors",
+  "save_precision": "fp16",
+  "save_state": false,
+  "save_state_on_train_end": false,
+  "scale_v_pred_loss_like_noise_pred": false,
+  "sdxl": false,
+  "seed": "1234",
+  "shuffle_caption": false,
+  "stop_text_encoder_training": 0,
+  "train_batch_size": 4,
+  "train_data_dir": "",
+  "use_wandb": false,
+  "v2": false,
+  "v_parameterization": false,
+  "v_pred_like_loss": 0,
+  "vae": "",
+  "vae_batch_size": 0,
+  "wandb_api_key": "",
+  "wandb_run_name": "",
+  "weighted_captions": false,
+  "xformers": "xformers"
+}

test/config/locon-AdamW8bit-masked_loss-toml.json

@@ -0,0 +1,137 @@
+{
+  "LoRA_type": "Standard",
+  "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": 1,
+  "conv_block_alphas": "",
+  "conv_block_dims": "",
+  "conv_dim": 1,
+  "dataset_config": "D:/kohya_ss/test/config/dataset-masked_loss.toml",
+  "debiased_estimation_loss": false,
+  "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,
+  "ip_noise_gamma_random_strength": false,
+  "keep_tokens": "0",
+  "learning_rate": 0.0005,
+  "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,
+  "masked_loss": true,
+  "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,
+  "multi_gpu": false,
+  "multires_noise_discount": 0,
+  "multires_noise_iterations": 0,
+  "network_alpha": 1,
+  "network_dim": 8,
+  "network_dropout": 0,
+  "noise_offset": 0.05,
+  "noise_offset_random_strength": false,
+  "noise_offset_type": "Original",
+  "num_cpu_threads_per_process": 2,
+  "num_machines": 1,
+  "num_processes": 1,
+  "optimizer": "AdamW8bit",
+  "optimizer_args": "",
+  "output_dir": "./test/output",
+  "output_name": "locon-AdamW8bit-masked_loss-toml",
+  "persistent_data_loader_workers": false,
+  "pretrained_model_name_or_path": "runwayml/stable-diffusion-v1-5",
+  "prior_loss_weight": 1.0,
+  "random_crop": false,
+  "rank_dropout": 0,
+  "rank_dropout_scale": false,
+  "reg_data_dir": "",
+  "rescaled": false,
+  "resume": "",
+  "sample_every_n_epochs": 0,
+  "sample_every_n_steps": 25,
+  "sample_prompts": "a painting of a gas mask , by darius kawasaki",
+  "sample_sampler": "euler_a",
+  "save_every_n_epochs": 1,
+  "save_every_n_steps": 0,
+  "save_last_n_steps": 0,
+  "save_last_n_steps_state": 0,
+  "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": 0,
+  "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": "",
+  "train_norm": false,
+  "train_on_input": true,
+  "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,
+  "v_pred_like_loss": 0,
+  "vae": "",
+  "vae_batch_size": 0,
+  "wandb_api_key": "",
+  "wandb_run_name": "",
+  "weighted_captions": false,
+  "xformers": "xformers"
+}