Merge pull request #83 from bmaltais/dev

Add new features from sd-script
2023-01-30 08:27:00 -05:00 · 2023-01-30 08:27:00 -05:00 · 9c5bdd1749
commit 9c5bdd1749
parent fdd1b02a26 c8f4c9d6e8
8 changed files with 103 additions and 12 deletions
--- a/README.md
+++ b/README.md
@ -1,6 +1,8 @@
 # Kohya's GUI

-This repository repository is providing a Gradio GUI for kohya's Stable Diffusion trainers found here: https://github.com/kohya-ss/sd-scripts. The GUI allow you to set the training parameters and generate and run the required CLI command to train the model.
+This repository repository is providing a Windows focussed Gradio GUI for kohya's Stable Diffusion trainers found here: https://github.com/kohya-ss/sd-scripts. The GUI allow you to set the training parameters and generate and run the required CLI command to train the model.
+
+If you run on Linux and would like to use the GUI there is now a port of it as a docker container. You can find the project here: https://github.com/P2Enjoy/kohya_ss-docker

 ## Tutorials

@ -55,7 +57,7 @@ This step is optional but can improve the learning speed for NVidia 30X0/40X0 ow

 Due to the filesize I can't host the DLLs needed for CUDNN 8.6 on Github, I strongly advise you download them for a speed boost in sample generation (almost 50% on 4090) you can download them from here: https://b1.thefileditch.ch/mwxKTEtelILoIbMbruuM.zip

-To install simply unzip the directory and place the cudnn_windows folder in the root of the kohya_diffusers_fine_tuning repo.
+To install simply unzip the directory and place the `cudnn_windows` folder in the root of the kohya_ss repo.

 Run the following command to install:

@ -141,6 +143,9 @@ Then redo the installation instruction within the kohya_ss venv.

 ## Change history

+* 2023/01/30 (v20.5.2):
+  - Add ``--lr_scheduler_num_cycles`` and ``--lr_scheduler_power`` options for ``train_network.py`` for cosine_with_restarts and polynomial learning rate schedulers. Thanks to mgz-dev!
+  - Fixed U-Net ``sample_size`` parameter to ``64`` when converting from SD to Diffusers format, in ``convert_diffusers20_original_sd.py``
 * 2023/01/27 (v20.5.1):
    - Fix issue: https://github.com/bmaltais/kohya_ss/issues/70
    - Fix issue https://github.com/bmaltais/kohya_ss/issues/71
--- a/library/common_gui.py
+++ b/library/common_gui.py
@ -520,7 +520,7 @@ def gradio_advanced_training():
            label='Shuffle caption', value=False
        )
        keep_tokens = gr.Slider(
-            label='Keen n tokens', value='0', minimum=0, maximum=32, step=1
+            label='Keep n tokens', value='0', minimum=0, maximum=32, step=1
        )
        use_8bit_adam = gr.Checkbox(label='Use 8bit adam', value=True)
        xformers = gr.Checkbox(label='Use xformers', value=True)
--- a/library/model_util.py
+++ b/library/model_util.py
@ -16,7 +16,7 @@ BETA_END = 0.0120
 UNET_PARAMS_MODEL_CHANNELS = 320
 UNET_PARAMS_CHANNEL_MULT = [1, 2, 4, 4]
 UNET_PARAMS_ATTENTION_RESOLUTIONS = [4, 2, 1]
-UNET_PARAMS_IMAGE_SIZE = 32  # unused
+UNET_PARAMS_IMAGE_SIZE = 64  # fixed from old invalid value `32`
 UNET_PARAMS_IN_CHANNELS = 4
 UNET_PARAMS_OUT_CHANNELS = 4
 UNET_PARAMS_NUM_RES_BLOCKS = 2
--- a/lora_gui.py
+++ b/lora_gui.py
@ -91,6 +91,7 @@ def save_configuration(
    max_data_loader_n_workers,
    network_alpha,
    training_comment, keep_tokens,
+    lr_scheduler_num_cycles, lr_scheduler_power,
 ):
    # Get list of function parameters and values
    parameters = list(locals().items())
@ -180,6 +181,7 @@ def open_configuration(
    max_data_loader_n_workers,
    network_alpha,
    training_comment, keep_tokens,
+    lr_scheduler_num_cycles, lr_scheduler_power,
 ):
    # Get list of function parameters and values
    parameters = list(locals().items())
@ -253,6 +255,7 @@ def train_model(
    max_data_loader_n_workers,
    network_alpha,
    training_comment, keep_tokens,
+    lr_scheduler_num_cycles, lr_scheduler_power,
 ):
    if pretrained_model_name_or_path == '':
        msgbox('Source model information is missing')
@ -395,6 +398,10 @@ def train_model(
        run_cmd += f' --gradient_accumulation_steps={int(gradient_accumulation_steps)}'
    if not output_name == '':
        run_cmd += f' --output_name="{output_name}"'
+    if not lr_scheduler_num_cycles == '':
+        run_cmd += f' --lr_scheduler_num_cycles="{lr_scheduler_num_cycles}"'
+    if not lr_scheduler_power == '':
+        run_cmd += f' --output_name="{lr_scheduler_power}"'

    run_cmd += run_cmd_training(
        learning_rate=learning_rate,
@ -646,6 +653,13 @@ def lora_tab(
                prior_loss_weight = gr.Number(
                    label='Prior loss weight', value=1.0
                )
+                lr_scheduler_num_cycles = gr.Textbox(
+                    label='LR number of cycles', placeholder='(Optional) For Cosine with restart and polynomial only'
+                )
+                
+                lr_scheduler_power = gr.Textbox(
+                    label='LR power', placeholder='(Optional) For Cosine with restart and polynomial only'
+                )
            (
                use_8bit_adam,
                xformers,
@ -736,6 +750,7 @@ def lora_tab(
        network_alpha,
        training_comment,
        keep_tokens,
+        lr_scheduler_num_cycles, lr_scheduler_power,
    ]

    button_open_config.click(
--- a/textual_inversion_gui.py
+++ b/textual_inversion_gui.py
@ -418,8 +418,8 @@ def train_model(
        use_8bit_adam=use_8bit_adam,
        keep_tokens=keep_tokens,
    )
-    run_cmd += f' --token_string={token_string}'
-    run_cmd += f' --init_word={init_word}'
+    run_cmd += f' --token_string="{token_string}"'
+    run_cmd += f' --init_word="{init_word}"'
    run_cmd += f' --num_vectors_per_token={num_vectors_per_token}'
    if not weights == '':
        run_cmd += f' --weights="{weights}"'
--- a/tools/convert_diffusers20_original_sd.py
+++ b/tools/convert_diffusers20_original_sd.py
@ -1,8 +1,4 @@
 # convert Diffusers v1.x/v2.0 model to original Stable Diffusion
-# v1: initial version
-# v2: support safetensors
-# v3: fix to support another format
-# v4: support safetensors in Diffusers

 import argparse
 import os
--- a/tools/cudann_1.8_install.py
+++ b/tools/cudann_1.8_install.py
@ -100,5 +100,7 @@ if os.name == "nt":
                if os.path.exists(dest_file):
                    shutil.copy2(src_file, cudnn_dest)
            print("Copied CUDNN 8.6 files to destination")
+    else: 
+        print(f"Installation Failed: \"{cudnn_src}\" could not be found. ")

            
--- a/train_network.py
+++ b/train_network.py
@ -35,6 +35,72 @@ def generate_step_logs(args: argparse.Namespace, current_loss, avr_loss, lr_sche
  return logs


+# Monkeypatch newer get_scheduler() function overridng current version of diffusers.optimizer.get_scheduler
+# code is taken from https://github.com/huggingface/diffusers diffusers.optimizer, commit d87cc15977b87160c30abaace3894e802ad9e1e6
+# Which is a newer release of diffusers than currently packaged with sd-scripts
+# This code can be removed when newer diffusers version (v0.12.1 or greater) is tested and implemented to sd-scripts
+
+from typing import Optional, Union
+from torch.optim import Optimizer
+from diffusers.optimization import SchedulerType, TYPE_TO_SCHEDULER_FUNCTION
+
+def get_scheduler_fix(
+    name: Union[str, SchedulerType],
+    optimizer: Optimizer,
+    num_warmup_steps: Optional[int] = None,
+    num_training_steps: Optional[int] = None,
+    num_cycles: int = 1,
+    power: float = 1.0,
+):
+    """
+    Unified API to get any scheduler from its name.
+    Args:
+        name (`str` or `SchedulerType`):
+            The name of the scheduler to use.
+        optimizer (`torch.optim.Optimizer`):
+            The optimizer that will be used during training.
+        num_warmup_steps (`int`, *optional*):
+            The number of warmup steps to do. This is not required by all schedulers (hence the argument being
+            optional), the function will raise an error if it's unset and the scheduler type requires it.
+        num_training_steps (`int``, *optional*):
+            The number of training steps to do. This is not required by all schedulers (hence the argument being
+            optional), the function will raise an error if it's unset and the scheduler type requires it.
+        num_cycles (`int`, *optional*):
+            The number of hard restarts used in `COSINE_WITH_RESTARTS` scheduler.
+        power (`float`, *optional*, defaults to 1.0):
+            Power factor. See `POLYNOMIAL` scheduler
+        last_epoch (`int`, *optional*, defaults to -1):
+            The index of the last epoch when resuming training.
+    """
+    name = SchedulerType(name)
+    schedule_func = TYPE_TO_SCHEDULER_FUNCTION[name]
+    if name == SchedulerType.CONSTANT:
+        return schedule_func(optimizer)
+
+    # All other schedulers require `num_warmup_steps`
+    if num_warmup_steps is None:
+        raise ValueError(f"{name} requires `num_warmup_steps`, please provide that argument.")
+
+    if name == SchedulerType.CONSTANT_WITH_WARMUP:
+        return schedule_func(optimizer, num_warmup_steps=num_warmup_steps)
+
+    # All other schedulers require `num_training_steps`
+    if num_training_steps is None:
+        raise ValueError(f"{name} requires `num_training_steps`, please provide that argument.")
+
+    if name == SchedulerType.COSINE_WITH_RESTARTS:
+        return schedule_func(
+            optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, num_cycles=num_cycles
+        )
+
+    if name == SchedulerType.POLYNOMIAL:
+        return schedule_func(
+            optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps, power=power
+        )
+
+    return schedule_func(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)
+
+
 def train(args):
  session_id = random.randint(0, 2**32)
  training_started_at = time.time()
@ -156,8 +222,11 @@ def train(args):
    print(f"override steps. steps for {args.max_train_epochs} epochs is / 指定エポックまでのステップ数: {args.max_train_steps}")

  # lr schedulerを用意する
-  lr_scheduler = diffusers.optimization.get_scheduler(
-      args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps, num_training_steps=args.max_train_steps * args.gradient_accumulation_steps)
+  # lr_scheduler = diffusers.optimization.get_scheduler(
+  lr_scheduler = get_scheduler_fix(
+      args.lr_scheduler, optimizer, num_warmup_steps=args.lr_warmup_steps, 
+      num_training_steps=args.max_train_steps * args.gradient_accumulation_steps,
+      num_cycles=args.lr_scheduler_num_cycles, power=args.lr_scheduler_power)

  # 実験的機能：勾配も含めたfp16学習を行う　モデル全体をfp16にする
  if args.full_fp16:
@ -445,6 +514,10 @@ if __name__ == '__main__':

  parser.add_argument("--unet_lr", type=float, default=None, help="learning rate for U-Net / U-Netの学習率")
  parser.add_argument("--text_encoder_lr", type=float, default=None, help="learning rate for Text Encoder / Text Encoderの学習率")
+  parser.add_argument("--lr_scheduler_num_cycles", type=int, default=1,
+                      help="Number of restarts for cosine scheduler with restarts / cosine with restartsスケジューラでのリスタート回数")
+  parser.add_argument("--lr_scheduler_power", type=float, default=1,
+                      help="Polynomial power for polynomial scheduler / polynomialスケジューラでのpolynomial power")

  parser.add_argument("--network_weights", type=str, default=None,
                      help="pretrained weights for network / 学習するネットワークの初期重み")