feat: support locon

LoCon

@@ -0,0 +1 @@
+Subproject commit 4ba42b9959f9deb426b7bec81ed80111b6db7056

install-cn.ps1

@@ -36,6 +36,9 @@ pip install -U -I --no-deps https://jihulab.com/api/v4/projects/82097/packages/p
 Check "xformers 安装失败。"
 pip install --upgrade lion-pytorch -i https://mirrors.bfsu.edu.cn/pypi/web/simple
 Check "Lion 优化器安装失败。"
+pip install --upgrade locon -i https://mirrors.bfsu.edu.cn/pypi/web/simple
+Check "locon 安装失败。"
+
 
 Write-Output "安装 bitsandbytes..."
 cp .\bitsandbytes_windows\*.dll ..\venv\Lib\site-packages\bitsandbytes\

install.bash

@@ -7,6 +7,6 @@ pip install torch==1.12.1+cu116 torchvision==0.13.1+cu116 --extra-index-url http
 pip install --upgrade -r requirements.txt
 pip install git+https://github.com/facebookresearch/xformers.git@0bad001ddd56c080524d37c84ff58d9cd030ebfd
 pip install triton==2.0.0.dev20221202
-pip install --upgrade lion-pytorch
+pip install --upgrade lion-pytorch locon
 
 echo "Install completed"

install.ps1

@@ -1,4 +1,4 @@
-$Env:HF_HOME="huggingface"
+$Env:HF_HOME = "huggingface"
 
 if (!(Test-Path -Path "venv")) {
     Write-Output  "Creating venv for python..."
@@ -17,5 +17,7 @@ cp .\bitsandbytes_windows\*.dll ..\venv\Lib\site-packages\bitsandbytes\
 cp .\bitsandbytes_windows\cextension.py ..\venv\Lib\site-packages\bitsandbytes\cextension.py
 cp .\bitsandbytes_windows\main.py ..\venv\Lib\site-packages\bitsandbytes\cuda_setup\main.py
 
+pip install --upgrade lion-pytorch locon
+
 Write-Output "Install completed"
 Read-Host | Out-Null ;

train.ps1

@@ -38,12 +38,17 @@ $persistent_data_loader_workers = 0 # persistent dataloader workers | 容易爆
 $use_8bit_adam = 1 # use 8bit adam optimizer | 使用 8bit adam 优化器节省显存，默认启用。部分 10 系老显卡无法使用，修改为 0 禁用。
 $use_lion = 0 # use lion optimizer | 使用 Lion 优化器
 
+# LoCon 训练设置
+$enable_locon_train = 0 # enable LoCon train | 启用 LoCon 训练 启用后 network_dim 和 network_alpha 应当选择较小的值，比如 2~16
+$conv_dim = 4 # conv dim | 类似于 network_dim，推荐为 4
+$conv_alpha = 4 # conv alpha | 类似于 network_alpha，可以采用与 conv_dim 一致或者更小的值
 
 # ============= DO NOT MODIFY CONTENTS BELOW | 请勿修改下方内容 =====================
 # Activate python venv
 .\venv\Scripts\activate
 
 $Env:HF_HOME = "huggingface"
+$network_module = "networks.lora"
 $ext_args = [System.Collections.ArrayList]::new()
 
 if ($train_unet_only) {
@@ -74,6 +79,13 @@ if ($persistent_data_loader_workers) {
   [void]$ext_args.Add("--persistent_data_loader_workers")
 }
 
+if ($enable_locon_train) {
+  $network_module = "locon.locon_kohya"
+  [void]$ext_args.Add("--network_args")
+  [void]$ext_args.Add("conv_dim=$conv_dim")
+  [void]$ext_args.Add("conv_alpha=$conv_alpha")
+}
+
 # run train
 accelerate launch --num_cpu_threads_per_process=8 "./sd-scripts/train_network.py" `
   --enable_bucket `
@@ -82,7 +94,7 @@ accelerate launch --num_cpu_threads_per_process=8 "./sd-scripts/train_network.py
   --output_dir="./output" `
   --logging_dir="./logs" `
   --resolution=$resolution `
-  --network_module=networks.lora `
+  --network_module=$network_module `
   --max_train_epochs=$max_train_epoches `
   --learning_rate=$lr `
   --unet_lr=$unet_lr `

train.sh

@@ -2,18 +2,18 @@
 
 # Train data path | 设置训练用模型、图片
 pretrained_model="./sd-models/model.ckpt" # base model path | 底模路径
-train_data_dir="./train/aki" # train dataset path | 训练数据集路径
-reg_data_dir="" # directory for regularization images | 正则化数据集路径，默认不使用正则化图像。
+train_data_dir="./train/aki"              # train dataset path | 训练数据集路径
+reg_data_dir=""                           # directory for regularization images | 正则化数据集路径，默认不使用正则化图像。
 
 # Train related params | 训练相关参数
-resolution="512,512" # image resolution w,h. 图片分辨率，宽,高。支持非正方形，但必须是 64 倍数。
-batch_size=1 # batch size
-max_train_epoches=10 # max train epoches | 最大训练 epoch
-save_every_n_epochs=2 # save every n epochs | 每 N 个 epoch 保存一次
-network_dim=32 # network dim | 常用 4~128，不是越大越好
-network_alpha=32 # network alpha | 常用与 network_dim 相同的值或者采用较小的值，如 network_dim的一半 防止下溢。默认值为 1，使用较小的 alpha 需要提升学习率。
-clip_skip=2 # clip skip | 玄学 一般用 2
-train_unet_only=0 # train U-Net only | 仅训练 U-Net，开启这个会牺牲效果大幅减少显存使用。6G显存可以开启
+resolution="512,512"      # image resolution w,h. 图片分辨率，宽,高。支持非正方形，但必须是 64 倍数。
+batch_size=1              # batch size
+max_train_epoches=10      # max train epoches | 最大训练 epoch
+save_every_n_epochs=2     # save every n epochs | 每 N 个 epoch 保存一次
+network_dim=32            # network dim | 常用 4~128，不是越大越好
+network_alpha=32          # network alpha | 常用与 network_dim 相同的值或者采用较小的值，如 network_dim的一半 防止下溢。默认值为 1，使用较小的 alpha 需要提升学习率。
+clip_skip=2               # clip skip | 玄学 一般用 2
+train_unet_only=0         # train U-Net only | 仅训练 U-Net，开启这个会牺牲效果大幅减少显存使用。6G显存可以开启
 train_text_encoder_only=0 # train Text Encoder only | 仅训练 文本编码器
 
 # Learning rate | 学习率
@@ -21,28 +21,33 @@ lr="1e-4"
 unet_lr="1e-4"
 text_encoder_lr="1e-5"
 lr_scheduler="cosine_with_restarts" # "linear", "cosine", "cosine_with_restarts", "polynomial", "constant", "constant_with_warmup"
-lr_warmup_steps=0 # warmup steps | 仅在 lr_scheduler 为 constant_with_warmup 时需要填写这个值
-lr_restart_cycles=1 # cosine_with_restarts restart cycles | 余弦退火重启次数，仅在 lr_scheduler 为 cosine_with_restarts 时起效。
+lr_warmup_steps=0                   # warmup steps | 仅在 lr_scheduler 为 constant_with_warmup 时需要填写这个值
+lr_restart_cycles=1                 # cosine_with_restarts restart cycles | 余弦退火重启次数，仅在 lr_scheduler 为 cosine_with_restarts 时起效。
 
 # Output settings | 输出设置
-output_name="aki" # output model name | 模型保存名称
+output_name="aki"           # output model name | 模型保存名称
 save_model_as="safetensors" # model save ext | 模型保存格式 ckpt, pt, safetensors
 
 # 其他设置
-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 最大分辨率
+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 优化器
+use_lion=0      # use lion optimizer | 使用 Lion 优化器
 
+# LoCon 训练设置
+enable_locon_train=0 # enable LoCon train | 启用 LoCon 训练 启用后 network_dim 和 network_alpha 应当选择较小的值，比如 2~16
+conv_dim=4           # conv dim | 类似于 network_dim，推荐为 4
+conv_alpha=4         # conv alpha | 类似于 network_alpha，可以采用与 conv_dim 一致或者更小的值
 
 # ============= DO NOT MODIFY CONTENTS BELOW | 请勿修改下方内容 =====================
 export HF_HOME="huggingface"
 export TF_CPP_MIN_LOG_LEVEL=3
 
+network_module="networks.lora"
 extArgs=()
 
 if [ $train_unet_only == 1 ]; then extArgs+=("--network_train_unet_only"); fi
@@ -59,6 +64,11 @@ if [ $use_lion == 1 ]; then extArgs+=("--use_lion_optimizer"); fi
 
 if [ $persistent_data_loader_workers == 1 ]; then extArgs+=("--persistent_data_loader_workers"); fi
 
+if [ $enable_locon_train == 1 ]; then
+  network_module="locon.locon_kohya"
+  extArgs+=("--network_args conv_dim=$conv_dim conv_alpha=$conv_alpha")
+fi
+
 accelerate launch --num_cpu_threads_per_process=8 "./sd-scripts/train_network.py" \
   --enable_bucket \
   --pretrained_model_name_or_path=$pretrained_model \
@@ -66,7 +76,7 @@ accelerate launch --num_cpu_threads_per_process=8 "./sd-scripts/train_network.py
   --output_dir="./output" \
   --logging_dir="./logs" \
   --resolution=$resolution \
-  --network_module=networks.lora \
+  --network_module=$network_module \
   --max_train_epochs=$max_train_epoches \
   --learning_rate=$lr \
   --unet_lr=$unet_lr \
@@ -90,4 +100,4 @@ 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 ${extArgs[@]}
+  --xformers --shuffle_caption ${extArgs[@]}