train_network.ps1

# https://github.com/cloneofsimo/lora
# https://github.com/kohya-ss/sd-scripts

##### Добавь --v2 в аргументы если загружаешь SD 2.x чекпоинт! #####

##### Начало конфига #####

$sd_scripts_dir = "X:\git-repos\sd-scripts\" # Путь к папке с репозиторием kohya-ss/sd-scripts

$ckpt = "X:\SD-models\checkpoint.safetensors" # Путь к чекпоинту (ckpt / safetensors)
$sd_v2_ckpt = 0 # Поставь '1' если загружаешь SD 2.x чекпоинт
$image_dir = "X:\training_data\img\" # Путь к папке с изображениями
$reg_dir = "X:\training_data\img_reg\" # Путь к папке с регуляризационными изображениями (можно указать на пустую папку, но путь обязательно должен быть указан)
$output_dir = "X:\LoRA\" # Директория сохранения LoRA чекпоинтов
$output_name = "my_LoRA_network_v1" # Название файла (расширение не нужно)

$learning_rate = 1e-4 # Рекомендуемое значение из официального репозитория LoRA
$train_batch_size = 1 # Сколько изображений тренировать одновременно. Чем больше значение, тем быстрее тренировка, но больше потребление видеопамяти
$resolution = 512 # Разрешение тренировки
$num_epochs = 5 # Число эпох
$save_every_n_epochs = 1 # Сохранять чекпоинт каждые n эпох
$save_last_n_epochs = 999 # Сохранить только последние n эпох
$scheduler = "cosine_with_restarts" # linear, cosine, cosine_with_restarts, polynomial, constant (по-умолчанию), constant_with_warmup
$network_dim = 128 # Размер нетворка. Чем больше значение, тем больше точность и размер выходного файла
$max_token_length = 150 # Максимальная длина токена. Возможные значения: None / 150 / 225. None = 75
                        # upd. "None" почему-то не работает
$clip_skip = 1 # https://github.com/AUTOMATIC1111/stable-diffusion-webui/wiki/Features#ignore-last-layers-of-clip-model

##### Конец конфига #####

function Is-Numeric ($value) { return $value -match "^[\d\.]+$" }

function Write-ColorOutput($ForegroundColor)
{
    $fc = $host.UI.RawUI.ForegroundColor
    $host.UI.RawUI.ForegroundColor = $ForegroundColor
    if ($args) { Write-Output $args }
    else { $input | Write-Output }
    $host.UI.RawUI.ForegroundColor = $fc
}

function Word-Ending($value)
{
    $ending = $value.ToString()
    if ($ending -ge "11" -and $ending -le "19") { return "й" }
    $ending = $ending.Substring([Math]::Max($ending.Length, 0) - 1)
    if ($ending -eq "1") { return "е" }
    if ($ending -ge "2" -and $ending -le "4") { return "я" }
    if (($ending -ge "5" -and $ending -le "9") -or $ending -eq "0") { return "й" }
}

Write-Output "Подсчет количества изображений в папках"
$total = 0
$is_structure_wrong = 0
$abort_script = 0
$iter = 0

Get-ChildItem -Path $image_dir -Directory | ForEach-Object {
    $parts = $_.Name.Split("_")
    if (!(Is-Numeric $parts[0]))
    {
        Write-ColorOutput red "Ошибка в $($_):`n`t$($parts[0]) не является числом"
        $is_structure_wrong = 1
        return
    }
    if ([int]$parts[0] -le 0)
    {
        Write-ColorOutput red "Ошибка в $($_):`n`tИмя папки с изображениями не может начинаться с '0'"
        $is_structure_wrong = 1
        return
    }
    $repeats = [int]$parts[0]
    $imgs = Get-ChildItem $_.FullName -Depth 0 -File -Include *.jpg, *.png, *.webp | Measure-Object | ForEach-Object { $_.Count }
    if ($iter -eq 0) { Write-Output "Обучающие изображения:" }
    $img_repeats = ($repeats * $imgs)
    Write-Output "`t$($parts[1]): $repeats повторени$(Word-Ending $repeats) * $imgs изображени$(Word-Ending $imgs) = $($img_repeats)"
    $total += $img_repeats
    $iter += 1
}

$iter = 0

if ($is_structure_wrong -eq 0) { Get-ChildItem -Path $reg_dir -Directory | % { if ($abort_script -ne "n") { ForEach-Object {
    $parts = $_.Name.Split("_")
    if (!(Is-Numeric $parts[0]))
    {
        Write-ColorOutput red "Ошибка в $($_):`n`t$($parts[0]) не является числом"
        $is_structure_wrong = 1
        return
    }
    if ([int]$parts[0] -le 0)
    {
        Write-ColorOutput red "Ошибка в $($_):`nИмя папки с изображениями не может начинаться '0'"
        $is_structure_wrong = 1
        return
    }
    $repeats = [int]$parts[0]
    $reg_imgs = Get-ChildItem $_.FullName -Depth 0 -File -Include *.jpg, *.png, *.webp | Measure-Object | ForEach-Object { $_.Count }
    if ($iter -eq 0) { Write-Output "Регуляризационные изображения:" }
    if ($reg_imgs -eq 0)
    {
        Write-ColorOutput darkyellow "Внимание: папка для регуляризационных изображений присутствует, но в ней ничего нет"
        do { $abort_script = Read-Host "Прервать выполнение скрипта? (y/N)" }
        until (($abort_script -eq "y") -or ($abort_script -ceq "N"))
        return
    }
    else
    {
        $img_repeats = ($repeats * $reg_imgs)
        Write-Output "`t$($parts[1]): $repeats повторени$(Word-Ending $repeats) * $reg_imgs изображени$(Word-Ending $reg_imgs) = $($img_repeats)"
        $iter += 1
    }
} } } }

if ($is_structure_wrong -eq 0 -and ($abort_script -eq "n" -or $abort_script -eq 0))
{
    if ($reg_imgs -gt 0)
    {
        $total *= 2
        Write-Output "Количество шагов увеличено вдвое: количество регуляризационных изображений больше 0"
    }

    Write-Output "Количество изображений с повторениями: $total"
    $max_training_steps = [int]($total / $train_batch_size * $num_epochs)
    Write-Output "Размер обучающей партии: $train_batch_size"
    Write-Output "Количество эпох: $num_epochs"
    Write-Output "Количество шагов: $total / $train_batch_size * $num_epochs = $max_training_steps"
    Write-Output "Выполнение скрипта...`n"

    $seed = Get-Random

    $script_origin = (get-location).path
    cd $sd_scripts_dir
    .\venv\Scripts\activate

    accelerate launch --num_cpu_threads_per_process 12 train_network.py `
    --network_module=networks.lora `
    --pretrained_model_name_or_path=$ckpt `
    --train_data_dir=$image_dir `
    --reg_data_dir=$reg_dir `
    --output_dir=$output_dir `
    --output_name=$output_name `
    --caption_extension=".txt" `
    --prior_loss_weight=1 `
    --resolution=$resolution `
    --enable_bucket `
    --min_bucket_reso=256 `
    --max_bucket_reso=1024 `
    --train_batch_size=$train_batch_size `
    --learning_rate=$learning_rate `
    --max_train_steps=$max_training_steps `
    --mixed_precision="fp16" `
    --save_precision="fp16" `
    --use_8bit_adam `
    --xformers `
    --save_every_n_epochs=$save_every_n_epochs `
    --save_last_n_epochs=$save_last_n_epochs `
    --save_model_as=safetensors `
    --clip_skip=$clip_skip `
    --seed=$seed `
    --network_dim=$network_dim `
    --cache_latents `
    --lr_scheduler=$scheduler `
    --max_token_length=$max_token_length `
    # --save_state `
    # --resume="save_state_location" `
    # --color_aug `
    # --flip_aug `
    # --gradient_checkpointing | пока не поддерживается `
    # --gradient_accumulation_steps=1 `
    # --shuffle_caption `
    # --keep_tokens=1 `
    # --vae="X:\VAE.vae.pt" `

    deactivate
    cd $script_origin
}

# 12.01.23 by anon