huzey commited on
Commit
604dab5
1 Parent(s): 69d2985

update gpu

Browse files
Files changed (1) hide show
  1. app.py +8 -7
app.py CHANGED
@@ -145,13 +145,13 @@ class MobileSAM(nn.Module):
145
 
146
  mobilesam = MobileSAM()
147
 
148
- @spaces.GPU(duration=30)
149
  def image_mobilesam_feature(
150
  images,
151
  node_type="block",
152
  layer=-1,
153
  ):
154
-
155
  if USE_CUDA:
156
  images = images.cuda()
157
 
@@ -249,13 +249,13 @@ class SAM(torch.nn.Module):
249
 
250
  sam = SAM()
251
 
252
- @spaces.GPU(duration=60)
253
  def image_sam_feature(
254
  images,
255
  node_type="block",
256
  layer=-1,
257
  ):
258
-
259
  if USE_CUDA:
260
  images = images.cuda()
261
 
@@ -331,9 +331,9 @@ class DiNOv2(torch.nn.Module):
331
 
332
  dinov2 = DiNOv2()
333
 
334
- @spaces.GPU(duration=30)
335
  def image_dino_feature(images, node_type="block", layer=-1):
336
-
337
  if USE_CUDA:
338
  images = images.cuda()
339
 
@@ -430,10 +430,11 @@ class CLIP(torch.nn.Module):
430
 
431
  clip = CLIP()
432
 
433
- @spaces.GPU(duration=30)
434
  def image_clip_feature(
435
  images, node_type="block", layer=-1
436
  ):
 
437
  if USE_CUDA:
438
  images = images.cuda()
439
 
 
145
 
146
  mobilesam = MobileSAM()
147
 
148
+ @spaces.GPU(duration=60)
149
  def image_mobilesam_feature(
150
  images,
151
  node_type="block",
152
  layer=-1,
153
  ):
154
+ global USE_CUDA
155
  if USE_CUDA:
156
  images = images.cuda()
157
 
 
249
 
250
  sam = SAM()
251
 
252
+ @spaces.GPU(duration=120)
253
  def image_sam_feature(
254
  images,
255
  node_type="block",
256
  layer=-1,
257
  ):
258
+ global USE_CUDA
259
  if USE_CUDA:
260
  images = images.cuda()
261
 
 
331
 
332
  dinov2 = DiNOv2()
333
 
334
+ @spaces.GPU(duration=60)
335
  def image_dino_feature(images, node_type="block", layer=-1):
336
+ global USE_CUDA
337
  if USE_CUDA:
338
  images = images.cuda()
339
 
 
430
 
431
  clip = CLIP()
432
 
433
+ @spaces.GPU(duration=60)
434
  def image_clip_feature(
435
  images, node_type="block", layer=-1
436
  ):
437
+ global USE_CUDA
438
  if USE_CUDA:
439
  images = images.cuda()
440