#!/bin/bash
# Flux LoRA Training Launcher Script
# Usage: ./start_training.sh [config_file]

source /workspace/venv/bin/activate

cd /workspace/sd-scripts

CONFIG_FILE="${1:-/workspace/flux-project/configs/train_config.toml}"

if [ ! -f "$CONFIG_FILE" ]; then
    echo "Error: Config file not found: $CONFIG_FILE"
    echo "Usage: $0 [path/to/config.toml]"
    exit 1
fi

echo "=============================================="
echo "Starting Flux LoRA Training"
echo "=============================================="
echo "Config: $CONFIG_FILE"
echo ""

# Display key config info
echo "Configuration Summary:"
grep -E "^(pretrained_model|train_data_dir|output_name|max_train_epochs|learning_rate|network_dim)" "$CONFIG_FILE" 2>/dev/null || true
echo ""

# Start training with accelerate
accelerate launch \
    --mixed_precision bf16 \
    --num_cpu_threads_per_process 4 \
    flux_train_network.py \
    --config_file "$CONFIG_FILE"

echo ""
echo "Training complete!"
echo "Output saved to: /workspace/flux-project/output/"