diff --git a/README.md b/README.md index a20bbf1..98018c9 100644 --- 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 diff --git a/library/common_gui.py b/library/common_gui.py index 816fd9d..2595540 100644 --- 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) diff --git a/library/model_util.py b/library/model_util.py index 6a1e656..778a4c2 100644 --- 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 diff --git a/lora_gui.py b/lora_gui.py index 49359e9..70dbfd0 100644 --- 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( diff --git a/textual_inversion_gui.py b/textual_inversion_gui.py index 5dd6015..0f21a8d 100644 --- 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}"' diff --git a/tools/convert_diffusers20_original_sd.py b/tools/convert_diffusers20_original_sd.py index a3cd03f..6c14284 100644 --- 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 diff --git a/tools/cudann_1.8_install.py b/tools/cudann_1.8_install.py index 75aef92..dec38a1 100644 --- 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. ") \ No newline at end of file diff --git a/train_network.py b/train_network.py index 8a8acc7..37a10f6 100644 --- 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 / 学習するネットワークの初期重み")