Update app.py and requirements.txt

app.py

@@ -1,4 +1,5 @@
 import shutil
+import subprocess
 import time
 from pathlib import Path
 from typing import Tuple
@@ -10,6 +11,20 @@ import torch
 import yaml
 from box import Box
 
+# Get the PyTorch and CUDA versions
+torch_version = torch.__version__.split("+")[0]  # Strips any "+cuXXX" suffix
+cuda_version = torch.version.cuda
+spconv_version = "-cu121" if cuda_version else "" 
+
+# Format CUDA version to match the URL convention (e.g., "cu118" for CUDA 11.8)
+if cuda_version:
+    cuda_version = f"cu{cuda_version.replace('.', '')}"
+else:
+    cuda_version = "cpu"  # Fallback in case CUDA is not available
+
+subprocess.run(f'pip install spconv{spconv_version}', shell=True)
+subprocess.run(f'pip install torch_scatter torch_cluster -f https://data.pyg.org/whl/torch-{torch_version}+{cuda_version}.html --no-cache-dir', shell=True)
+
 from src.data.datapath import Datapath
 from src.data.dataset import DatasetConfig, UniRigDatasetModule
 from src.data.extract import extract_builtin, get_files

requirements.txt

@@ -17,18 +17,4 @@ scipy
 matplotlib
 plotly
 pyyaml
-# PyTorch and related packages - ensure compatibility
-torch==2.7.0
-torchvision
-torchaudio
-# PyTorch Geometric ecosystem packages
---find-links https://data.pyg.org/whl/torch-2.7.0+cu126.html
-torch-scatter
-torch-sparse
-torch-cluster
-torch-spline-conv
-torch-geometric
-# Flash attention for PyTorch 2.7
 https://github.com/mjun0812/flash-attention-prebuild-wheels/releases/download/v0.0.8/flash_attn-2.7.4.post1+cu126torch2.7-cp310-cp310-linux_x86_64.whl
-# Sparse convolution
-spconv-cu121