diff --git a/LoCon b/LoCon new file mode 160000 index 0000000..4ba42b9 --- /dev/null +++ b/LoCon @@ -0,0 +1 @@ +Subproject commit 4ba42b9959f9deb426b7bec81ed80111b6db7056 diff --git a/install-cn.ps1 b/install-cn.ps1 index 9000e38..5321c95 100644 --- a/install-cn.ps1 +++ b/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\ diff --git a/install.bash b/install.bash index 08b5def..de4def6 100644 --- a/install.bash +++ b/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" diff --git a/install.ps1 b/install.ps1 index 166d621..7b6b8a2 100644 --- a/install.ps1 +++ b/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 ; \ No newline at end of file diff --git a/train.ps1 b/train.ps1 index 4e12dda..df71cf3 100644 --- a/train.ps1 +++ b/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 ` diff --git a/train.sh b/train.sh index 8595fbb..9662c32 100644 --- a/train.sh +++ b/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[@]} \ No newline at end of file + --xformers --shuffle_caption ${extArgs[@]}