#!/bin/bash
set -e

source /home/fabian/ComfyUI/venv/bin/activate

export HSA_OVERRIDE_GFX_VERSION=10.1.0
export HSA_ENABLE_SDMA=0
export HIP_VISIBLE_DEVICES=0
export ROCM_PATH=/opt/rocm
export PYTORCH_ROCM_ARCH=gfx1010
export FORCE_CUDA=1
export TORCH_CUDA_ARCH_LIST=""

# Uninstall old torchvision
pip uninstall torchvision -y 2>/dev/null || true

# Clone torchvision matching PyTorch 2.5
cd /tmp
rm -rf torchvision_build
git clone --depth 1 --branch v0.20.0 https://github.com/pytorch/vision.git torchvision_build
cd torchvision_build

# Build and install
python setup.py install 2>&1 | tail -20

echo "=== DONE ==="
python -c "import torchvision; print('torchvision version:', torchvision.__version__)"