diff --git a/fine_tune.py b/fine_tune.py index 3ba63063..13241bc6 100644 --- a/fine_tune.py +++ b/fine_tune.py @@ -158,6 +158,13 @@ def train(args): 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 diff --git a/library/train_util.py b/library/train_util.py index 441838e5..63868f98 100644 --- a/library/train_util.py +++ b/library/train_util.py @@ -1389,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", @@ -1424,8 +1426,8 @@ def add_training_arguments(parser: argparse.ArgumentParser, support_dreambooth: 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="load models to VRAM instead of RAM (for machines which have bigger VRAM than RAM such as Colab and Kaggle)") - + 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 parser.add_argument("--prior_loss_weight", type=float, default=1.0, diff --git a/train_db.py b/train_db.py index e4f1e54c..1903c4c4 100644 --- a/train_db.py +++ b/train_db.py @@ -124,6 +124,13 @@ def train(args): 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 diff --git a/train_network.py b/train_network.py index e29e0174..b41a52a9 100644 --- a/train_network.py +++ b/train_network.py @@ -156,10 +156,12 @@ def train(args): # モデルを読み込む text_encoder, vae, unet, _ = train_util.load_target_model(args, weight_dtype) + # work on low-ram device if args.lowram: text_encoder.to("cuda") unet.to("cuda") + # モデルに xformers とか memory efficient attention を組み込む train_util.replace_unet_modules(unet, args.mem_eff_attn, args.xformers) @@ -214,6 +216,13 @@ def train(args): 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 diff --git a/train_textual_inversion.py b/train_textual_inversion.py index 010bd04b..ffec0516 100644 --- a/train_textual_inversion.py +++ b/train_textual_inversion.py @@ -207,6 +207,13 @@ def train(args): 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