$lr_warmup_ratio = 0.05
$train_batch_size = 10
$num_epochs = 1
$save_every_n_epochs=1
$scheduler="cosine_with_restarts"
$network_dim=128
$text_encoder_lr=1.5e-5
$unet_lr=1.5e-4

# Config End

Clear-Host

Write-Output "Measuring folders:"
$total = 0
$folders = Get-ChildItem -Path $image_dir -Directory
$folders | ForEach-Object {
    $parts = $_.Name.Split("_")
    if(($parts | Measure-Object | ForEach-Object{$_.Count}) -ne 2)
    {
        Return
    }
    $repeats = [int]$parts[0]
    $imgs = Get-ChildItem $_.FullName -Recurse -File -Include *.png, *.bmp, *.gif, *.jpg,*.jpeg, *.webp | Measure-Object | ForEach-Object{$_.Count}
    $img_repeats = ($repeats * $imgs)
    Write-Output "`t$($parts[1]): $repeats repeats * $imgs images = $($img_repeats)"
    $total += $img_repeats
}
Write-Output "Total images with repeats: $total"
$mts = [int]($total / $train_batch_size * $num_epochs)
$lr_warmup_steps = [math]::Round($mts * $lr_warmup_ratio)
Write-Output "Max training steps $total / $train_batch_size * $num_epochs = $mts"

.\venv\Scripts\activate

accelerate launch --num_cpu_threads_per_process 16 train_network.py `
    --cache_latents `
    --enable_bucket `
    --use_8bit_adam `
    --xformers `
    --shuffle_caption `
    --pretrained_model_name_or_path=$ckpt `
    --train_data_dir=$image_dir `
    --reg_data_dir=$reg_dir `
    --resolution=512,512 `
    --output_dir=$output `
    --train_batch_size=$train_batch_size `
    --lr_scheduler=$scheduler `
    --lr_warmup_steps=$lr_warmup_steps `
    --max_train_steps=$mts `
    --use_8bit_adam `
    --xformers `
    --mixed_precision=fp16 `
    --save_every_n_epochs=$save_every_n_epochs `
    --seed=31337 `
    --save_precision=fp16 `
    --logging_dir="" `
    --caption_extension=.txt `
    --save_model_as=safetensors `
    --network_module=networks.lora `
    --text_encoder_lr=$text_encoder_lr `
    --unet_lr=$unet_lr `
    --network_dim=$network_dim `
    --clip_skip=2 `