#!/bin/bash
# Install GFPGAN face enhancement into the ComfyUI venv.
# Applies two patches for Python 3.13 + newer torchvision compatibility.
# Usage: bash tour-comfy/install_gfpgan.sh
set -e

SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" && pwd )"
source "$SCRIPT_DIR/env.sh"
source "$VENV/bin/activate"

PYTHON="$VENV/bin/python"
PIP="$VENV/bin/pip"

echo "[gfpgan] Step 1 — install basicsr (with Python 3.13 patch) ..."
TMPDIR=$(mktemp -d)
curl -sL "https://pypi.io/packages/source/b/basicsr/basicsr-1.4.2.tar.gz" -o "$TMPDIR/basicsr-1.4.2.tar.gz"
tar -xzf "$TMPDIR/basicsr-1.4.2.tar.gz" -C "$TMPDIR"

# Patch 1: fix get_version() — exec() doesn't update locals() in Python 3
"$PYTHON" - <<'PYPATCH'
import sys, re
setup = sys.argv[1]
with open(setup) as f:
    content = f.read()
old = "def get_version():\n    with open(version_file, 'r') as f:\n        exec(compile(f.read(), version_file, 'exec'))\n    return locals()['__version__']"
new = "def get_version():\n    if not os.path.exists(version_file):\n        write_version_py()\n    globs = {}\n    with open(version_file, 'r') as f:\n        exec(compile(f.read(), version_file, 'exec'), globs)\n    return globs['__version__']"
if old in content:
    content = content.replace(old, new)
    with open(setup, 'w') as f:
        f.write(content)
    print('  Patched setup.py get_version()')
else:
    print('  setup.py pattern not found, skipping patch')
PYPATCH
"$TMPDIR/basicsr-1.4.2/setup.py" -- "$TMPDIR/basicsr-1.4.2/setup.py" 2>/dev/null || true
"$PYTHON" - "$TMPDIR/basicsr-1.4.2/setup.py" <<'PYPATCH'
import sys, re, os
setup = sys.argv[1]
with open(setup) as f:
    content = f.read()
old = "def get_version():\n    with open(version_file, 'r') as f:\n        exec(compile(f.read(), version_file, 'exec'))\n    return locals()['__version__']"
new = "def get_version():\n    if not os.path.exists(version_file):\n        write_version_py()\n    globs = {}\n    with open(version_file, 'r') as f:\n        exec(compile(f.read(), version_file, 'exec'), globs)\n    return globs['__version__']"
if old in content:
    content = content.replace(old, new)
    with open(setup, 'w') as f:
        f.write(content)
    print('  Patched setup.py get_version()')
else:
    print('  setup.py already patched or pattern changed, skipping')
PYPATCH

"$PIP" install "$TMPDIR/basicsr-1.4.2/" --no-build-isolation --no-deps -q
rm -rf "$TMPDIR"

echo "[gfpgan] Step 2 — install facexlib and gfpgan ..."
"$PIP" install facexlib gfpgan -q

# Patch 2: fix torchvision functional_tensor import (removed in newer torchvision)
DEGR_PY="$VENV/lib/python3.13/site-packages/basicsr/data/degradations.py"
if [ -f "$DEGR_PY" ]; then
    if grep -q "functional_tensor" "$DEGR_PY"; then
        sed -i 's/from torchvision.transforms.functional_tensor import rgb_to_grayscale/from torchvision.transforms.functional import rgb_to_grayscale/' "$DEGR_PY"
        echo "[gfpgan] Patched degradations.py functional_tensor import"
    fi
fi

# Pre-download the model
MODEL_DIR="$HOME/.gfpgan/weights"
MODEL_PATH="$MODEL_DIR/GFPGANv1.4.pth"
mkdir -p "$MODEL_DIR"
if [ ! -f "$MODEL_PATH" ]; then
    echo "[gfpgan] Downloading GFPGANv1.4.pth (~333 MB) ..."
    wget -q --show-progress \
        "https://github.com/TencentARC/GFPGAN/releases/download/v1.3.4/GFPGANv1.4.pth" \
        -O "$MODEL_PATH.tmp"
    mv "$MODEL_PATH.tmp" "$MODEL_PATH"
    echo "[gfpgan] Model saved to $MODEL_PATH"
else
    echo "[gfpgan] Model already present: $MODEL_PATH"
fi

echo ""
echo "[gfpgan] Done. Restart the API (start_api.sh) to enable face enhancement."