support lion optimizer

install-cn.ps1

@@ -34,6 +34,8 @@ pip install --upgrade -r requirements.txt -i https://mirrors.bfsu.edu.cn/pypi/we
 Check "其他依赖安装失败。"
 pip install -U -I --no-deps https://jihulab.com/api/v4/projects/82097/packages/pypi/files/e8508fe14c8f2552a822f5e6f5620b24fdd4ba3129c2a31a39b56425bcc023bc/xformers-0.0.14.dev0+torch12-cp310-cp310-win_amd64.whl
 Check "xformers 安装失败。"
+pip install --upgrade -r lion-pytorch -i https://mirrors.bfsu.edu.cn/pypi/web/simple
+Check "Lion 优化器安装失败。"
 
 Write-Output "安装 bitsandbytes..."
 cp .\bitsandbytes_windows\*.dll ..\venv\Lib\site-packages\bitsandbytes\

train.ps1

@@ -30,6 +30,11 @@ $save_model_as = "safetensors" # model save ext | 模型保存格式 ckpt, pt, s
 $network_weights = "" # pretrained weights for LoRA network | 若需要从已有的 LoRA 模型上继续训练，请填写 LoRA 模型路径。
 $min_bucket_reso = 256 # arb min resolution | arb 最小分辨率
 $max_bucket_reso = 1024 # arb max resolution | arb 最大分辨率
+$persistent_data_loader_workers = 0 # persistent dataloader workers | 容易爆显存，保留加载训练集的worker，减少每个 epoch 之间的停顿
+
+# 优化器设置
+$use_8bit_adam = 1 # use 8bit adam optimizer | 使用 8bit adam 优化器节省显存，默认启用。部分 10 系老显卡无法使用，修改为 0 禁用。
+$use_lion = 0 # use lion optimizer | 使用 Lion 优化器
 
 
 # ============= DO NOT MODIFY CONTENTS BELOW | 请勿修改下方内容 =====================
@@ -51,6 +56,18 @@ if ($network_weights) {
   [void]$ext_args.Add("--network_weights=" + $network_weights)
 }
 
+if ($use_8bit_adam) {
+  [void]$ext_args.Add("--use_8bit_adam")
+}
+
+if ($use_lion) {
+  [void]$ext_args.Add("--use_lion_optimizer")
+}
+
+if ($persistent_data_loader_workers) {
+  [void]$ext_args.Add("--persistent_data_loader_workers")
+}
+
 # run train
 accelerate launch --num_cpu_threads_per_process=8 "./sd-scripts/train_network.py" `
   --enable_bucket `
@@ -82,7 +99,7 @@ accelerate launch --num_cpu_threads_per_process=8 "./sd-scripts/train_network.py
   --save_model_as=$save_model_as `
   --min_bucket_reso=$min_bucket_reso `
   --max_bucket_reso=$max_bucket_reso `
-  --xformers --shuffle_caption --use_8bit_adam $ext_args
+  --xformers --shuffle_caption $ext_args
 
 Write-Output "Train finished"
 Read-Host | Out-Null ;