#!/bin/bash
set -e

# Make pip downloads more resilient: the large CUDA/torch wheels (hundreds of MB)
# occasionally get cut off mid-stream on the build network (IncompleteRead),
# which fails the whole startup. Retry harder and allow a longer read timeout.
export PIP_DEFAULT_TIMEOUT=120
export PIP_RETRIES=10

# --- Build prerequisites -----------------------------------------------------
# setuptools 82 (released 2026-02-08) removed the legacy `pkg_resources` module,
# which breaks the source builds of mmcv / mmgen (their setup.py still imports
# it). Keep a pre-82 setuptools, plus the tools needed to build native exts.
pip install "setuptools<82" wheel ninja

# Install torch first. ZeroGPU now runs on NVIDIA RTX PRO 6000 Blackwell
# (compute capability sm_120), which is ONLY supported by CUDA 12.8+ builds of
# PyTorch -- the default cu126 wheels raise "no kernel image is available for
# execution on the device". Use the cu128 build (torch 2.8.x is the minimum
# version supported by HF ZeroGPU). torch/torchvision/triton/nvidia-* are
# intentionally NOT pinned in my_req.txt so this install owns the CUDA stack,
# and mmcv/mmgen below compile against this exact torch.
pip install --index-url https://download.pytorch.org/whl/cu128 \
  torch==2.8.0 torchvision==0.23.0

# --- piFlow (lakonlab) -------------------------------------------------------
# Install ONLY the `lakonlab` package, not piFlow's own dependencies: piFlow's
# requirements.txt pins (gradio 4.x, transformers 4.54, diffusers 0.35, ...)
# conflict with this Space's curated my_req.txt below. Dependencies are provided
# by my_req.txt instead.
#
# Also drop the `-e` (editable) flag: an editable VCS install clones into
# ./src/lakonlab and, when that path already exists, pip falls back to an
# interactive prompt ("(i)gnore/(w)ipe/(b)ackup") -> EOFError in the container.
pip install --no-deps \
  "lakonlab @ git+https://github.com/Lakonik/piFlow.git@b1ef16e5e305251bccdfeac2a0e3d0ef339b974a"

# --- Remaining dependencies --------------------------------------------------
# Disable build isolation so mmcv/mmgen build against the torch + setuptools<82
# installed above, instead of a fresh isolated build env that would pull the
# pkg_resources-less setuptools 82.
pip install --no-build-isolation -r my_req.txt