#!/usr/bin/env bash
# runpod-ltx-models.sh  —  Fase 2: custom nodes faltantes + modelos LTX-2.3 FFLF + workflow
# Se corre en la terminal de Jupyter del Pod:
#   curl -s https://dl.ramsesvii.com/runpod-ltx-models.sh | bash
# Lanza la descarga pesada en tmux (sesión 'dl'); sobrevive si se cierra Jupyter.
set +e
COMFY="/workspace/ComfyUI"
log(){ echo -e "\n\033[1;36m>> $*\033[0m"; }

# ---------- a) custom nodes faltantes (rápido, inline) ----------
log "Custom nodes faltantes (rgthree, Anything-Everywhere, Impact-Pack, Easy-Use)"
clone_node(){ local url="$1" name="$2" d="$COMFY/custom_nodes/$2"
  if [ -d "$d/.git" ]; then git -C "$d" pull --ff-only >/dev/null 2>&1
  else echo "  clonando $name"; git clone --depth 1 "$url" "$d"; fi
  [ -f "$d/requirements.txt" ] && pip install -q -r "$d/requirements.txt"; }
clone_node https://github.com/rgthree/rgthree-comfy            rgthree-comfy
clone_node https://github.com/chrisgoringe/cg-use-everywhere   cg-use-everywhere
clone_node https://github.com/ltdrdata/ComfyUI-Impact-Pack     ComfyUI-Impact-Pack
clone_node https://github.com/yolain/ComfyUI-Easy-Use          ComfyUI-Easy-Use

# ---------- b) workflow a su carpeta ----------
log "Workflow FFLF Seed Hunter v1.6 -> user/default/workflows/"
mkdir -p "$COMFY/user/default/workflows"
curl -s https://dl.ramsesvii.com/ltx23FFLF_v16.json -o "$COMFY/user/default/workflows/ltx23FFLF_SeedHunter_v16.json"

# ---------- c) lanzar descarga de modelos en tmux ----------
log "Lanzando descarga de modelos en tmux (sesión 'dl')"
cat > /workspace/_dl_models.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 MODELOS $(date -u) ############"
# transformer fp8 (lo mas pesado)
HFDL Kijai/LTX2.3_comfy diffusion_models/ltx-2.3-22b-distilled-1.1_transformer_only_fp8_scaled.safetensors "$COMFY/models/diffusion_models"
# text encoder gemma fp8 + projection
HFDL Comfy-Org/ltx-2   split_files/text_encoders/gemma_3_12B_it_fp8_scaled.safetensors "$COMFY/models/text_encoders"
HFDL Kijai/LTX2.3_comfy text_encoders/ltx-2.3_text_projection_bf16.safetensors "$COMFY/models/text_encoders"
# vaes
HFDL Kijai/LTX2.3_comfy vae/LTX23_video_vae_bf16.safetensors "$COMFY/models/vae"
HFDL Kijai/LTX2.3_comfy vae/LTX23_audio_vae_bf16.safetensors "$COMFY/models/vae"
HFDL Kijai/LTX2.3_comfy vae/taeltx2_3.safetensors            "$COMFY/models/vae"
# lora (subcarpeta LTX2 como pide el workflow)
HFDL Kijai/LTX2.3_comfy loras/LTX-2.3-OmniNFT-RL-Lora_bf16.safetensors "$COMFY/models/loras/LTX2"
# upscaler espacial x2 -> carpeta dedicada 'latent_upscale_models' (lo lee LatentUpscaleModelLoader)
HFDL Lightricks/LTX-2.3 ltx-2.3-spatial-upscaler-x2-1.1.safetensors "$COMFY/models/latent_upscale_models"

echo; echo "==== ARCHIVOS DESCARGADOS ===="
find "$COMFY/models" -type f \( -name '*.safetensors' \) -printf '%10s  %p\n' | awk '{printf "  %6.2f GB  %s\n",$1/1e9,$2}'

echo; echo ">> Reiniciando ComfyUI para tomar nodos+modelos..."
tmux kill-session -t comfy 2>/dev/null
tmux new-session -d -s comfy "cd $COMFY && python main.py --listen 0.0.0.0 --port 8188 2>&1 | tee /workspace/comfyui.log"
echo ">> LISTO. ComfyUI reiniciando. La URL del tunel sigue siendo la misma."
WORKER
chmod +x /workspace/_dl_models.sh
tmux kill-session -t dl 2>/dev/null
tmux new-session -d -s dl "bash /workspace/_dl_models.sh 2>&1 | tee /workspace/dl-models.log"

echo
echo "=================================================================="
echo "  ⏬ Descarga de ~44 GB corriendo en tmux 'dl'."
echo "  Monitorear:   tmux attach -t dl       (Ctrl-b luego d para salir)"
echo "  o:            tail -f /workspace/dl-models.log"
echo "  Cuando termine, reinicia ComfyUI solo. Avisá a Claude."
echo "=================================================================="