echo "Installing PyTorch with CUDA support..."
pip install torch==2.2.0+cu121 torchvision==0.17.0+cu121 torchaudio==2.2.0+cu121 \
--extra-index-url https://download.pytorch.org/whl/cu121

echo "Reinstalling NumPy with specific version..."
pip install --force-reinstall -v "numpy==1.25.2"


# # Compile and install DISO package
# echo "Compiling and installing DISO..."
# pip install ninja
# mkdir -p diso_package
# cd diso_package
# git clone https://github.com/SarahWeiii/diso.git
# cd diso
# pip install -e .
# cd ..
# cd ..



echo "Installing other requirements..."
pip install -r requirements.txt


echo "Installing pre-compiled DISO wheel package..."
huggingface-cli download --resume-download ZhiyuanthePony/TriplaneTurbo \
  --include "diso-0.1.4-*.whl" \
  --local-dir ./diso_package \
  --local-dir-use-symlinks False

pip install ./diso_package/diso-0.1.4-*.whl
echo "Setup completed successfully!"

echo "Installing compatible dependency versions..."
pip uninstall -y pydantic
pip install pydantic==1.10.8  # Install compatible older version

# Ensure Gradio and other dependencies are installed correctly
pip install "gradio>=4.0.0,<5.0.0"
pip install "fastapi<0.103.0"  # Ensure compatible FastAPI version