Merge pull request #209 from bmaltais/dev

v20.7.4

README-ja.md

@@ -64,6 +64,12 @@ cp .\bitsandbytes_windows\main.py .\venv\Lib\site-packages\bitsandbytes\cuda_set
 accelerate config
 ```
 
+<!-- 
+pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 --extra-index-url https://download.pytorch.org/whl/cu117
+pip install --use-pep517 --upgrade -r requirements.txt
+pip install -U -I --no-deps xformers==0.0.16
+-->
+
 コマンドプロンプトでは以下になります。
 
 

README.md

@@ -163,9 +163,23 @@ This will store your a backup file with your current locally installed pip packa
 
 ## Change History
 
-* 2023/02/16 (v20.7.3)
+* 2023/02/19 (v20.7.4):
+    - Add `--use_lion_optimizer` to each training script to use [Lion optimizer](https://github.com/lucidrains/lion-pytorch).
+        - Please install Lion optimizer with `pip install lion-pytorch` (it is not in ``requirements.txt`` currently.)
+    - Add `--lowram` option to `train_network.py`. Load models to VRAM instead of VRAM (for machines which have bigger VRAM than RAM such as Colab and Kaggle). Thanks to Isotr0py!
+        - Default behavior (without lowram) has reverted to the same as before 14 Feb.
+    - Fixed git commit hash to be set correctly regardless of the working directory. Thanks to vladmandic!
+* 2023/02/15 (v20.7.3):
+    - Update upgrade.ps1 script
+    - Integrate new kohya sd-script
     - Noise offset is recorded to the metadata. Thanks to space-nuko!
     - Show the moving average loss to prevent loss jumping in `train_network.py` and `train_db.py`. Thanks to shirayu!
+    - Add support with multi-gpu trainining for `train_network.py`. Thanks to Isotr0py!
+    - Add `--verbose` option for `resize_lora.py`. For details, see [this PR](https://github.com/kohya-ss/sd-scripts/pull/179). Thanks to mgz-dev!
+    - Git commit hash is added to the metadata for LoRA. Thanks to space-nuko!
+    - Add `--noise_offset` option for each training scripts. 
+        - Implementation of https://www.crosslabs.org//blog/diffusion-with-offset-noise 
+        - This option may improve ability to generate darker/lighter images. May work with LoRA.
 * 2023/02/11 (v20.7.2):
     - `lora_interrogator.py` is added in `networks` folder. See `python networks\lora_interrogator.py -h` for usage.
         - For LoRAs where the activation word is unknown, this script compares the output of Text Encoder after applying LoRA to that of unapplied to find out which token is affected by LoRA. Hopefully you can figure out the activation word. LoRA trained with captions does not seem to be able to interrogate.

dreambooth_gui.py

@@ -89,6 +89,7 @@ def save_configuration(
     random_crop,
     bucket_reso_steps,
     caption_dropout_every_n_epochs, caption_dropout_rate,
+    optimizer,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -179,6 +180,7 @@ def open_configuration(
     random_crop,
     bucket_reso_steps,
     caption_dropout_every_n_epochs, caption_dropout_rate,
+    optimizer,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -253,6 +255,7 @@ def train_model(
     random_crop,
     bucket_reso_steps,
     caption_dropout_every_n_epochs, caption_dropout_rate,
+    optimizer,
 ):
     if pretrained_model_name_or_path == '':
         msgbox('Source model information is missing')
@@ -397,6 +400,7 @@ def train_model(
         seed=seed,
         caption_extension=caption_extension,
         cache_latents=cache_latents,
+        optimizer=optimizer
     )
 
     run_cmd += run_cmd_advanced_training(
@@ -541,6 +545,7 @@ def dreambooth_tab(
             seed,
             caption_extension,
             cache_latents,
+            optimizer,
         ) = gradio_training(
             learning_rate_value='1e-5',
             lr_scheduler_value='cosine',
@@ -668,6 +673,7 @@ def dreambooth_tab(
         random_crop,
         bucket_reso_steps,
         caption_dropout_every_n_epochs, caption_dropout_rate,
+        optimizer,
     ]
 
     button_open_config.click(

fine_tune.py

@@ -161,6 +161,13 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
       raise ImportError("No bitsand bytes / bitsandbytesがインストールされていないようです")
     print("use 8-bit Adam optimizer")
     optimizer_class = bnb.optim.AdamW8bit
+  elif args.use_lion_optimizer:
+    try:
+      import lion_pytorch
+    except ImportError:
+      raise ImportError("No lion_pytorch / lion_pytorch がインストールされていないようです")
+    print("use Lion optimizer")
+    optimizer_class = lion_pytorch.Lion
   else:
     optimizer_class = torch.optim.AdamW
 
@@ -272,6 +279,9 @@ def fn_recursive_set_mem_eff(module: torch.nn.Module):
 
         # Sample noise that we'll add to the latents
         noise = torch.randn_like(latents, device=latents.device)
+        if args.noise_offset:
+          # https://www.crosslabs.org//blog/diffusion-with-offset-noise
+          noise += args.noise_offset * torch.randn((latents.shape[0], latents.shape[1], 1, 1), device=latents.device)
 
         # Sample a random timestep for each image
         timesteps = torch.randint(0, noise_scheduler.config.num_train_timesteps, (b_size,), device=latents.device)

finetune_gui.py

@@ -85,6 +85,7 @@ def save_configuration(
     random_crop,
     bucket_reso_steps,
     caption_dropout_every_n_epochs, caption_dropout_rate,
+    optimizer,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -181,6 +182,7 @@ def open_config_file(
     random_crop,
     bucket_reso_steps,
     caption_dropout_every_n_epochs, caption_dropout_rate,
+    optimizer,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -262,6 +264,7 @@ def train_model(
     random_crop,
     bucket_reso_steps,
     caption_dropout_every_n_epochs, caption_dropout_rate,
+    optimizer,
 ):
     # create caption json file
     if generate_caption_database:
@@ -386,6 +389,7 @@ def train_model(
         seed=seed,
         caption_extension=caption_extension,
         cache_latents=cache_latents,
+        optimizer=optimizer,
     )
 
     run_cmd += run_cmd_advanced_training(
@@ -564,6 +568,7 @@ def finetune_tab():
             seed,
             caption_extension,
             cache_latents,
+            optimizer,
         ) = gradio_training(learning_rate_value='1e-5')
         with gr.Row():
             dataset_repeats = gr.Textbox(label='Dataset repeats', value=40)
@@ -661,6 +666,7 @@ def finetune_tab():
         random_crop,
         bucket_reso_steps,
         caption_dropout_every_n_epochs, caption_dropout_rate,
+        optimizer,
     ]
 
     button_run.click(train_model, inputs=settings_list)

library/common_gui.py

@@ -445,6 +445,7 @@ def gradio_training(
             value=2,
         )
         seed = gr.Textbox(label='Seed', value=1234)
+        cache_latents = gr.Checkbox(label='Cache latent', value=True)
     with gr.Row():
         learning_rate = gr.Textbox(
             label='Learning rate', value=learning_rate_value
@@ -464,7 +465,15 @@ def gradio_training(
         lr_warmup = gr.Textbox(
             label='LR warmup (% of steps)', value=lr_warmup_value
         )
-        cache_latents = gr.Checkbox(label='Cache latent', value=True)
+        optimizer = gr.Dropdown(
+            label='Optimizer',
+            choices=[
+                'AdamW',
+                'Lion',
+            ],
+            value="AdamW",
+            interactive=True,
+        )
     return (
         learning_rate,
         lr_scheduler,
@@ -478,6 +487,7 @@ def gradio_training(
         seed,
         caption_extension,
         cache_latents,
+        optimizer,
     )
 
 
@@ -512,10 +522,34 @@ 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 '',
     ]
     run_cmd = ''.join(options)
     return run_cmd
 
+# # This function takes a dictionary of keyword arguments and returns a string that can be used to run a command-line training script
+# def run_cmd_training(**kwargs):
+#     arg_map = {
+#         'learning_rate': ' --learning_rate="{}"',
+#         'lr_scheduler': ' --lr_scheduler="{}"',
+#         'lr_warmup_steps': ' --lr_warmup_steps="{}"',
+#         'train_batch_size': ' --train_batch_size="{}"',
+#         'max_train_steps': ' --max_train_steps="{}"',
+#         'save_every_n_epochs': ' --save_every_n_epochs="{}"',
+#         'mixed_precision': ' --mixed_precision="{}"',
+#         'save_precision': ' --save_precision="{}"',
+#         'seed': ' --seed="{}"',
+#         'caption_extension': ' --caption_extension="{}"',
+#         'cache_latents': ' --cache_latents',
+#         'optimizer': ' --use_lion_optimizer' if kwargs.get('optimizer') == 'Lion' else '',
+#     }
+
+#     options = [arg_map[key].format(value) for key, value in kwargs.items() if key in arg_map and value]
+
+#     cmd = ''.join(options)
+
+#     return cmd
+
 
 def gradio_advanced_training():
     with gr.Row():
@@ -664,3 +698,34 @@ def run_cmd_advanced_training(**kwargs):
     ]
     run_cmd = ''.join(options)
     return run_cmd
+
+# def run_cmd_advanced_training(**kwargs):
+#     arg_map = {
+#         'max_train_epochs': ' --max_train_epochs="{}"',
+#         'max_data_loader_n_workers': ' --max_data_loader_n_workers="{}"',
+#         'max_token_length': ' --max_token_length={}' if int(kwargs.get('max_token_length', 75)) > 75 else '',
+#         'clip_skip': ' --clip_skip={}' if int(kwargs.get('clip_skip', 1)) > 1 else '',
+#         'resume': ' --resume="{}"',
+#         'keep_tokens': ' --keep_tokens="{}"' if int(kwargs.get('keep_tokens', 0)) > 0 else '',
+#         'caption_dropout_every_n_epochs': ' --caption_dropout_every_n_epochs="{}"' if int(kwargs.get('caption_dropout_every_n_epochs', 0)) > 0 else '',
+#         'caption_dropout_rate': ' --caption_dropout_rate="{}"' if float(kwargs.get('caption_dropout_rate', 0)) > 0 else '',
+#         'bucket_reso_steps': ' --bucket_reso_steps={:d}' if int(kwargs.get('bucket_reso_steps', 64)) >= 1 else '',
+#         'save_state': ' --save_state',
+#         'mem_eff_attn': ' --mem_eff_attn',
+#         'color_aug': ' --color_aug',
+#         'flip_aug': ' --flip_aug',
+#         'shuffle_caption': ' --shuffle_caption',
+#         'gradient_checkpointing': ' --gradient_checkpointing',
+#         'full_fp16': ' --full_fp16',
+#         'xformers': ' --xformers',
+#         'use_8bit_adam': ' --use_8bit_adam',
+#         'persistent_data_loader_workers': ' --persistent_data_loader_workers',
+#         'bucket_no_upscale': ' --bucket_no_upscale',
+#         'random_crop': ' --random_crop',
+#     }
+
+#     options = [arg_map[key].format(value) for key, value in kwargs.items() if key in arg_map and value]
+
+#     cmd = ''.join(options)
+
+#     return cmd

library/train_util.py

@@ -12,6 +12,7 @@
 import os
 import random
 import hashlib
+import subprocess
 from io import BytesIO
 
 from tqdm import tqdm
@@ -299,7 +300,7 @@ def dropout_tags(tokens):
         if self.shuffle_keep_tokens is None:
           if self.shuffle_caption:
             random.shuffle(tokens)
-          
+
           tokens = dropout_tags(tokens)
         else:
           if len(tokens) > self.shuffle_keep_tokens:
@@ -308,7 +309,7 @@ def dropout_tags(tokens):
 
             if self.shuffle_caption:
               random.shuffle(tokens)
-            
+
             tokens = dropout_tags(tokens)
 
             tokens = keep_tokens + tokens
@@ -1100,6 +1101,13 @@ def addnet_hash_safetensors(b):
   return hash_sha256.hexdigest()
 
 
+def get_git_revision_hash() -> str:
+  try:
+    return subprocess.check_output(['git', 'rev-parse', 'HEAD'], cwd=os.path.dirname(__file__)).decode('ascii').strip()
+  except:
+    return "(unknown)"
+
+
 # flash attention forwards and backwards
 
 # https://arxiv.org/abs/2205.14135
@@ -1381,6 +1389,8 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth:
                       help="max token length of text encoder (default for 75, 150 or 225) / text encoderのトークンの最大長（未指定で75、150または225が指定可）")
   parser.add_argument("--use_8bit_adam", action="store_true",
                       help="use 8bit Adam optimizer (requires bitsandbytes) / 8bit Adamオプティマイザを使う（bitsandbytesのインストールが必要）")
+  parser.add_argument("--use_lion_optimizer", action="store_true",
+                      help="use Lion optimizer (requires lion-pytorch) / Lionオプティマイザを使う（ lion-pytorch のインストールが必要）")
   parser.add_argument("--mem_eff_attn", action="store_true",
                       help="use memory efficient attention for CrossAttention / CrossAttentionに省メモリ版attentionを使う")
   parser.add_argument("--xformers", action="store_true",
@@ -1413,6 +1423,10 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth:
                       help="scheduler to use for learning rate / 学習率のスケジューラ: linear, cosine, cosine_with_restarts, polynomial, constant (default), constant_with_warmup")
   parser.add_argument("--lr_warmup_steps", type=int, default=0,
                       help="Number of steps for the warmup in the lr scheduler (default is 0) / 学習率のスケジューラをウォームアップするステップ数（デフォルト0）")
+  parser.add_argument("--noise_offset", type=float, default=None,
+                      help="enable noise offset with this value (if enabled, around 0.1 is recommended) / Noise offsetを有効にしてこの値を設定する（有効にする場合は0.1程度を推奨）")
+  parser.add_argument("--lowram", action="store_true",
+                      help="enable low RAM optimization. e.g. load models to VRAM instead of RAM (for machines which have bigger VRAM than RAM such as Colab and Kaggle) / メインメモリが少ない環境向け最適化を有効にする。たとえばVRAMにモデルを読み込むなど（ColabやKaggleなどRAMに比べてVRAMが多い環境向け）")
 
   if support_dreambooth:
     # DreamBooth training
@@ -1620,9 +1634,6 @@ def get_hidden_states(args: argparse.Namespace, input_ids, tokenizer, text_encod
   else:
     enc_out = text_encoder(input_ids, output_hidden_states=True, return_dict=True)
     encoder_hidden_states = enc_out['hidden_states'][-args.clip_skip]
-    if weight_dtype is not None:
-      # this is required for additional network training
-      encoder_hidden_states = encoder_hidden_states.to(weight_dtype)
     encoder_hidden_states = text_encoder.text_model.final_layer_norm(encoder_hidden_states)
 
   # bs*3, 77, 768 or 1024
@@ -1649,6 +1660,10 @@ def get_hidden_states(args: argparse.Namespace, input_ids, tokenizer, text_encod
       states_list.append(encoder_hidden_states[:, -1].unsqueeze(1))                         # <EOS>
       encoder_hidden_states = torch.cat(states_list, dim=1)
 
+  if weight_dtype is not None:
+    # this is required for additional network training
+    encoder_hidden_states = encoder_hidden_states.to(weight_dtype)
+
   return encoder_hidden_states
 
 

lora_gui.py

@@ -100,6 +100,7 @@ def save_configuration(
     random_crop,
     bucket_reso_steps,
     caption_dropout_every_n_epochs, caption_dropout_rate,
+    optimizer,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -197,6 +198,7 @@ def open_configuration(
     random_crop,
     bucket_reso_steps,
     caption_dropout_every_n_epochs, caption_dropout_rate,
+    optimizer,
 ):
     # Get list of function parameters and values
     parameters = list(locals().items())
@@ -278,6 +280,7 @@ def train_model(
     random_crop,
     bucket_reso_steps,
     caption_dropout_every_n_epochs, caption_dropout_rate,
+    optimizer,
 ):  
     if pretrained_model_name_or_path == '':
         msgbox('Source model information is missing')
@@ -457,6 +460,7 @@ def train_model(
         seed=seed,
         caption_extension=caption_extension,
         cache_latents=cache_latents,
+        optimizer=optimizer,
     )
 
     run_cmd += run_cmd_advanced_training(
@@ -609,6 +613,7 @@ def lora_tab(
             seed,
             caption_extension,
             cache_latents,
+            optimizer,
         ) = gradio_training(
             learning_rate_value='0.0001',
             lr_scheduler_value='cosine',
@@ -778,6 +783,7 @@ def lora_tab(
         random_crop,
         bucket_reso_steps,
         caption_dropout_every_n_epochs, caption_dropout_rate,
+        optimizer,
     ]
 
     button_open_config.click(