Adding GUI support for new features

2023-02-22 20:32:57 -05:00 · 2023-02-22 20:32:57 -05:00 · bf0344ba9e
commit bf0344ba9e
parent 2a5fb346d5
8 changed files with 60 additions and 22 deletions
--- a/README.md
+++ b/README.md
@ -164,8 +164,10 @@ This will store your a backup file with your current locally installed pip packa
 ## Change History

 * 2023/02/22 (v20.8.0):
+    - Add gui support for optimizers: `AdamW, AdamW8bit, Lion, SGDNesterov, SGDNesterov8bit, DAdaptation, AdaFactor`
+    - Add gui support for `--noise_offset`
    - Refactor optmizer options. Thanks to mgz-dev!
-        - Add `--optimizer_type` option for each training script. Please see help. Japanese documentation is [here](https://github.com/kohya-ss/sd-scripts/blob/main/train_network_README-ja.md#%E3%82%AA%E3%83%97%E3%83%86%E3%82%A3%E3%83%9E%E3%82%A4%E3%82%B6%E3%81%AE%E6%8C%87%E5%AE%9A%E3%81%AB%E3%81%A4%E3%81%84%E3%81%A6).
+        - Add `--optimizer_type` option for each training script. Please see help. Japanese documentation is [here](https://github-com.translate.goog/kohya-ss/sd-scripts/blob/main/train_network_README-ja.md?_x_tr_sl=fr&_x_tr_tl=en&_x_tr_hl=en-US&_x_tr_pto=wapp#%E3%82%AA%E3%83%97%E3%83%86%E3%82%A3%E3%83%9E%E3%82%A4%E3%82%B6%E3%81%AE%E6%8C%87%E5%AE%9A%E3%81%AB%E3%81%A4%E3%81%84%E3%81%A6).
        - `--use_8bit_adam` and `--use_lion_optimizer` options also work and will override the options above for backward compatibility.
    - Add SGDNesterov and its 8bit.
    - Add [D-Adaptation](https://github.com/facebookresearch/dadaptation) optimizer. Thanks to BootsofLagrangian and all! 
--- a/dreambooth_gui.py
+++ b/dreambooth_gui.py
@ -90,6 +90,7 @@ def save_configuration(
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
    optimizer,
+    optimizer_args,noise_offset,
 ):
    # Get list of function parameters and values
    parameters = list(locals().items())
@ -188,6 +189,7 @@ def open_configuration(
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
    optimizer,
+    optimizer_args,noise_offset,
 ):
    # Get list of function parameters and values
    parameters = list(locals().items())
@ -263,6 +265,7 @@ def train_model(
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
    optimizer,
+    optimizer_args,noise_offset,
 ):
    if pretrained_model_name_or_path == '':
        msgbox('Source model information is missing')
@ -407,7 +410,8 @@ def train_model(
        seed=seed,
        caption_extension=caption_extension,
        cache_latents=cache_latents,
-        optimizer=optimizer
+        optimizer=optimizer,
+        optimizer_args=optimizer_args,
    )

    run_cmd += run_cmd_advanced_training(
@ -432,6 +436,7 @@ def train_model(
        bucket_reso_steps=bucket_reso_steps,
        caption_dropout_every_n_epochs=caption_dropout_every_n_epochs,
        caption_dropout_rate=caption_dropout_rate,
+        noise_offset=noise_offset,
    )

    print(run_cmd)
@ -552,7 +557,7 @@ def dreambooth_tab(
            seed,
            caption_extension,
            cache_latents,
-            optimizer,
+            optimizer,optimizer_args,
        ) = gradio_training(
            learning_rate_value='1e-5',
            lr_scheduler_value='cosine',
@ -610,7 +615,7 @@ def dreambooth_tab(
                bucket_no_upscale,
                random_crop,
                bucket_reso_steps,
-                caption_dropout_every_n_epochs, caption_dropout_rate,
+                caption_dropout_every_n_epochs, caption_dropout_rate,noise_offset,
            ) = gradio_advanced_training()
            color_aug.change(
                color_aug_changed,
@ -680,7 +685,7 @@ def dreambooth_tab(
        random_crop,
        bucket_reso_steps,
        caption_dropout_every_n_epochs, caption_dropout_rate,
-        optimizer,
+        optimizer,optimizer_args,noise_offset,
    ]

    button_open_config.click(
--- a/finetune_gui.py
+++ b/finetune_gui.py
@ -85,7 +85,7 @@ def save_configuration(
    random_crop,
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
-    optimizer,
+    optimizer,optimizer_args,noise_offset,
 ):
    # Get list of function parameters and values
    parameters = list(locals().items())
@ -189,7 +189,7 @@ def open_config_file(
    random_crop,
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
-    optimizer,
+    optimizer,optimizer_args,noise_offset,
 ):
    # Get list of function parameters and values
    parameters = list(locals().items())
@ -271,7 +271,7 @@ def train_model(
    random_crop,
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
-    optimizer,
+    optimizer,optimizer_args,noise_offset,
 ):
    # create caption json file
    if generate_caption_database:
@ -397,6 +397,7 @@ def train_model(
        caption_extension=caption_extension,
        cache_latents=cache_latents,
        optimizer=optimizer,
+        optimizer_args=optimizer_args,
    )

    run_cmd += run_cmd_advanced_training(
@ -421,6 +422,7 @@ def train_model(
        bucket_reso_steps=bucket_reso_steps,
        caption_dropout_every_n_epochs=caption_dropout_every_n_epochs,
        caption_dropout_rate=caption_dropout_rate,
+        noise_offset=noise_offset,
    )

    print(run_cmd)
@ -575,7 +577,7 @@ def finetune_tab():
            seed,
            caption_extension,
            cache_latents,
-            optimizer,
+            optimizer,optimizer_args,
        ) = gradio_training(learning_rate_value='1e-5')
        with gr.Row():
            dataset_repeats = gr.Textbox(label='Dataset repeats', value=40)
@ -607,7 +609,7 @@ def finetune_tab():
                bucket_no_upscale,
                random_crop,
                bucket_reso_steps,
-                caption_dropout_every_n_epochs, caption_dropout_rate,
+                caption_dropout_every_n_epochs, caption_dropout_rate,noise_offset,
            ) = gradio_advanced_training()
            color_aug.change(
                color_aug_changed,
@ -673,7 +675,7 @@ def finetune_tab():
        random_crop,
        bucket_reso_steps,
        caption_dropout_every_n_epochs, caption_dropout_rate,
-        optimizer,
+        optimizer,optimizer_args,noise_offset,
    ]

    button_run.click(train_model, inputs=settings_list)
--- a/library/common_gui.py
+++ b/library/common_gui.py
@ -469,11 +469,20 @@ def gradio_training(
            label='Optimizer',
            choices=[
                'AdamW',
+                'AdamW8bit',
+                'Adafactor',
+                'DAdaptation',
                'Lion',
+                'SGDNesterov',
+                'SGDNesterov8bit'
            ],
            value="AdamW",
            interactive=True,
        )
+    with gr.Row():
+        optimizer_args = gr.Textbox(
+            label='Optimizer extra arguments', placeholder='(Optional) eg: relative_step=True scale_parameter=True warmup_init=True'
+        )
    return (
        learning_rate,
        lr_scheduler,
@ -488,6 +497,7 @@ def gradio_training(
        caption_extension,
        cache_latents,
        optimizer,
+        optimizer_args,
    )


@ -522,7 +532,9 @@ def run_cmd_training(**kwargs):
        if kwargs.get('caption_extension')
        else '',
        ' --cache_latents' if kwargs.get('cache_latents') else '',
-        ' --use_lion_optimizer' if kwargs.get('optimizer') == 'Lion' else '',
+        # ' --use_lion_optimizer' if kwargs.get('optimizer') == 'Lion' else '',
+        f' --optimizer_type="{kwargs.get("optimizer", "AdamW")}"',
+        f' --optimizer_args {kwargs.get("optimizer_args", "")}' if not kwargs.get('optimizer_args') == '' else '',
    ]
    run_cmd = ''.join(options)
    return run_cmd
@ -597,6 +609,10 @@ def gradio_advanced_training():
        random_crop = gr.Checkbox(
            label='Random crop instead of center crop', value=False
        )
+        noise_offset = gr.Textbox(
+            label='Noise offset (0 - 1)', placeholder='(Oprional) eg: 0.1'
+        )
+        
    with gr.Row():
        caption_dropout_every_n_epochs = gr.Number(
            label="Dropout caption every n epochs",
@ -644,7 +660,7 @@ def gradio_advanced_training():
        bucket_no_upscale,
        random_crop,
        bucket_reso_steps,
-        caption_dropout_every_n_epochs, caption_dropout_rate,
+        caption_dropout_every_n_epochs, caption_dropout_rate,noise_offset,
    )


@ -695,6 +711,9 @@ def run_cmd_advanced_training(**kwargs):
        else '',
        ' --bucket_no_upscale' if kwargs.get('bucket_no_upscale') else '',
        ' --random_crop' if kwargs.get('random_crop') else '',
+        f' --noise_offset={float(kwargs.get("noise_offset", 0))}'
+        if not kwargs.get('noise_offset', '') == ''
+        else '',
    ]
    run_cmd = ''.join(options)
    return run_cmd
--- a/library/train_util.py
+++ b/library/train_util.py
@ -1633,7 +1633,7 @@ def get_optimizer(args, trainable_params):
    if optimizer_kwargs["relative_step"]:
      print(f"relative_step is true / relative_stepがtrueです")
      if lr != 0.0:
-        print(f"learning rate is used as initial_lr / 指定したlearning rateはinitial_lrとして使用されます")
+        print(f"learning rate is used as initial_lr / 指定したlearning rate はinitial_lrとして使用されます: {lr}")
      args.learning_rate = None

      # trainable_paramsがgroupだった時の処理：lrを削除する
--- a/lora_gui.py
+++ b/lora_gui.py
@ -101,6 +101,7 @@ def save_configuration(
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
    optimizer,
+    optimizer_args,noise_offset,
 ):
    # Get list of function parameters and values
    parameters = list(locals().items())
@ -206,6 +207,7 @@ def open_configuration(
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
    optimizer,
+    optimizer_args,noise_offset,
 ):
    # Get list of function parameters and values
    parameters = list(locals().items())
@ -288,6 +290,7 @@ def train_model(
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
    optimizer,
+    optimizer_args,noise_offset,
 ):  
    if pretrained_model_name_or_path == '':
        msgbox('Source model information is missing')
@ -468,6 +471,7 @@ def train_model(
        caption_extension=caption_extension,
        cache_latents=cache_latents,
        optimizer=optimizer,
+        optimizer_args=optimizer_args,
    )

    run_cmd += run_cmd_advanced_training(
@ -492,6 +496,7 @@ def train_model(
        bucket_reso_steps=bucket_reso_steps,
        caption_dropout_every_n_epochs=caption_dropout_every_n_epochs,
        caption_dropout_rate=caption_dropout_rate,
+        noise_offset=noise_offset,
    )

    print(run_cmd)
@ -621,6 +626,7 @@ def lora_tab(
            caption_extension,
            cache_latents,
            optimizer,
+            optimizer_args,
        ) = gradio_training(
            learning_rate_value='0.0001',
            lr_scheduler_value='cosine',
@ -708,7 +714,7 @@ def lora_tab(
                bucket_no_upscale,
                random_crop,
                bucket_reso_steps,
-                caption_dropout_every_n_epochs, caption_dropout_rate,
+                caption_dropout_every_n_epochs, caption_dropout_rate,noise_offset,
            ) = gradio_advanced_training()
            color_aug.change(
                color_aug_changed,
@ -791,6 +797,7 @@ def lora_tab(
        bucket_reso_steps,
        caption_dropout_every_n_epochs, caption_dropout_rate,
        optimizer,
+        optimizer_args,noise_offset,
    ]

    button_open_config.click(
--- a/textual_inversion_gui.py
+++ b/textual_inversion_gui.py
@ -95,7 +95,7 @@ def save_configuration(
    random_crop,
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
-    optimizer,
+    optimizer,optimizer_args,noise_offset,
 ):
    # Get list of function parameters and values
    parameters = list(locals().items())
@ -199,7 +199,7 @@ def open_configuration(
    random_crop,
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
-    optimizer,
+    optimizer,optimizer_args,noise_offset,
 ):
    # Get list of function parameters and values
    parameters = list(locals().items())
@ -280,7 +280,7 @@ def train_model(
    random_crop,
    bucket_reso_steps,
    caption_dropout_every_n_epochs, caption_dropout_rate,
-    optimizer,
+    optimizer,optimizer_args,noise_offset,
 ):
    if pretrained_model_name_or_path == '':
        msgbox('Source model information is missing')
@ -441,6 +441,7 @@ def train_model(
        caption_extension=caption_extension,
        cache_latents=cache_latents,
        optimizer=optimizer,
+        optimizer_args=optimizer_args,
    )

    run_cmd += run_cmd_advanced_training(
@ -465,6 +466,7 @@ def train_model(
        bucket_reso_steps=bucket_reso_steps,
        caption_dropout_every_n_epochs=caption_dropout_every_n_epochs,
        caption_dropout_rate=caption_dropout_rate,
+        noise_offset=noise_offset,
    )
    run_cmd += f' --token_string="{token_string}"'
    run_cmd += f' --init_word="{init_word}"'
@ -630,7 +632,7 @@ def ti_tab(
            seed,
            caption_extension,
            cache_latents,
-            optimizer,
+            optimizer,optimizer_args,
        ) = gradio_training(
            learning_rate_value='1e-5',
            lr_scheduler_value='cosine',
@ -688,7 +690,7 @@ def ti_tab(
                bucket_no_upscale,
                random_crop,
                bucket_reso_steps,
-                caption_dropout_every_n_epochs, caption_dropout_rate,
+                caption_dropout_every_n_epochs, caption_dropout_rate,noise_offset,
            ) = gradio_advanced_training()
            color_aug.change(
                color_aug_changed,
@ -764,7 +766,7 @@ def ti_tab(
        random_crop,
        bucket_reso_steps,
        caption_dropout_every_n_epochs, caption_dropout_rate,
-        optimizer,
+        optimizer,optimizer_args,noise_offset,
    ]

    button_open_config.click(
--- a/train_network.py
+++ b/train_network.py
@ -36,7 +36,8 @@ def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_sche
    logs["lr/unet"] = float(lr_scheduler.get_last_lr()[-1])          # may be same to textencoder

  if args.optimizer_type.lower() == "DAdaptation".lower():  # tracking d*lr value of unet.
-    logs["lr/d*lr"] = lr_scheduler.optimizers[-1].param_groups[0]['d']*lr_scheduler.optimizers[-1].param_groups[0]['lr']
+    logs["lr/d*lr-textencoder"] = lr_scheduler.optimizers[-1].param_groups[0]['d']*lr_scheduler.optimizers[-1].param_groups[0]['lr']
+    logs["lr/d*lr-unet"] = lr_scheduler.optimizers[-1].param_groups[1]['d']*lr_scheduler.optimizers[-1].param_groups[1]['lr']

  return logs