train functions and templates

training_scripts/t35-template

@@ -0,0 +1,149 @@
+#!/usr/bin/env zsh
+set -e -o pipefail
+
+source "$HOME/toolkit/zsh/train_functions.zsh"
+
+local name_default="by_momowaaai-v1s6000"
+NAME="${NAME:-"$name_default"}"
+
+# Extract steps from name
+STEPS=$(extract_steps_from_name "$NAME" "4096")
+# remove the suffix (eg. -v1s6000) from NAME
+NAME=$(echo "$NAME" | awk -F'-' '{print $1}')
+
+local training_dir_default="${HOME}/datasets/${NAME}"
+TRAINING_DIR=${TRAINING_DIR:-"$training_dir_default"}
+
+OUTPUT_DIR="${HOME}/output_dir"
+
+SD_SCRIPT="${SD_SCRIPT:-sd3_train_network.py}"
+SD_REPO="${SD_REPO:-$HOME/source/repos/sd-scripts-sd3}"
+
+echo "\e[35mNAME\e[0m: $NAME, \e[35mTRAINING_DIR\e[0m: $TRAINING_DIR, \e[35mOUTPUT_DIR\e[0m: $OUTPUT_DIR"
+
+[[ ! -d "$TRAINING_DIR" ]] && echo "ERROR: Training directory not found" && exit 1
+if [[ -d "$OUTPUT_DIR/$NAME" ]]; then
+    echo "ERROR: Output directory already exists: $OUTPUT_DIR/$NAME"
+    exit 1
+fi
+
+# Initialize conda and activate environment
+setup_conda_env "sdscripts"
+
+# Store the commits hashes for libraries, copy the script to the output directory
+LYCORIS_REPO=$(get_lycoris_repo) # Python package editable path
+store_commits_hashes "$SD_REPO" "$LYCORIS_REPO"
+
+# alpha=1 @ dim=16 is the same lr than alpha=4 @ dim=256
+# --min_snr_gamma=1
+args=(
+    # ⚠️  TODO: Benchmark...
+    #--debiased_estimation_loss
+    # ⚠️  TODO: What does this do? Does it even work?
+    #--max_token_length=225
+    --clip_g_dropout_rate=0.0
+    --t5_dropout_rate=0.0
+    --enable_scaled_pos_embed
+    # Keep Tokens
+    --keep_tokens=1
+    --keep_tokens_separator="|||"
+    # Model
+    --pretrained_model_name_or_path=/home/kade/ComfyUI/models/checkpoints/sd3.5_large.safetensors
+    --clip_l=/home/kade/ComfyUI/models/clip/clip_l.safetensors
+    --clip_g=/home/kade/ComfyUI/models/clip/clip_g.safetensors
+    --t5xxl=/home/kade/ComfyUI/models/clip/t5xxl_fp16.safetensors
+    # Output, logging
+    --output_dir="$OUTPUT_DIR/$NAME"
+    --output_name="$NAME"
+    --log_prefix="$NAME-"
+    --log_with=tensorboard
+    --logging_dir="$OUTPUT_DIR/logs"
+    --seed=1728871242
+    --fp8_base
+    # Dataset
+    --dataset_config="$TRAINING_DIR/config.toml"
+    #--train_data_dir="$TRAINING_DIR"
+    --dataset_repeats=1
+    --resolution="1024,1024"
+    --enable_bucket
+    --bucket_reso_steps=64
+    --min_bucket_reso=128
+    --max_bucket_reso=2048
+    --flip_aug
+    --shuffle_caption
+    --cache_latents
+    --cache_latents_to_disk
+    --max_data_loader_n_workers=8
+    --persistent_data_loader_workers
+    # Network config
+    --network_dim=64
+    # ⚠️  TODO: Plot
+    --network_alpha=0.0625
+    --network_module="lycoris.kohya"
+    --network_args
+               "preset=full"
+               "decompose_both=False"
+               "rank_dropout=0"
+               "module_dropout=0"
+               "use_tucker=True"
+               "use_scalar=False"
+               "rank_dropout_scale=False"
+               "algo=lokr"
+               "bypass_mode=False"
+               "factor=16"
+               "dora_wd=True"
+               "train_norm=False"
+    --network_dropout=0
+    # Optimizer config
+    --optimizer_type=ClybW
+    --train_batch_size=14
+    #--gradient_accumulation_steps=1
+    --max_grad_norm=1
+    --gradient_checkpointing
+    #--scale_weight_norms=1
+    # LR Scheduling
+    --max_train_steps=$STEPS
+    #--lr_warmup_steps=100
+    # NOTE: 0.0004 if its anything like FLUX..
+    --learning_rate=0.0005
+    --unet_lr=0.0002
+    --text_encoder_lr=0.0001
+    --lr_scheduler="cosine"
+    --lr_scheduler_args="num_cycles=0.375"
+    # Noise
+    --multires_noise_iterations=12
+    --multires_noise_discount=0.4
+    #--min_snr_gamma=1
+    # Optimization, details
+    --no_half_vae
+    --sdpa
+    --mixed_precision="bf16"
+    # Saving
+    --save_model_as="safetensors"
+    --save_precision="fp16"
+    --save_every_n_steps=100
+    # Saving States
+    #--save_state
+    # Either resume from a saved state
+    #--resume="$OUTPUT_DIR/wolflink-vfucks400" # Resume from saved state
+    #--skip_until_initial_step
+    # Or from a checkpoint
+    #--network_weights="$OUTPUT_DIR/wolflink-vfucks400/wolflink-vfucks400-step00000120.safetensors" # Resume from checkpoint (not needed with state, i think)
+    #--initial_step=120
+    # Sampling
+    --sample_every_n_steps=25
+    --sample_prompts="$TRAINING_DIR/sample-prompts.txt"
+    --sample_sampler="euler_a"
+    --sample_at_first
+    --caption_extension=".txt"
+)
+
+run_training_script "$SD_REPO/$SD_SCRIPT" "${args[@]}" "$@"
+
+# Delete the output directory if it is empty
+# that is if there is no samples under $OUTPUT_DIR/$NAME/samples/ and no safetensors under $OUTPUT_DIR/$NAME/
+if [[ ! -e "$OUTPUT_DIR/$NAME/samples/*" ]] && [[ ! -e "$OUTPUT_DIR/$NAME"/*.safetensors ]]; then
+    echo "No samples or model files found, deleting empty output directory"
+    rm -rf "$OUTPUT_DIR/$NAME"
+fi
+

training_scripts/tp-template

@@ -0,0 +1,132 @@
+#!/usr/bin/env zsh
+set -e -o pipefail
+
+local name_default="by_"
+NAME="${NAME:-"$name_default"}"
+
+local training_dir_default="${HOME}/datasets/${NAME}"
+TRAINING_DIR=${TRAINING_DIR:-"$training_dir_default"}
+
+OUTPUT_DIR="${HOME}/output_dir"
+
+SD_SCRIPT="${SD_SCRIPT:-sdxl_train_network.py}"
+SD_REPO="${SD_REPO:-$HOME/source/repos/sd-scripts-sd3}"
+
+echo "\e[35mNAME\e[0m: $NAME, \e[35mTRAINING_DIR\e[0m: $TRAINING_DIR, \e[35mOUTPUT_DIR\e[0m: $OUTPUT_DIR"
+
+[[ ! -d "$TRAINING_DIR" ]] && echo "ERROR: Training directory not found" && exit 1
+if [[ -d "$OUTPUT_DIR/$NAME" ]]; then
+    echo "ERROR: Output directory already exists: $OUTPUT_DIR/$NAME"
+    exit 1
+fi
+
+source "$HOME/toolkit/zsh/train_functions.zsh"
+
+# Extract steps from name
+STEPS=$(extract_steps_from_name "$NAME" "4096")
+
+# Initialize conda and activate environment
+setup_conda_env "sdscripts"
+
+# Store the commits hashes for libraries, copy the script to the output directory
+LYCORIS_REPO=$(get_lycoris_repo) # Python package editable path
+store_commits_hashes "$SD_REPO" "$LYCORIS_REPO"
+
+# alpha=1 @ dim=16 is the same lr than alpha=4 @ dim=256
+# --min_snr_gamma=1
+args=(
+    # ⚠️  TODO: Benchmark...
+    --debiased_estimation_loss
+    # ⚠️  TODO: What does this do? Does it even work?
+    --max_token_length=225
+    # Keep Tokens
+    --keep_tokens=1
+    --keep_tokens_separator="|||"
+    # Model
+    --pretrained_model_name_or_path=/home/kade/ComfyUI/models/checkpoints/ponyDiffusionV6XL_v6StartWithThisOne.safetensors
+    # Output, logging
+    --output_dir="$OUTPUT_DIR/$NAME"
+    --output_name="$NAME"
+    --log_prefix="$NAME-"
+    --log_with=tensorboard
+    --logging_dir="$OUTPUT_DIR/logs"
+    --seed=1728871242
+    # Dataset
+    --train_data_dir="$TRAINING_DIR"
+    --dataset_repeats=1
+    --resolution="1024,1024"
+    --enable_bucket
+    --bucket_reso_steps=64
+    --min_bucket_reso=256
+    --max_bucket_reso=2048
+    --flip_aug
+    --shuffle_caption
+    --cache_latents
+    --cache_latents_to_disk
+    --max_data_loader_n_workers=8
+    --persistent_data_loader_workers
+    # Network config
+    --network_dim=100000
+    # ⚠️  TODO: Plot
+    --network_alpha=64
+    --network_module="lycoris.kohya"
+    --network_args
+               "preset=full"
+               "conv_dim=100000"
+               "decompose_both=False"
+               "conv_alpha=64"
+               "rank_dropout=0"
+               "module_dropout=0"
+               "use_tucker=True"
+               "use_scalar=False"
+               "rank_dropout_scale=False"
+               "algo=lokr"
+               "bypass_mode=False"
+               "factor=16"
+               "dora_wd=True"
+               "train_norm=False"
+    --network_dropout=0
+    # Optimizer config
+    --optimizer_type=ClybW
+    --train_batch_size=14
+    #--gradient_accumulation_steps=1
+    --max_grad_norm=1
+    --gradient_checkpointing
+    #--scale_weight_norms=1
+    # LR Scheduling
+    --max_train_steps=$STEPS
+    --lr_warmup_steps=0
+    --learning_rate=0.0003
+    --unet_lr=0.0003
+    --text_encoder_lr=0.00015
+    --lr_scheduler="cosine"
+    --lr_scheduler_args="num_cycles=0.375"
+    # Noise
+    --multires_noise_iterations=12
+    --multires_noise_discount=0.4
+    #--min_snr_gamma=1
+    # Optimization, details
+    --no_half_vae
+    --sdpa
+    --mixed_precision="bf16"
+    # Saving
+    --save_model_as="safetensors"
+    --save_precision="fp16"
+    --save_every_n_steps=100
+    # Sampling
+    --sample_every_n_steps=10
+    --sample_prompts="$TRAINING_DIR/sample-prompts.txt"
+    --sample_sampler="euler_a"
+    --sample_at_first
+    --caption_extension=".txt"
+)
+
+run_training_script "$SD_REPO/$SD_SCRIPT" "${args[@]}" "$@"
+
+# Delete the output directory if it is empty
+# that is if there is no samples under $OUTPUT_DIR/$NAME/samples/ and no safetensors under $OUTPUT_DIR/$NAME/
+if [[ ! -e "$OUTPUT_DIR/$NAME/samples/*" ]] && [[ ! -e "$OUTPUT_DIR/$NAME"/*.safetensors ]]; then
+    echo "No samples or model files found, deleting empty output directory"
+    rm -rf "$OUTPUT_DIR/$NAME"
+fi
+

zsh/train_functions.zsh

@@ -0,0 +1,205 @@
+# Functions for sd-scripts training scripts
+
+# Executes a training script located at the specified path with the provided arguments.
+#
+# Parameters:
+# - script_path: The path to the training script to be executed.
+# - args_array: An array of arguments to be passed to the training script.
+#
+# Behavior:
+# - Changes the current directory to the directory of the script.
+# - If the DEBUG environment variable is set, it prints the working directory and the arguments.
+# - Executes the script using `wat python` and captures the exit code.
+# - Returns to the original directory before exiting.
+#
+# Returns:
+# - The exit code of the executed script.
+run_training_script() {
+    local script_path="$1"
+    local args_array=("${@:2}") # Get all arguments after the first one
+
+    # Store the current directory
+    local current_dir=$(pwd)
+
+    # Change to script directory
+    local script_dir=$(dirname "$script_path")
+    local script_name=$(basename "$script_path")
+    cd "$script_dir" || return 1
+
+    # Test if the script exists
+    [[ ! -f "$script_name" ]] && echo "\e[31mERROR\e[0m: Script not found: $script_name" && return 1
+
+    echo "Working directory: $(pwd)\nRunning $script_name arguments:"
+    for arg in "${args_array[@]}"; do
+        echo "  $arg"
+    done
+
+    if [[ -n "$DEBUG" ]]; then
+        echo "This was a dry run, exiting."
+        local exit_code=0
+    else
+        python "$(basename "$script_path")" "${args_array[@]}"
+        local exit_code=$?
+    fi
+
+    # Return to original directory
+    cd "$current_dir"
+
+    return $exit_code
+}
+
+# Sets up and activates a specified Conda environment.
+#
+# Parameters:
+# - env_name: The name of the Conda environment to activate.
+# - conda_path: (Optional) The path to the Conda installation. Defaults to $HOME/miniconda3.
+#
+# Behavior:
+# - Checks if the environment name is provided and if the Conda installation exists.
+# - Initializes Conda for the current shell session.
+# - Activates the specified Conda environment and verifies its activation.
+#
+# Returns:
+# - 0 on success, or 1 if any error occurs (e.g., missing environment name, Conda installation not found, activation failure).
+setup_conda_env() {
+    local env_name="$1"
+    [[ -z "$env_name" ]] && echo "\e[31mERROR\e[0m: Environment name required" && return 1
+
+    local conda_path="${2:-$HOME/miniconda3}"
+    [[ ! -d "$conda_path" ]] && echo "\e[31mERROR\e[0m: Conda installation not found at $conda_path" && return 1
+
+    # Initialize conda for the shell session
+    if __conda_setup="$("$conda_path/bin/conda" 'shell.zsh' 'hook' 2>/dev/null)" && eval "$__conda_setup"; then
+        unset __conda_setup
+    else
+        echo "\e[31mERROR\e[0m: Failed to initialize conda environment" && return 1
+    fi
+
+    # Activate conda environment
+    conda activate "$env_name"
+    if [ $? -ne 0 ]; then
+        echo "\e[31mERROR\e[0m: Failed to activate conda environment: $env_name"
+        return 1
+    fi
+    echo "Conda environment: $CONDA_PREFIX"
+
+    # Verify environment activation
+    if ! conda env list | grep -q "^${env_name} "; then
+        echo "ERROR: Environment $env_name not found"
+        return 1
+    fi
+}
+
+# Extracts the number of steps from a given name or uses a default value if not found.
+#
+# Parameters:
+# - name: The name from which to extract the step count.
+# - default_steps: The default number of steps to use if extraction fails.
+#
+# Behavior:
+# - Uses a regular expression to find a number at the end of the name.
+# - If no valid step count is found, it uses the default value and prints a warning.
+# - If the extracted step count is less than 1, it also uses the default value and prints a warning.
+# - Prints a confirmation message if a valid step count is found.
+#
+# Returns:
+# - The extracted or default step count.
+extract_steps_from_name() {
+    local name="$1"
+    local default_steps="$2"
+
+    # More robust step extraction with validation
+    local steps=$(echo "$name" | grep -oE '[0-9]+$')
+    if [[ -z "$steps" ]]; then
+        steps="$default_steps"
+        echo "⚠️  No step count found in NAME. Using default: \e[35m$steps\e[0m" >&2
+    elif ((steps < 1)); then
+        echo "⚠️  Invalid step count. Using default: \e[35m$default_steps\e[0m" >&2
+        steps="$default_steps"
+    fi
+
+    echo "$steps"
+}
+
+# Stores the commit hashes of specified Git repositories and copies the script to an output directory.
+#
+# Parameters:
+# - output_dir: The directory where the commit hashes will be stored.
+# - repo_path: One or more paths to Git repositories.
+#
+# Behavior:
+# - Creates the output directory if it does not exist.
+# - Copies the current script to the output directory.
+# - Iterates over each repository path, checking if it is a valid Git repository.
+# - Retrieves the current commit SHA for each repository and writes it to an output file.
+# - Generates a SHA-1 hash of the script and appends it to the output file.
+#
+# Returns:
+# - 0 on success, or 1 if any error occurs during the process (e.g., Git command failure, not a Git repository).
+store_commits_hashes() {
+    # Construct the output directory path
+    local output_dir="$OUTPUT_DIR/$NAME"
+    # Define the path for the output file
+    local output_file="$output_dir/repos.git"
+    # If DEBUG is set, print the output directory path
+    [[ -n "$DEBUG" ]] && echo "Output directory: $output_dir"
+    # Create the output directory if it doesn't exist
+    [[ ! -d "$output_dir" ]] && mkdir -p "$output_dir"
+    # Create or truncate the output file
+    : >"$output_file"
+
+    local summary=""
+    local res=0
+
+    for repo_path in "$@"; do
+        local repo_name=$(basename "$repo_path")
+        if [[ -d "$repo_path/.git" ]]; then
+            if local commit_sha=$(git -C "$repo_path" rev-parse HEAD 2>/dev/null); then
+                # Get the checked-out branch
+                if local branch_name=$(git -C "$repo_path" rev-parse --abbrev-ref HEAD 2>/dev/null); then
+                    echo "$repo_path: ($branch_name) $commit_sha" >>"$output_file"
+                    summary+="✓ $repo_name: $commit_sha ($branch_name) $repo_path\n"
+                else
+                    echo "$repo_path: $commit_sha (Failed to get branch)" >>"$output_file"
+                    summary+="⚠️  $repo_name: $commit_sha (Failed to get branch) $repo_path\n"
+                    res=1
+                fi
+            else
+                echo "$repo_path: Git command failed" >>"$output_file"
+                summary+="⚠️  $repo_name: Git command failed $repo_path\n"
+                res=1
+            fi
+        else
+            echo "$repo_path: Not a git repository" >>"$output_file"
+            summary+="⚠️  $repo_name: Not a git repository $repo_path\n"
+            res=1
+        fi
+    done
+
+    # Copy the script to the output directory
+    local script_path=$(readlink -f "$ZSH_SCRIPT")
+    cp "$script_path" "$output_dir/$(basename "$script_path")"
+    [[ -n "$DEBUG" ]] && echo "Copied $script_path to $output_dir"
+
+    # Add script hash with error handling
+    if local script_sha=$(sha1sum "$script_path" | cut -f1 -d' '); then
+        echo "$script_path: $script_sha" >>"$output_file"
+        summary+="✓ $ZSH_SCRIPT: $script_sha\n"
+    else
+        echo "$script_path: Failed to generate SHA-1" >>"$output_file"
+        summary+="⚠️  Failed to generate script SHA-1\n"
+        res=1
+    fi
+
+    echo -e "$summary"
+    return $res
+}
+
+get_lycoris_repo() {
+    python -c """
+import importlib.util
+import pathlib
+spec = importlib.util.find_spec('lycoris')
+print(pathlib.Path(spec.origin).parent.parent)
+    """
+}