test

app.py

@@ -12,6 +12,28 @@ import csv
 import spaces
 # from triton.fb import build_paths
 
+from torch.utils.cpp_extension import load
+
+# Compile and load the CUDA extension
+adaptive_conv_cuda_impl = load(
+    name='adaptive_conv_cuda_impl',
+    sources=[
+        'featup/adaptive_conv_cuda/adaptive_conv_cuda.cpp',
+        'featup/adaptive_conv_cuda/adaptive_conv_kernel.cu'
+    ],
+    extra_cflags=['-O3'],
+    extra_cuda_cflags=['-O3']
+)
+
+# Compile and load the C++ extension
+adaptive_conv_cpp_impl = load(
+    name='adaptive_conv_cpp_impl',
+    sources=['featup/adaptive_conv_cuda/adaptive_conv.cpp'],
+    extra_cflags=['-O3']
+)
+
+print(adaptive_conv_cuda_impl, adaptive_conv_cpp_impl)
+
 
 def plot_feats(image, lr, hr):
     from featup.util import pca, remove_axes
@@ -89,46 +111,8 @@ model_option = gr.Radio(options, value="dino16",
                         label='Choose a backbone to upsample')
 
 
-def find_cuda_home():
-    try:
-        # Define the search string and the directory
-        search_string = "CUDA"
-        search_directory = "/usr"
-
-        # Use subprocess to run the grep command
-        command = ['grep', '-r', search_string, search_directory]
-        output = subprocess.check_output(command).decode()
-
-        print(output)
-        for line in output.split('\n'):
-            if 'Cuda compilation tools' in line:
-                version = line.split()[-1]
-                return f"/usr/local/cuda-{version.split('.')[0]}.{version.split('.')[1]}"
-    except Exception as e:
-        print(f"Error finding CUDA_HOME: {e}")
-        return None
-    
-
-
-print(torch.cuda.is_available())
-from torch.utils.cpp_extension import _find_cuda_home
-cuda_home = _find_cuda_home()
-if cuda_home is None:
-    raise EnvironmentError("CUDA_HOME could not be found.")
-
-
 @spaces.GPU
 def upsample_features(image, model_option):
-
-    os.environ["CUDA_HOME"] = cuda_home
-    print(os.environ["CUDA_HOME"])
-    os.environ['PATH'] = '/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin'
-    os.environ['LD_LIBRARY_PATH'] = '/usr/local/nvidia/lib:/usr/local/nvidia/lib64'
-
-    # Install the required package from GitHub
-    subprocess.check_call(
-        ["pip", "install", "git+https://github.com/mhamilton723/FeatUp"])
-
     from featup.util import norm, unnorm
     models = {o: torch.hub.load("mhamilton723/FeatUp", o) for o in options}