#!/usr/bin/env bash
# Install Lens-Turbo dependencies into the ComfyUI venv.
# Covers v2 README:
#   - 模式 A: transformers / huggingface_hub 升级
#   - 模式 D: Python 3.13 老 wheel 编不过 -> --override
set -euo pipefail

VENV_PY="${VENV_PY:-/mnt/modelall/ComfyUI/.venv/bin/python}"
TMP="${TMPDIR:-/mnt/modelall/tmp}"
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")"/.. && pwd)"

mkdir -p "$TMP"

echo "[install_deps] freezing current env ..."
"$VENV_PY" -m uv pip freeze > "$TMP/freeze_before_lens_$(date +%Y%m%d_%H%M).txt" || true

OVERRIDES="$TMP/lens_overrides.txt"
cat > "$OVERRIDES" <<EOF
# Python 3.13 wheel compatibility (v2 README §3 模式 D)
sentencepiece>=0.2.0
pyyaml>=6.0.2

# Lens upstream needs reasonably new transformers/diffusers
transformers>=4.55.0
huggingface_hub>=0.26
diffusers>=0.31.0
EOF

echo "[install_deps] installing requirements with override ..."

if [ -f "$HERE/lens_source/requirements.txt" ]; then
  REQ="$HERE/lens_source/requirements.txt"
elif [ -f "$HERE/requirements.txt" ]; then
  REQ="$HERE/requirements.txt"
else
  echo "ERROR: no requirements.txt found in $HERE/lens_source or $HERE" >&2
  exit 1
fi

echo "[install_deps] using requirements file: $REQ"
echo "[install_deps] overrides: $OVERRIDES"

uv pip install --python "$VENV_PY" \
    -r "$REQ" \
    --override "$OVERRIDES" \
    --no-cache-dir

# Always make sure the bare essentials are present even if the requirements
# file is missing them.
uv pip install --python "$VENV_PY" \
    "transformers>=4.55.0" \
    "diffusers>=0.31.0" \
    "huggingface_hub>=0.26" \
    "accelerate>=0.34" \
    "safetensors>=0.4" \
    "sentencepiece>=0.2.0" \
    "pyyaml>=6.0.2" \
    "pillow>=10.0"

echo "[install_deps] uninstalling broken `kernels` lib (0.14.0 has trailing-space user_agent bug; RTX 3090 can't use MXFP4 native anyway — transformers will gracefully dequantize to bf16)"
uv pip uninstall --python "$VENV_PY" kernels || true

echo "[install_deps] verifying import ..."
"$VENV_PY" -c "
import sys
sys.path.insert(0, '$HERE/lens_source')
import lens
from lens import LensPipeline
print('[OK] lens version:', getattr(lens, '__version__', 'n/a'))
print('[OK] LensPipeline:', LensPipeline)
"

echo "[install_deps] done. Restart ComfyUI."