abreza commited on
Commit
173c22a
1 Parent(s): d6247cd
Files changed (1) hide show
  1. app.py +8 -1
app.py CHANGED
@@ -88,6 +88,7 @@ image_input = gr.Image(label="Choose an image to featurize",
88
  model_option = gr.Radio(options, value="dino16",
89
  label='Choose a backbone to upsample')
90
 
 
91
  def find_cuda_home():
92
  try:
93
  # Define the search string and the directory
@@ -107,10 +108,16 @@ def find_cuda_home():
107
  print(f"Error finding CUDA_HOME: {e}")
108
  return None
109
 
 
110
  @spaces.GPU
111
  def upsample_features(image, model_option):
 
112
  from torch.utils.cpp_extension import _find_cuda_home
113
- os.environ["CUDA_HOME"] = _find_cuda_home()
 
 
 
 
114
  print(os.environ["CUDA_HOME"])
115
  os.environ['PATH'] = '/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin'
116
  os.environ['LD_LIBRARY_PATH'] = '/usr/local/nvidia/lib:/usr/local/nvidia/lib64'
 
88
  model_option = gr.Radio(options, value="dino16",
89
  label='Choose a backbone to upsample')
90
 
91
+
92
  def find_cuda_home():
93
  try:
94
  # Define the search string and the directory
 
108
  print(f"Error finding CUDA_HOME: {e}")
109
  return None
110
 
111
+
112
  @spaces.GPU
113
  def upsample_features(image, model_option):
114
+ print(torch.cuda.is_available())
115
  from torch.utils.cpp_extension import _find_cuda_home
116
+ cuda_home = _find_cuda_home()
117
+ if cuda_home is None:
118
+ raise EnvironmentError("CUDA_HOME could not be found.")
119
+
120
+ os.environ["CUDA_HOME"] = cuda_home
121
  print(os.environ["CUDA_HOME"])
122
  os.environ['PATH'] = '/usr/local/nvidia/bin:/usr/local/cuda/bin:/usr/local/sbin:/usr/local/bin:/usr/sbin:/usr/bin:/sbin:/bin'
123
  os.environ['LD_LIBRARY_PATH'] = '/usr/local/nvidia/lib:/usr/local/nvidia/lib64'