Update app.py

app.py

@@ -20,12 +20,13 @@ clf_1 = pipeline(model=model_1, task="image-classification", image_processor=ima
 
 # Load the second model
 model_2_path = "Heem2/AI-vs-Real-Image-Detection"
-clf_2 = pipeline("image-classification", model=model_2_path)
+clf_2 = pipeline("image-classification", model=model_2_path, device=device)
 
 # Define class names for both models
 class_names_1 = ['artificial', 'real']
 class_names_2 = ['AI Image', 'Real Image']  # Adjust if the second model has different classes
 
+@spaces.GPU(duration=30)
 def predict_image(img, confidence_threshold):
     # Ensure the image is a PIL Image
     if not isinstance(img, Image.Image):