#!/usr/bin/env bash
# runpod-klein-models.sh  —  modelos FLUX.2 Klein 9B (generacion de imagen) para el Pod LTX
# Se corre en la terminal de Jupyter del Pod:
#   curl -s https://dl.ramsesvii.com/runpod-klein-models.sh | bash
# NO reinicia ComfyUI: Klein usa nodos del core; los modelos aparecen al refrescar la UI.
# Descarga ~17 GB en tmux (sesion 'dlklein'); sobrevive si se cierra Jupyter.
set +e
COMFY="/workspace/ComfyUI"
log(){ echo -e "\n\033[1;36m>> $*\033[0m"; }

cat > /workspace/_dl_klein.sh <<'WORKER'
#!/usr/bin/env bash
set +e
COMFY="/workspace/ComfyUI"
export HF_HUB_ENABLE_HF_TRANSFER=1
pip install -q --upgrade "huggingface_hub[hf_transfer]" 2>/dev/null
HFDL(){ # repo  archivo_remoto  carpeta_destino  [nombre_local]
  local repo="$1" rfile="$2" dest="$3" name="${4:-$(basename "$2")}"
  mkdir -p "$dest"; local out="$dest/$name"
  if [ -f "$out" ]; then echo "   ya existe: $name"; return; fi
  echo ">> bajando $name  ($repo)"
  hf download "$repo" "$rfile" --local-dir "$dest/.hfdl" && mv "$dest/.hfdl/$rfile" "$out" && rm -rf "$dest/.hfdl"
}
echo "############ DESCARGA KLEIN 9B $(date -u) ############"
df_av(){ df -h /workspace | tail -1; }
echo "Espacio /workspace antes:"; df_av

# DiT distilled fp8 (~8.8 GB) -> diffusion_models
HFDL silveroxides/FLUX.2-dev-fp8_scaled flux-2-klein-9b-fp8mixed.safetensors "$COMFY/models/diffusion_models"
# text encoder qwen3-8b fp8 (~8.1 GB) -> text_encoders
HFDL Comfy-Org/vae-text-encorder-for-flux-klein-9b split_files/text_encoders/qwen_3_8b_fp8mixed.safetensors "$COMFY/models/text_encoders"
# VAE flux2 (~321 MB) -> vae
HFDL Comfy-Org/vae-text-encorder-for-flux-klein-9b split_files/vae/flux2-vae.safetensors "$COMFY/models/vae"

echo; echo "Espacio /workspace despues:"; df_av
echo; echo "==== KLEIN: archivos presentes ===="
for f in \
  "$COMFY/models/diffusion_models/flux-2-klein-9b-fp8mixed.safetensors" \
  "$COMFY/models/text_encoders/qwen_3_8b_fp8mixed.safetensors" \
  "$COMFY/models/vae/flux2-vae.safetensors"; do
  if [ -f "$f" ]; then printf "  OK  %6.2f GB  %s\n" "$(stat -c%s "$f" | awk '{print $1/1e9}')" "$f"
  else printf "  FALTA  %s\n" "$f"; fi
done
echo; echo ">> LISTO. En ComfyUI: boton de refrescar (o tecla R) para que aparezcan en los loaders."
echo ">> Avisale a Claude para verificar por la API."
WORKER
chmod +x /workspace/_dl_klein.sh

log "Lanzando descarga de Klein (~17 GB) en tmux 'dlklein'"
tmux kill-session -t dlklein 2>/dev/null
tmux new-session -d -s dlklein "bash /workspace/_dl_klein.sh 2>&1 | tee /workspace/dl-klein.log"

echo
echo "=================================================================="
echo "  Descarga de ~17 GB (Klein 9B) corriendo en tmux 'dlklein'."
echo "  Monitorear:   tail -f /workspace/dl-klein.log"
echo "  o:            tmux attach -t dlklein   (Ctrl-b luego d para salir)"
echo "  NO reinicia ComfyUI. Al terminar: refresca la UI y avisa a Claude."
echo "=================================================================="