test

app.py

@@ -5,7 +5,9 @@ import glob
 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
@@ -14,7 +16,6 @@ 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(
@@ -98,7 +99,8 @@ def upsample_features(image, model_option):
     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.')
+        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()
@@ -107,15 +109,14 @@ def upsample_features(image, model_option):
     # 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', '')
+    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
+    subprocess.check_call(
+        ["pip", "install", "git+https://github.com/mhamilton723/FeatUp"])
 
+    models = {o: torch.hub.load("mhamilton723/FeatUp", o) for o in options}
 
     # Image preprocessing
     input_size = 224