#!/usr/bin/env bash
# Runs inside the RunPod container on startup. Installs deps, fetches
# the training script from HF (MYD repo is private), and runs it.
set -euxo pipefail

mkdir -p /workspace/orpheus
cd /workspace/orpheus

nvidia-smi > /workspace/orpheus/nvidia-smi.txt 2>&1 || true

# Install deps (try cache first)
pip install --upgrade pip
pip install --no-cache-dir \
  "torch==2.4.0" \
  "transformers>=4.45" \
  "accelerate>=0.34" \
  "peft>=0.13" \
  "datasets>=3.0,<4.0" \
  "huggingface_hub>=0.26" \
  "hf_transfer>=0.1.8" \
  "snac>=1.2.1" \
  "librosa>=0.10" \
  "soundfile>=0.12" \
  "jiwer>=3.0" \
  "sentencepiece" \
  "protobuf" \
  "pyarrow>=16" \
  "torchaudio" 2>&1 | tail -30

export HF_HUB_ENABLE_HF_TRANSFER=1
export HUGGINGFACE_HUB_TOKEN="${HF_TOKEN:-}"

# Fetch training script from HF (public, no gating)
SCRIPT_URL="${TRAIN_SCRIPT_URL:-https://huggingface.co/oridror/orpheus-hebrew-trainer/resolve/main/train.py}"
curl -sSL "$SCRIPT_URL" -o /workspace/orpheus/train.py
echo "Downloaded training script ($(wc -l < /workspace/orpheus/train.py) lines)"

# Run training (tee into a log for after-the-fact debugging)
python -u /workspace/orpheus/train.py 2>&1 | tee /workspace/orpheus/run.log || true

# Mark done
echo "TRAINING_FINISHED_$(date -u +%s)" > /workspace/orpheus/status.txt

# Push run.log to the progress repo for retrieval
python -c "
import os, sys
from huggingface_hub import HfApi, create_repo
token=os.environ.get('HUGGINGFACE_HUB_TOKEN')
if token:
  create_repo('oridror/orpheus-hebrew-progress', token=token, repo_type='model', exist_ok=True, private=False)
  api=HfApi(token=token)
  for name in ['run.log','status.txt','nvidia-smi.txt']:
    p=f'/workspace/orpheus/{name}'
    if os.path.exists(p):
      try:
        api.upload_file(path_or_fileobj=p, path_in_repo=name, repo_id='oridror/orpheus-hebrew-progress', repo_type='model', token=token)
      except Exception as e:
        print('upload fail',name,e)
" || true

# Keep container alive briefly for any async uploads, then exit
sleep 120