#!/usr/bin/env bash
# =============================================================================
# runpod-gpu-train-setup.sh — ai-toolkit (ostris) en el pod, para entrenar LoRAs
#
# Se corre DENTRO del pod (por SSH desde el VPS o desde Jupyter):
#   curl -fsSL https://dl.ramsesvii.com/runpod-gpu-train-setup.sh | bash
#
# Idempotente. Todo en /workspace (persiste): venv propio /workspace/venv-train con
# --system-site-packages (reusa el torch de la imagen; ai-toolkit y ComfyUI no
# comparten venv porque sus dependencias chocan). Datasets en /workspace/datasets,
# salidas en /workspace/loras/personajes.
# =============================================================================
set -u
W=/workspace; AIT=$W/ai-toolkit; VENV=$W/venv-train; ST=$W/status
mkdir -p "$W/datasets" "$W/loras/personajes" "$ST"
say() { echo "[$(date -u +%FT%TZ)] $*"; }

if [ ! -d "$AIT/.git" ]; then
  say "clonando ai-toolkit"
  git clone --depth 1 https://github.com/ostris/ai-toolkit "$AIT"
else
  say "ai-toolkit ya existe"
fi

PYTAG="$(python3 -c 'import sys;print(f"py{sys.version_info[0]}.{sys.version_info[1]}")')-${PYTORCH_VERSION:-torch}"
if [ ! -x "$VENV/bin/python" ] || [ "$(cat "$ST/venv-train-tag" 2>/dev/null)" != "$PYTAG" ]; then
  say "creando venv-train"
  rm -rf "$VENV"
  python3 -m venv --system-site-packages "$VENV" 2>/dev/null \
    || { python3 -m pip install -q virtualenv && python3 -m virtualenv -q --system-site-packages "$VENV"; }
  echo "$PYTAG" >"$ST/venv-train-tag"; rm -f "$ST/train-deps-ok"
fi
PIP="$VENV/bin/python -m pip"
if [ ! -f "$ST/train-deps-ok" ]; then
  say "instalando dependencias de ai-toolkit (varios minutos)"
  $PIP install -q --upgrade pip
  # requirements.txt incluye requirements_base.txt por ruta relativa: instalar desde su carpeta.
  # Ninguno de los dos fija torch (lo aporta la imagen vía --system-site-packages).
  ( cd "$AIT" && $PIP install -q -r requirements.txt ) && touch "$ST/train-deps-ok" || say "FALLÓ pip de ai-toolkit (ver arriba)"
  # con --system-site-packages, safetensors queda con metadata rota (version=None) y transformers no importa
  $PIP install -q --ignore-installed --no-deps "safetensors>=0.4.3"
fi
"$VENV/bin/python" -c "import torch, diffusers, transformers; print('torch', torch.__version__, 'cuda', torch.cuda.is_available(), '| diffusers', diffusers.__version__)"
say "listo. Entrenar: cd $AIT && $VENV/bin/python run.py /workspace/datasets/<job>/config.yaml"