Update app.py

app.py

@@ -38,6 +38,15 @@ def install_cuda_toolkit():
     os.environ["TORCH_CUDA_ARCH_LIST"] = "8.0;8.6"
     print("==> finfish install")
 install_cuda_toolkit()
+@spaces.GPU
+def check_gpu():
+    os.environ['CUDA_HOME'] = '/usr/local/cuda-12.1'
+    os.environ['PATH'] += ':/usr/local/cuda-12.1/bin'
+    # os.environ['LD_LIBRARY_PATH'] += ':/usr/local/cuda-12.1/lib64'
+    os.environ['LD_LIBRARY_PATH'] = "/usr/local/cuda-12.1/lib64:" + os.environ.get('LD_LIBRARY_PATH', '')
+    subprocess.run(['nvidia-smi'])  # 测试 CUDA 是否可用
+    print(f"torch.cuda.is_available:{torch.cuda.is_available()}")
+check_gpu()
 import torch
 import numpy as np
 from PIL import Image