#!/bin/bash
set -e 

echo "Installing system dependencies..."
apt-get update
apt-get install -y git

echo "Starting SimpleTuner environment setup..."
git clone https://github.com/bghira/SimpleTuner.git /workspace/SimpleTuner
cd /workspace/SimpleTuner

pip install uv

echo "Setting up Python 3.12 environment..."
uv python install 3.12
uv venv --python 3.12 /workspace/venv
source /workspace/venv/bin/activate

echo "Installing dependencies..."
uv pip install -e .
uv pip install peft bitsandbytes wandb huggingface_hub

echo "Migrating configurations..."
cp /data/config.env /workspace/SimpleTuner/config.env
cp /data/dataset.json /workspace/SimpleTuner/dataset.json

source config.env

echo "Authenticating Hugging Face natively..."
mkdir -p /root/.cache/huggingface
echo -n "$HF_TOKEN" > /root/.cache/huggingface/token

echo "Launching accelerate pipeline..."
accelerate launch \
  --num_processes=1 \
  --mixed_precision=$MIXED_PRECISION \
  train.py \
  --model_type="lora" \
  --model_family="flux" \
  --pretrained_model_name_or_path="$MODEL_NAME" \
  --data_backend_config="dataset.json" \
  --output_dir="$OUTPUT_DIR" \
  --train_batch_size=$TRAIN_BATCH_SIZE \
  --gradient_accumulation_steps=$GRADIENT_ACCUMULATION_STEPS \
  --learning_rate=$LEARNING_RATE \
  --lr_scheduler="$LR_SCHEDULER" \
  --lr_warmup_steps=$LR_WARMUP_STEPS \
  --max_train_steps=$MAX_TRAIN_STEPS \
  --checkpointing_steps=$CHECKPOINTING_STEPS \
  --mixed_precision="$MIXED_PRECISION" \
  --optimizer="$OPTIMIZER" \
  --gradient_checkpointing=true \
  --lora_rank=$LORA_RANK \
  --lora_alpha=$LORA_ALPHA \
  --validation_prompt="$VALIDATION_PROMPT" \
  --validation_steps=$VALIDATION_STEPS \
  --validation_resolution=$RESOLUTION \
  --push_to_hub \
  --hub_model_id="ctan-dev/flux-klein-output"