update app.py

app.py

@@ -10,6 +10,10 @@ import time
 
 import gradio as gr
 
+use_cuda = torch.cuda.is_available()
+
+print("CUDA is available:", use_cuda)
+
 
 class SAM(torch.nn.Module):
     def __init__(self, checkpoint="/data/sam_model/sam_vit_b_01ec64.pth", **kwargs):
@@ -50,7 +54,8 @@ class SAM(torch.nn.Module):
 
         self.image_encoder = sam.image_encoder
         self.image_encoder.eval()
-        # self.image_encoder = self.image_encoder.cuda()
+        if use_cuda:
+            self.image_encoder = self.image_encoder.cuda()
 
     @torch.no_grad()
     def forward(self, x: torch.Tensor) -> torch.Tensor:
@@ -98,8 +103,9 @@ def image_sam_feature(
     outputs = []
     for i, image in enumerate(images):
         torch_image = transform(image)
+        if use_cuda:
+            torch_image = torch_image.cuda()
         attn_output, mlp_output, block_output = feat_extractor(
-            # torch_image.unsqueeze(0).cuda()
             torch_image.unsqueeze(0)
         )
         out_dict = {
@@ -120,7 +126,8 @@ class DiNOv2(torch.nn.Module):
         self.dinov2 = torch.hub.load("facebookresearch/dinov2", ver)
         self.dinov2.requires_grad_(False)
         self.dinov2.eval()
-        # self.dinov2 = self.dinov2.cuda()
+        if use_cuda:
+            self.dinov2 = self.dinov2.cuda()
 
         def new_block_forward(self, x: torch.Tensor) -> torch.Tensor:
             def attn_residual_func(x):
@@ -173,8 +180,9 @@ def image_dino_feature(images, resolution=(448, 448), node_type="block", layer=-
     outputs = []
     for i, image in enumerate(images):
         torch_image = transform(image)
+        if use_cuda:
+            torch_image = torch_image.cuda()
         attn_output, mlp_output, block_output = feat_extractor(
-            # torch_image.unsqueeze(0).cuda()
             torch_image.unsqueeze(0)
         )
         out_dict = {
@@ -199,8 +207,9 @@ 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()
-        # self.model = self.model.cuda()
-        
+        if use_cuda:
+            self.model = self.model.cuda()
+
         def new_forward(
             self,
             hidden_states: torch.Tensor,
@@ -277,8 +286,9 @@ def image_clip_feature(
     outputs = []
     for i, image in enumerate(images):
         torch_image = transform(image)
+        if use_cuda:
+            torch_image = torch_image.cuda()
         attn_output, mlp_output, block_output = feat_extractor(
-            # torch_image.unsqueeze(0).cuda()
             torch_image.unsqueeze(0)
         )
         out_dict = {
@@ -321,7 +331,7 @@ def compute_ncut(
     eigvecs, eigvals = NCUT(
         num_eig=num_eig,
         num_sample=num_sample_ncut,
-        # device="cuda:0",
+        device="cuda" if use_cuda else "cpu",
         affinity_focal_gamma=affinity_focal_gamma,
         knn=knn_ncut,
     ).fit_transform(features.reshape(-1, features.shape[-1]))