nielsr HF staff commited on
Commit
7baa4fb
1 Parent(s): 149b7f2

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +3 -4
app.py CHANGED
@@ -1,6 +1,6 @@
1
  import torch
2
  import gradio as gr
3
- from transformers import OwlViTProcessor, OwlViTForObjectDetection
4
 
5
 
6
  # Use GPU if available
@@ -9,9 +9,8 @@ if torch.cuda.is_available():
9
  else:
10
  device = torch.device("cpu")
11
 
12
- model = OwlViTForObjectDetection.from_pretrained("nielsr/owlv2-base-patch16-ensemble").to(device)
13
- model.eval()
14
- processor = OwlViTProcessor.from_pretrained("nielsr/owlv2-base-patch16-ensemble")
15
 
16
 
17
  def query_image(img, text_queries, score_threshold):
 
1
  import torch
2
  import gradio as gr
3
+ from transformers import Owlv2Processor, Owlv2ForObjectDetection
4
 
5
 
6
  # Use GPU if available
 
9
  else:
10
  device = torch.device("cpu")
11
 
12
+ model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16-ensemble").to(device)
13
+ processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16-ensemble")
 
14
 
15
 
16
  def query_image(img, text_queries, score_threshold):