#!/usr/bin/env bash
#
# setup_kontext.sh — Descarga FLUX.1-Kontext-dev (edición de imágenes por texto)
# Reutiliza el ComfyUI, el venv y el login de Hugging Face que ya tienes de setup.sh.
#
# Uso:
#   cd ~/Desktop/Flux2
#   chmod +x setup_kontext.sh
#   ./setup_kontext.sh
#
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
COMFY_DIR="$SCRIPT_DIR/ComfyUI"

if [[ ! -d "$COMFY_DIR/venv" ]]; then
  echo "ERROR: no encuentro ComfyUI/venv. Ejecuta primero ./setup.sh"
  exit 1
fi

cd "$COMFY_DIR"
# shellcheck disable=SC1091
source venv/bin/activate

echo "--------------------------------------------------------------------"
echo " FLUX.1-Kontext-dev también está protegido (gated). Necesitas aceptar"
echo " su licencia (es un repo distinto al de FLUX.1-dev):"
echo "   https://huggingface.co/black-forest-labs/FLUX.1-Kontext-dev"
echo " Pulsa 'Agree and access repository' con tu sesión iniciada."
echo "--------------------------------------------------------------------"

# Login solo si hiciera falta
if ! hf auth whoami >/dev/null 2>&1; then
  echo "Inicia sesión en Hugging Face (pega tu token Read cuando lo pida):"
  hf auth login
fi

echo "Descargando el modelo Kontext (~24 GB, paciencia)..."
hf download black-forest-labs/FLUX.1-Kontext-dev flux1-kontext-dev.safetensors \
  --local-dir "$COMFY_DIR/models/diffusion_models"

# Copiar el workflow de 2 imágenes a ComfyUI
if [[ -f "$SCRIPT_DIR/flux_kontext_2img_workflow.json" ]]; then
  mkdir -p "$COMFY_DIR/user/default/workflows"
  cp "$SCRIPT_DIR/flux_kontext_2img_workflow.json" "$COMFY_DIR/user/default/workflows/" 2>/dev/null || true
fi

echo ""
echo "============================================"
echo "  ¡Kontext listo!"
echo "============================================"
echo "Arranca ComfyUI con:   ./run.sh"
echo "y carga el workflow:   flux_kontext_2img_workflow.json"
echo "============================================"