Update app.py

app.py

@@ -3,14 +3,19 @@ import torch
 import numpy as np
 from transformers import AutoProcessor, AutoModel
 from PIL import Image
-from decord import VideoReader, cpu
+from decord import VideoReader, cpu, gpu
 
 MODEL_NAME = "microsoft/xclip-base-patch16-zero-shot"
 CLIP_LEN = 32
 
-# Load model and processor once
+# Check for GPU availability
+device = "cuda" if torch.cuda.is_available() else "cpu"
+print (device)
+
+# Load model and processor once and move them to the GPU
 processor = AutoProcessor.from_pretrained(MODEL_NAME)
-model = AutoModel.from_pretrained(MODEL_NAME)
+model = AutoModel.from_pretrained(MODEL_NAME).to(device)
+model.eval()
 
 def sample_uniform_frame_indices(clip_len, seg_len):
     if seg_len < clip_len:
@@ -23,7 +28,9 @@ def sample_uniform_frame_indices(clip_len, seg_len):
     return np.array(indices).astype(np.int64)
 
 def read_video_decord(file_path, indices):
-    vr = VideoReader(file_path, num_threads=1, ctx=cpu(0))
+    # Use GPU for video decoding if available
+    vr_ctx = cpu(0)
+    vr = VideoReader(file_path, num_threads=1, ctx=vr_ctx)
     video = vr.get_batch(indices).asnumpy()
     return video
 
@@ -55,6 +62,9 @@ def model_interface(uploaded_video, activity):
         padding=True,
     )
 
+    # Move inputs to GPU
+    inputs = {name: tensor.to(device) for name, tensor in inputs.items()}
+
     with torch.no_grad():
         outputs = model(**inputs)
 
@@ -66,13 +76,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())
+        logit = float(logits_per_video[0][i].cpu())
         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
 
     return concatenated_image, results_probs, results_logits, [ likely_label , likely_probability ]