bmay commited on
Commit
b194803
·
verified ·
1 Parent(s): 2594b39

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -3
app.py CHANGED
@@ -1,14 +1,14 @@
1
  import gradio as gr
 
2
  import torch
3
  import numpy as np
4
  from transformers import AutoModel
5
  from theia.decoding import load_feature_stats, prepare_depth_decoder, prepare_mask_generator, decode_everything
6
 
7
- device = "cuda:0" if torch.cuda.is_available() else "cpu"
8
-
9
  def run_theia(image):
10
  theia_model = AutoModel.from_pretrained("theaiinstitute/theia-base-patch16-224-cdiv", trust_remote_code=True)
11
- theia_model = theia_model.to(device)
12
  target_model_names = [
13
  "google/vit-huge-patch14-224-in21k",
14
  "facebook/dinov2-large",
 
1
  import gradio as gr
2
+ import spaces
3
  import torch
4
  import numpy as np
5
  from transformers import AutoModel
6
  from theia.decoding import load_feature_stats, prepare_depth_decoder, prepare_mask_generator, decode_everything
7
 
8
+ @spaces.GPU
 
9
  def run_theia(image):
10
  theia_model = AutoModel.from_pretrained("theaiinstitute/theia-base-patch16-224-cdiv", trust_remote_code=True)
11
+ theia_model = theia_model.to('cuda')
12
  target_model_names = [
13
  "google/vit-huge-patch14-224-in21k",
14
  "facebook/dinov2-large",