#!/usr/bin/env bash
#
# setup_krea.sh — Descarga FLUX.1-Krea-dev (fp8 scaled) para edición fotorrealista.
# Repo abierto de Comfy-Org: NO necesita aceptar licencia ni token.
# Reutiliza el ComfyUI y el venv de setup.sh.
#
# Uso:
#   cd ~/Desktop/Flux2
#   chmod +x setup_krea.sh
#   ./setup_krea.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 " IMPORTANTE (Mac): se usa la versión bf16 completa, porque la fp8"
echo " NO funciona en la GPU de Apple (MPS no soporta el tipo Float8)."
echo " Acepta antes la licencia de Krea (repo gated, distinto al de dev):"
echo "   https://huggingface.co/black-forest-labs/FLUX.1-Krea-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 FLUX.1-Krea-dev (bf16 completo, ~24 GB; paciencia)..."
hf download black-forest-labs/FLUX.1-Krea-dev flux1-krea-dev.safetensors \
  --local-dir "$COMFY_DIR/models/diffusion_models"

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

echo ""
echo "============================================"
echo "  ¡Krea listo!"
echo "============================================"
echo "Arranca ComfyUI:   ./run.sh"
echo "y carga:           flux_krea_workflow.json"
echo "============================================"