update gpu

app.py

@@ -11,11 +11,11 @@ import time
 
 import gradio as gr
 
-use_cuda = torch.cuda.is_available()
+import spaces
 
-# use_cuda = False
+USE_CUDA = torch.cuda.is_available()
 
-print("CUDA is available:", use_cuda)
+print("CUDA is available:", USE_CUDA)
 
 class MobileSAM(nn.Module):
     def __init__(self, **kwargs):
@@ -32,7 +32,7 @@ class MobileSAM(nn.Module):
             with open(sam_checkpoint, 'wb') as f:
                 f.write(r.content)
 
-        device = 'cuda' if use_cuda else 'cpu'
+        device = 'cuda' if USE_CUDA else 'cpu'
 
         mobile_sam = sam_model_registry[model_type](checkpoint=sam_checkpoint)
 
@@ -135,6 +135,7 @@ class MobileSAM(nn.Module):
                 block_outputs.append(blk.block_output)
         return attn_outputs, mlp_outputs, block_outputs
 
+mobilesam = MobileSAM()
 
 def image_mobilesam_feature(
     images,
@@ -152,13 +153,15 @@ def image_mobilesam_feature(
     )
     
 
-    feat_extractor = MobileSAM()
+    feat_extractor = mobilesam
+    if USE_CUDA:
+        feat_extractor = feat_extractor.cuda()
 
     # attn_outputs, mlp_outputs, block_outputs = [], [], []
     outputs = []
     for i, image in enumerate(images):
         torch_image = transform(image)
-        if use_cuda:
+        if USE_CUDA:
             torch_image = torch_image.cuda()
         attn_output, mlp_output, block_output = feat_extractor(
             torch_image.unsqueeze(0)
@@ -172,15 +175,25 @@ def image_mobilesam_feature(
         out = out[layer]
         outputs.append(out.cpu())
     outputs = torch.cat(outputs, dim=0)
+    
+    mobilesam = mobilesam.cpu()
     return outputs
 
 
 
 class SAM(torch.nn.Module):
-    def __init__(self, checkpoint="/data/sam_model/sam_vit_b_01ec64.pth", **kwargs):
+    def __init__(self, **kwargs):
         super().__init__(**kwargs)
         from segment_anything import sam_model_registry, SamPredictor
         from segment_anything.modeling.sam import Sam
+        
+        checkpoint = "sam_vit_b_01ec64.pth"
+        if not os.path.exists(checkpoint):
+            checkpoint_url = 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth'
+            import requests
+            r = requests.get(checkpoint_url)
+            with open(checkpoint, 'wb') as f:
+                f.write(r.content)
 
         sam: Sam = sam_model_registry["vit_b"](checkpoint=checkpoint)
 
@@ -215,7 +228,7 @@ class SAM(torch.nn.Module):
 
         self.image_encoder = sam.image_encoder
         self.image_encoder.eval()
-        if use_cuda:
+        if USE_CUDA:
             self.image_encoder = self.image_encoder.cuda()
 
     @torch.no_grad()
@@ -234,6 +247,7 @@ class SAM(torch.nn.Module):
         block_outputs = torch.stack(block_outputs)
         return attn_outputs, mlp_outputs, block_outputs
 
+sam = SAM()
 
 def image_sam_feature(
     images,
@@ -249,22 +263,16 @@ def image_sam_feature(
             transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225]),
         ]
     )
-    
-    checkpoint = "sam_vit_b_01ec64.pth"
-    if not os.path.exists(checkpoint):
-        checkpoint_url = 'https://dl.fbaipublicfiles.com/segment_anything/sam_vit_b_01ec64.pth'
-        import requests
-        r = requests.get(checkpoint_url)
-        with open(checkpoint, 'wb') as f:
-            f.write(r.content)
 
-    feat_extractor = SAM(checkpoint=checkpoint)
+    feat_extractor = sam
+    if USE_CUDA:
+        feat_extractor = feat_extractor.cuda()
 
     # attn_outputs, mlp_outputs, block_outputs = [], [], []
     outputs = []
     for i, image in enumerate(images):
         torch_image = transform(image)
-        if use_cuda:
+        if USE_CUDA:
             torch_image = torch_image.cuda()
         attn_output, mlp_output, block_output = feat_extractor(
             torch_image.unsqueeze(0)
@@ -278,6 +286,9 @@ def image_sam_feature(
         out = out[layer]
         outputs.append(out.cpu())
     outputs = torch.cat(outputs, dim=0)
+    
+    sam = sam.cpu()
+    
     return outputs
 
 
@@ -287,7 +298,7 @@ class DiNOv2(torch.nn.Module):
         self.dinov2 = torch.hub.load("facebookresearch/dinov2", ver)
         self.dinov2.requires_grad_(False)
         self.dinov2.eval()
-        if use_cuda:
+        if USE_CUDA:
             self.dinov2 = self.dinov2.cuda()
 
         def new_block_forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -325,6 +336,7 @@ class DiNOv2(torch.nn.Module):
         block_outputs = torch.stack(block_outputs)
         return attn_outputs, mlp_outputs, block_outputs
 
+dinov2 = DiNOv2()
 
 def image_dino_feature(images, resolution=(448, 448), node_type="block", layer=-1):
 
@@ -336,12 +348,14 @@ def image_dino_feature(images, resolution=(448, 448), node_type="block", layer=-
         ]
     )
 
-    feat_extractor = DiNOv2()
+    feat_extractor = dinov2
+    if USE_CUDA:
+        feat_extractor = feat_extractor.cuda()
 
     outputs = []
     for i, image in enumerate(images):
         torch_image = transform(image)
-        if use_cuda:
+        if USE_CUDA:
             torch_image = torch_image.cuda()
         attn_output, mlp_output, block_output = feat_extractor(
             torch_image.unsqueeze(0)
@@ -356,6 +370,8 @@ def image_dino_feature(images, resolution=(448, 448), node_type="block", layer=-
         outputs.append(out.cpu())
     outputs = torch.cat(outputs, dim=0)
     outputs = rearrange(outputs[:, 5:, :], "b (h w) c -> b h w c", h=32, w=32)
+    
+    dinov2 = dinov2.cpu()
     return outputs
 
 
@@ -368,7 +384,7 @@ class CLIP(torch.nn.Module):
         model = CLIPModel.from_pretrained("openai/clip-vit-base-patch16")
         # processor = CLIPProcessor.from_pretrained("openai/clip-vit-base-patch16")
         self.model = model.eval()
-        if use_cuda:
+        if USE_CUDA:
             self.model = self.model.cuda()
 
         def new_forward(
@@ -424,6 +440,7 @@ class CLIP(torch.nn.Module):
         block_outputs = torch.stack(block_outputs)
         return attn_outputs, mlp_outputs, block_outputs
 
+clip = CLIP()
 
 def image_clip_feature(
     images, resolution=(224, 224), node_type="block", layer=-1
@@ -442,12 +459,14 @@ def image_clip_feature(
         ]
     )
 
-    feat_extractor = CLIP()
+    feat_extractor = clip
+    if USE_CUDA:
+        feat_extractor = feat_extractor.cuda()
 
     outputs = []
     for i, image in enumerate(images):
         torch_image = transform(image)
-        if use_cuda:
+        if USE_CUDA:
             torch_image = torch_image.cuda()
         attn_output, mlp_output, block_output = feat_extractor(
             torch_image.unsqueeze(0)
@@ -461,6 +480,8 @@ def image_clip_feature(
         out = out[layer]
         outputs.append(out.cpu())
     outputs = torch.cat(outputs, dim=0)
+    
+    clip = clip.cpu()
     return outputs
 
 
@@ -505,6 +526,27 @@ def compute_hash(*args, **kwargs):
     return hasher.hexdigest()
 
 
+@spaces.GPU(duration=30)
+def run_model_on_image(image, model_name="sam", node_type="block", layer=-1):
+    global USE_CUDA
+    USE_CUDA = True
+    
+    if model_name == "SAM(sam_vit_b)":
+        if not USE_CUDA:
+            gr.warning("GPU not detected. Running SAM on CPU, ~30s/image.")
+        result = image_sam_feature([image], node_type=node_type, layer=layer)
+    elif model_name == 'MobileSAM':
+        result = image_mobilesam_feature([image], node_type=node_type, layer=layer)
+    elif model_name == "DiNO(dinov2_vitb14_reg)":
+        result = image_dino_feature([image], node_type=node_type, layer=layer)
+    elif model_name == "CLIP(openai/clip-vit-base-patch16)":
+        result = image_clip_feature([image], node_type=node_type, layer=layer)
+    else:
+        raise ValueError(f"Model {model_name} not supported.")
+    
+    USE_CUDA = False
+    return result
+
 def extract_features(images, model_name="sam", node_type="block", layer=-1):
     # Compute the cache key
     cache_key = compute_hash(images, model_name, node_type, layer)
@@ -514,20 +556,7 @@ def extract_features(images, model_name="sam", node_type="block", layer=-1):
         print("Cache hit!")
         return cache[cache_key]
     
-    
-    # Compute the result if not in cache
-    if model_name == "SAM(sam_vit_b)":
-        if not use_cuda:
-            gr.warning("GPU not detected. Running SAM on CPU, ~30s/image.")
-        result = image_sam_feature(images, node_type=node_type, layer=layer)
-    elif model_name == 'MobileSAM':
-        result = image_mobilesam_feature(images, node_type=node_type, layer=layer)
-    elif model_name == "DiNO(dinov2_vitb14_reg)":
-        result = image_dino_feature(images, node_type=node_type, layer=layer)
-    elif model_name == "CLIP(openai/clip-vit-base-patch16)":
-        result = image_clip_feature(images, node_type=node_type, layer=layer)
-    else:
-        raise ValueError(f"Model {model_name} not supported.")
+    result = run_model_on_image(images[0], model_name=model_name, node_type=node_type, layer=layer)
 
     # Store the result in the cache
     cache[cache_key] = result
@@ -550,11 +579,11 @@ def compute_ncut(
     eigvecs, eigvals = NCUT(
         num_eig=num_eig,
         num_sample=num_sample_ncut,
-        device="cuda" if use_cuda else "cpu",
+        device="cpu",
         affinity_focal_gamma=affinity_focal_gamma,
         knn=knn_ncut,
     ).fit_transform(features.reshape(-1, features.shape[-1]))
-    print(f"NCUT time: {time.time() - start:.2f}s")
+    print(f"NCUT time (cpu): {time.time() - start:.2f}s")
     
     start = time.time()
     X_3d, rgb = rgb_from_tsne_3d(
@@ -563,7 +592,7 @@ def compute_ncut(
         perplexity=perplexity,
         knn=knn_tsne,
     )
-    print(f"t-SNE time: {time.time() - start:.2f}s")
+    print(f"t-SNE time (cpu): {time.time() - start:.2f}s")
     
     # print("input shape:", features.shape)
     # print("output shape:", rgb.shape)
@@ -613,7 +642,7 @@ def main_fn(
     features = extract_features(
         images, model_name=model_name, node_type=node_type, layer=layer
     )
-    print(f"Feature extraction time: {time.time() - start:.2f}s")
+    print(f"Feature extraction time (gpu): {time.time() - start:.2f}s")
     
     rgb = compute_ncut(
         features,