Update app.py

app.py

@@ -3,15 +3,17 @@ import torch
 import numpy as np
 from transformers import AutoProcessor, AutoModel
 from PIL import Image
-import cv2
+import cv2 
 
 MODEL_NAME = "microsoft/xclip-base-patch16-zero-shot"
 CLIP_LEN = 32
 
-# Load model and processor once
-processor = AutoProcessor.from_pretrained(MODEL_NAME)
-model = AutoModel.from_pretrained(MODEL_NAME)
+# Check if GPU is available and set the device
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
+# Load model and processor once and move them to the device
+processor = AutoProcessor.from_pretrained(MODEL_NAME)
+model = AutoModel.from_pretrained(MODEL_NAME).to(device)
 
 def get_video_length(file_path):
     cap = cv2.VideoCapture(file_path)
@@ -49,13 +51,7 @@ def concatenate_frames(frames, clip_len):
     for i in range(rows):
         x_offset = 0
         for j in range(cols):
-            img_array = next(frame_iter)
-            
-            # Handling rank-4 tensor
-            if len(img_array.shape) == 4:
-                img_array = img_array[0]
-                
-            img = Image.fromarray(img_array)
+            img = Image.fromarray(next(frame_iter))
             combined_image.paste(img, (x_offset, y_offset))
             x_offset += frames[0].shape[1]
         y_offset += frames[0].shape[0]
@@ -75,6 +71,11 @@ def model_interface(uploaded_video, activity):
         padding=True,
     )
 
+    # Move the tensors to the same device as the model
+    for key, value in inputs.items():
+        if isinstance(value, torch.Tensor):
+            inputs[key] = value.to(device)
+
     with torch.no_grad():
         outputs = model(**inputs)
 
@@ -86,13 +87,13 @@ def model_interface(uploaded_video, activity):
     max_prob_index = torch.argmax(probs[0]).item()
     for i in range(len(activities_list)):
         current_activity = activities_list[i]
-        prob = float(probs[0][i])
-        logit = float(logits_per_video[0][i])
+        prob = float(probs[0][i].cpu())  # Move tensor data to CPU for further processing
+        logit = float(logits_per_video[0][i].cpu())  # Move tensor data to CPU for further processing
         results_probs.append((current_activity, f"Probability: {prob * 100:.2f}%"))
         results_logits.append((current_activity, f"Raw Score: {logit:.2f}"))
 
     likely_label = activities_list[max_prob_index]
-    likely_probability = float(probs[0][max_prob_index]) * 100
+    likely_probability = float(probs[0][max_prob_index].cpu()) * 100  # Move tensor data to CPU
 
     return concatenated_image, results_probs, results_logits, [likely_label, likely_probability]