update gpu

app.py

@@ -145,6 +145,7 @@ class MobileSAM(nn.Module):
 
 mobilesam = MobileSAM()
 
+@spaces.GPU(duration=30)
 def image_mobilesam_feature(
     images,
     node_type="block",
@@ -247,6 +248,7 @@ class SAM(torch.nn.Module):
 
 sam = SAM()
 
+@spaces.GPU(duration=60)
 def image_sam_feature(
     images,
     node_type="block",
@@ -327,6 +329,7 @@ class DiNOv2(torch.nn.Module):
 
 dinov2 = DiNOv2()
 
+@spaces.GPU(duration=30)
 def image_dino_feature(images, node_type="block", layer=-1):
 
     if USE_CUDA:
@@ -424,6 +427,7 @@ class CLIP(torch.nn.Module):
 
 clip = CLIP()
 
+@spaces.GPU(duration=30)
 def image_clip_feature(
     images, node_type="block", layer=-1
 ):
@@ -495,7 +499,6 @@ def compute_hash(*args, **kwargs):
     return hasher.hexdigest()
 
 
-@spaces.GPU(duration=30)
 def run_model_on_image(images, model_name="sam", node_type="block", layer=-1):
     global USE_CUDA
     USE_CUDA = True
@@ -516,12 +519,12 @@ def run_model_on_image(images, model_name="sam", node_type="block", layer=-1):
     USE_CUDA = False
     return result
 
-def extract_features(images, model_name="mobilesam", node_type="block", layer=-1):
+def extract_features(images, model_name="MobileSAM", node_type="block", layer=-1):
     resolution_dict = {
-        "mobilesam": (1024, 1024),
-        "sam(sam_vit_b)": (1024, 1024),
-        "dinov2(dinov2_vitb14_reg)": (448, 448),
-        "clip(openai/clip-vit-base-patch16)": (224, 224),
+        "MobileSAM": (1024, 1024),
+        "SAM(sam_vit_b)": (1024, 1024),
+        "DiNO(dinov2_vitb14_reg)": (448, 448),
+        "CLIP(openai/clip-vit-base-patch16)": (224, 224),
     }
     images = transform_images(images, resolution=resolution_dict[model_name])