#!/usr/bin/env bash
set -euo pipefail

VENV="${CLEF_VENV:-/workspace/clef-venv}"
CUDA_HOME="${CUDA_HOME:-/usr/local/cuda}"
CUDA_HOME="$(readlink -f "${CUDA_HOME}")"

export CUDA_HOME
export PATH="${CUDA_HOME}/bin:${PATH}"
export CPATH="${CUDA_HOME}/targets/x86_64-linux/include${CPATH:+:${CPATH}}"
export LIBRARY_PATH="${CUDA_HOME}/targets/x86_64-linux/lib/stubs${LIBRARY_PATH:+:${LIBRARY_PATH}}"
export TRITON_PTXAS_PATH="${TRITON_PTXAS_PATH:-${CUDA_HOME}/bin/ptxas}"
export TORCH_CUDA_ARCH_LIST="${TORCH_CUDA_ARCH_LIST:-9.0}"
export MAX_JOBS="${MAX_JOBS:-8}"

python -m venv --system-site-packages "${VENV}"
"${VENV}/bin/python" -m pip install --upgrade pip
"${VENV}/bin/python" -m pip install --upgrade --no-cache-dir \
  transformers==5.12.1 \
  huggingface_hub \
  safetensors \
  pillow \
  accelerate \
  nvidia-cutlass-dsl==4.8.0 \
  apache-tvm-ffi==0.1.14.post1 \
  flash-linear-attention==0.5.2
"${VENV}/bin/python" -m pip install --no-cache-dir --no-build-isolation \
  causal-conv1d==1.7.0

if [[ "${CLEF_INSTALL_INT8:-0}" == "1" ]]; then
  "${VENV}/bin/python" -m pip install --no-cache-dir --no-deps comfy-kitchen==0.2.31
fi

"${VENV}/bin/python" - <<'PY'
import torch
from transformers.utils.import_utils import (
    is_causal_conv1d_available,
    is_flash_linear_attention_available,
)

print("torch", torch.__version__, torch.version.cuda)
print("flash_linear_attention", is_flash_linear_attention_available())
print("causal_conv1d", is_causal_conv1d_available())
PY