update gpu

app.py

@@ -13,7 +13,6 @@ import gradio as gr
 import spaces
 
 USE_CUDA = torch.cuda.is_available()
-
 print("CUDA is available:", USE_CUDA)
 
 def transform_images(images, resolution=(1024, 1024)):
@@ -142,7 +141,6 @@ class MobileSAM(nn.Module):
 
 mobilesam = MobileSAM()
 
-@spaces.GPU(duration=60)
 def image_mobilesam_feature(
     images,
     node_type="block",
@@ -172,10 +170,9 @@ def image_mobilesam_feature(
         }
         out = out_dict[node_type]
         out = out[layer]
-        outputs.append(out.cpu())
+        outputs.append(out)
     outputs = torch.cat(outputs, dim=0)
     
-    mobilesam = mobilesam.cpu()
     return outputs
 
 
@@ -246,7 +243,6 @@ class SAM(torch.nn.Module):
 
 sam = SAM()
 
-@spaces.GPU(duration=120)
 def image_sam_feature(
     images,
     node_type="block",
@@ -274,10 +270,9 @@ def image_sam_feature(
         }
         out = out_dict[node_type]
         out = out[layer]
-        outputs.append(out.cpu())
+        outputs.append(out)
     outputs = torch.cat(outputs, dim=0)
     
-    sam = sam.cpu()
     
     return outputs
 
@@ -326,7 +321,6 @@ class DiNOv2(torch.nn.Module):
 
 dinov2 = DiNOv2()
 
-@spaces.GPU(duration=60)
 def image_dino_feature(images, node_type="block", layer=-1):
     global USE_CUDA
     if USE_CUDA:
@@ -350,11 +344,10 @@ def image_dino_feature(images, node_type="block", layer=-1):
         }
         out = out_dict[node_type]
         out = out[layer]
-        outputs.append(out.cpu())
+        outputs.append(out)
     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
 
 
@@ -423,7 +416,6 @@ class CLIP(torch.nn.Module):
 
 clip = CLIP()
 
-@spaces.GPU(duration=60)
 def image_clip_feature(
     images, node_type="block", layer=-1
 ):
@@ -449,10 +441,9 @@ def image_clip_feature(
         }
         out = out_dict[node_type]
         out = out[layer]
-        outputs.append(out.cpu())
+        outputs.append(out)
     outputs = torch.cat(outputs, dim=0)
     
-    clip = clip.cpu()
     return outputs
 
 
@@ -514,7 +505,6 @@ def run_model_on_image(images, model_name="sam", node_type="block", layer=-1):
     else:
         raise ValueError(f"Model {model_name} not supported.")
     
-    USE_CUDA = False
     return result
 
 def extract_features(images, model_name="MobileSAM", node_type="block", layer=-1):
@@ -550,14 +540,14 @@ def compute_ncut(
     knn_tsne=10,
     num_sample_tsne=1000,
     perplexity=500,
-):
+):        
     from ncut_pytorch import NCUT, rgb_from_tsne_3d
 
     start = time.time()
     eigvecs, eigvals = NCUT(
         num_eig=num_eig,
         num_sample=num_sample_ncut,
-        device="cpu",
+        device="cuda" if USE_CUDA else "cpu",
         affinity_focal_gamma=affinity_focal_gamma,
         knn=knn_ncut,
     ).fit_transform(features.reshape(-1, features.shape[-1]))
@@ -595,7 +585,7 @@ def to_pil_images(images):
         for image in images
     ]
 
-
+@spaces.GPU(duration=60)
 def main_fn(
     images,
     model_name="SAM(sam_vit_b)",