# https://github.com/67372a/sd-scripts
# Config Start
# txt

$ckpt = "A:\abcdefg\noobaiXLNAIXL_vPred10Version.safetensors"; # base checkpoint to finetune

$img_dir = "A:\abcdefg\v1";
# dataset_folder
#	10_A, images in A will repeat 10x
#	2_B, images in B will repeat 2x

$reg_dir = "A:\abcdefg\reg_empty"; #optional, just point this to an empty folder if you don't care.

$output = "A:\noobvpredv10\v1_out"; # safetensors output folder

$graph_dir = "A:\noobvpredv10\v1_out\graph"; # edm2 graph output folder

$model_name="style_noobvpredv10_abcdefg_v1"

$learning_rate = 3e-5
$lr_warmup_ratio = 0.25
$train_batch_size = 2
$num_epochs = 16
$save_every_n_epochs=1
$scheduler="library.LoraEasyCustomOptimizer.CosineAnnealingWarmRestarts.CosineAnnealingWarmRestarts"
$T_0_Multiplier = 5
$clip_threshold = 1.25
# Config End

Clear-Host

Write-Output "Measuring folders:"
$total = 0
$folders = Get-ChildItem -Path $img_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)
Write-Output "Max training steps $total / $train_batch_size * $num_epochs = $mts"
$lr_warmup_steps = [math]::Round($mts * $lr_warmup_ratio)
Write-Output "Warmup steps: $lr_warmup_steps, Batch size: $train_batch_size, Epochs: $num_epochs"
$first_cycle_max_steps = [math]::Round($mts / $num_epochs) * $T_0_Multiplier
Write-Output "Warmup steps: $lr_warmup_steps, First cycle max steps: $first_cycle_max_steps, Batch size: $train_batch_size, Epochs: $num_epochs"


.\venv\Scripts\activate

accelerate launch --num_cpu_threads_per_process 12 sdxl_train.py `
    --pretrained_model_name_or_path=$ckpt `
    --train_data_dir=$img_dir `
    --output_dir=$output `
    --caption_extension=".txt" `
    --shuffle_caption `
    --max_grad_norm=1.0 `
    --output_name=$model_name `
    --resolution=1024 `
    --enable_bucket `
    --min_bucket_reso=512 `
    --max_bucket_reso=2048 `
    --train_batch_size=$train_batch_size `
    --learning_rate=$learning_rate `
    --max_train_steps=$mts `
    --mixed_precision="bf16" `
    --save_precision="fp16" `
    --optimizer_type=adafactor `
    --optimizer_args scale_parameter=False relative_step=False warmup_init=False clip_threshold=$clip_threshold `
    --sdpa `
    --save_every_n_epochs=$save_every_n_epochs `
    --save_model_as=safetensors `
    --seed=22 `
    --fused_backward_pass `
    --v_parameterization `
    --zero_terminal_snr `
    --sangoi_loss_modifier `
    --sangoi_loss_modifier_max_snr=15 `
    --immiscible_noise=4096 `
    --loss_type="log_cosh" `
    --edm2_loss_weighting `
    --edm2_loss_weighting_laplace `
    --edm2_loss_weighting_num_channels=128 `
    --edm2_loss_weighting_optimizer="torch.optim.AdamW" `
    --edm2_loss_weighting_optimizer_lr=1.5e-2 `
    --edm2_loss_weighting_lr_scheduler `
    --edm2_loss_weighting_lr_scheduler_warmup_percent=0.10 `
    --edm2_loss_weighting_generate_graph `
    --edm2_loss_weighting_generate_graph_output_dir=$graph_dir `
    --edm2_loss_weighting_max_grad_norm=2.0 `
    --edm2_loss_weighting_use_float64 `
    --loss_related_use_float64 `
    --disable_cuda_reduced_precision_operations `
    --block_lr=0.0,3.00e-5,3.00e-5,3.00e-5,3.00e-5,3.00e-5,3.00e-5,3.00e-5,3.00e-5,3.00e-5,2.95e-5,2.95e-5,2.95e-5,2.90e-5,2.85e-5,2.80e-5,2.75e-5,2.70e-5,8e-6,2.60e-5,8e-6,2.30e-5,2.45e-5 `
    <# --color_aug #> `
    <# --flip_aug #> `
    --logging_dir="$output\logs" `
    --max_token_length=225 `
    --cache_latents `
    --cache_latents_to_disk `
    --no_half_vae `
    --gradient_checkpointing `
    --persistent_data_loader_workers `
    --lr_scheduler_type=$scheduler `
    --lr_scheduler_num_cycles=2 `
    --lr_scheduler_args "first_cycle_max_steps=$first_cycle_max_steps" "cycle_multiplier=1" "gamma=0.710" "warmup_steps=$lr_warmup_steps" `
    #--vae=$vae `
    #--lr_warmup_steps=$lr_warmup_steps `