test

app.py

@@ -2,31 +2,10 @@ import os
 import subprocess
 import glob
 
-# Find all CUDA directories that match /usr/local/cuda*
-cuda_dirs = glob.glob('/usr/local/cuda*')
-
-if not cuda_dirs:
-    raise EnvironmentError('No CUDA installation found. Please install CUDA or set CUDA_HOME manually.')
-
-# Assume the highest version of CUDA is the one to use
-cuda_dirs.sort()
-cuda_home = cuda_dirs[-1]
-
-# Set the CUDA_HOME environment variable
-os.environ['CUDA_HOME'] = cuda_home
-os.environ['PATH'] = os.environ['CUDA_HOME'] + '/bin:' + os.environ['PATH']
-os.environ['LD_LIBRARY_PATH'] = os.environ['CUDA_HOME'] + '/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')
-
-# Install the required package from GitHub
-subprocess.check_call(["pip", "install", "git+https://github.com/mhamilton723/FeatUp"])
-
-
 import matplotlib.pyplot as plt
 import torch
 import torchvision.transforms as T
-from PIL import Image
 import gradio as gr
-from featup.util import norm, unnorm, pca, remove_axes
 from pytorch_lightning import seed_everything
 import os
 import requests
@@ -35,6 +14,7 @@ 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
     seed_everything(0)
     [lr_feats_pca, hr_feats_pca], _ = pca(
@@ -113,6 +93,30 @@ models = {o: torch.hub.load("mhamilton723/FeatUp", o) for o in options}
 
 @spaces.GPU
 def upsample_features(image, model_option):
+
+    # Find all CUDA directories that match /usr/local/cuda*
+    cuda_dirs = glob.glob('/usr/local/cuda*')
+
+    if not cuda_dirs:
+        raise EnvironmentError('No CUDA installation found. Please install CUDA or set CUDA_HOME manually.')
+
+    # Assume the highest version of CUDA is the one to use
+    cuda_dirs.sort()
+    cuda_home = cuda_dirs[-1]
+
+    # Set the CUDA_HOME environment variable
+    os.environ['CUDA_HOME'] = cuda_home
+    os.environ['PATH'] = os.environ['CUDA_HOME'] + '/bin:' + os.environ['PATH']
+    os.environ['LD_LIBRARY_PATH'] = os.environ['CUDA_HOME'] + '/lib64:' + os.environ.get('LD_LIBRARY_PATH', '')
+
+    # Install the required package from GitHub
+    subprocess.check_call(["pip", "install", "git+https://github.com/mhamilton723/FeatUp"])
+
+
+
+    from featup.util import norm, unnorm
+
+
     # Image preprocessing
     input_size = 224
     transform = T.Compose([