test

app.py

@@ -1,7 +1,5 @@
 import os
 import subprocess
-from setuptools import setup, find_packages
-from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
 
 import matplotlib.pyplot as plt
 import torch
@@ -13,6 +11,7 @@ import requests
 import csv
 import spaces
 
+
 def plot_feats(image, lr, hr):
     from featup.util import pca, remove_axes
     assert len(image.shape) == len(lr.shape) == len(hr.shape) == 3
@@ -92,6 +91,12 @@ model_option = gr.Radio(options, value="dino16",
 @spaces.GPU(duration=120)
 def upsample_features(image, model_option):
     with torch.no_grad():
+        from setuptools import setup, find_packages
+        from torch.utils.cpp_extension import BuildExtension, CUDAExtension, CppExtension
+        from torch.utils.cpp_extension import _find_cuda_home
+
+        os.environ['CUDA_HOME'] = _find_cuda_home()
+        print(os.environ['CUDA_HOME'])
         setup(
             name='featup',
             version='0.1.2',