reab5555 commited on
Commit
3f18cbf
·
verified ·
1 Parent(s): 7e7ddb7

Update app.py

Browse files
Files changed (1) hide show
  1. app.py +2 -8
app.py CHANGED
@@ -11,18 +11,12 @@ import tempfile
11
  import shutil
12
 
13
  # Check if CUDA is available, otherwise use CPU
14
- device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
15
 
16
  processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
17
  model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16")
18
 
19
- # Try to move model to GPU and use half precision
20
- try:
21
- model = model.to(device).half()
22
- except RuntimeError:
23
- print("GPU out of memory, using CPU instead")
24
- device = torch.device("cpu")
25
- model = model.to(device)
26
 
27
  @spaces.GPU(duration=120)
28
  def process_video(video_path, target, progress=gr.Progress()):
 
11
  import shutil
12
 
13
  # Check if CUDA is available, otherwise use CPU
14
+ device = "cuda"
15
 
16
  processor = Owlv2Processor.from_pretrained("google/owlv2-base-patch16")
17
  model = Owlv2ForObjectDetection.from_pretrained("google/owlv2-base-patch16")
18
 
19
+ model = model.to(device).half()
 
 
 
 
 
 
20
 
21
  @spaces.GPU(duration=120)
22
  def process_video(video_path, target, progress=gr.Progress()):