Update app.py

app.py

@@ -1,8 +1,9 @@
-import gradio as gr
 import cv2
 from PIL import Image
 from transformers import AutoImageProcessor, AutoModelForImageClassification
 import torch
+import gradio as gr
+from spaces import GPU  # ✅ import this to use the decorator
 
 # Load model and processor
 model_name = "dima806/ai_vs_real_image_detection"
@@ -10,11 +11,15 @@ processor = AutoImageProcessor.from_pretrained(model_name)
 model = AutoModelForImageClassification.from_pretrained(model_name)
 model.eval()
 
+# Use GPU if available
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+model = model.to(device)
+
+@GPU  # ✅ This activates GPU on Hugging Face ZeroGPU Spaces
 def analyze_video(video):
     cap = cv2.VideoCapture(video)
     frame_num = 0
     frame_interval = 60  
-    batch_size = 12
     frames_to_process = []
 
     while cap.isOpened():
@@ -34,16 +39,15 @@ def analyze_video(video):
     if not frames_to_process:
         return "No frames extracted."
 
-    # Extract just the images for the processor
     frame_numbers, pil_images = zip(*frames_to_process)
 
-    # Batch inference
     inputs = processor(images=pil_images, return_tensors="pt")
+    inputs = {k: v.to(device) for k, v in inputs.items()}  # ✅ Move inputs to GPU if available
+
     with torch.no_grad():
         outputs = model(**inputs)
         predictions = torch.argmax(outputs.logits, dim=1).tolist()
 
-    # Convert predictions to labels
     results = []
     for frame_idx, pred in zip(frame_numbers, predictions):
         label = model.config.id2label[pred]
@@ -51,12 +55,11 @@ def analyze_video(video):
 
     return "\n".join(results)
 
-# Gradio UI
+# Gradio UI (note: NO gpu=True here)
 gr.Interface(
     fn=analyze_video,
     inputs=gr.Video(label="Upload a video"),
     outputs=gr.Textbox(label="Detection Results"),
     title="AI Frame Detector",
-    description="Detects whether frames in a video are AI-generated or real.",
-    gpu=True
+    description="Detects whether frames in a video are AI-generated or real."
 ).launch()