Update app.py

app.py

@@ -4,12 +4,16 @@ from transformers import pipeline, AutoImageProcessor, Swinv2ForImageClassificat
 from torchvision import transforms
 import torch
 from PIL import Image
+import warnings
+
+# Suppress warnings
+warnings.filterwarnings("ignore", category=UserWarning, message="Using a slow image processor as `use_fast` is unset")
 
 # Ensure using GPU if available
 device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
 # Load the first model and processor
-image_processor_1 = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy")
+image_processor_1 = AutoImageProcessor.from_pretrained("haywoodsloan/ai-image-detector-deploy", use_fast=True)
 model_1 = Swinv2ForImageClassification.from_pretrained("haywoodsloan/ai-image-detector-deploy")
 model_1 = model_1.to(device)
 clf_1 = pipeline(model=model_1, task="image-classification", image_processor=image_processor_1, device=device)
@@ -22,6 +26,7 @@ clf_2 = pipeline("image-classification", model=model_2_path)
 class_names_1 = ['artificial', 'real']
 class_names_2 = ['AI Image', 'Real Image']  # Adjust if the second model has different classes
 
+
 def predict_image(img, confidence_threshold):
     # Ensure the image is a PIL Image
     if not isinstance(img, Image.Image):
@@ -67,10 +72,10 @@ def predict_image(img, confidence_threshold):
                 result_2[class_name] = 0.0
         
         # Check if either class meets the confidence threshold
-        if result_2['artificial'] >= confidence_threshold:
-            label_2 = f"Label: artificial, Confidence: {result_2['artificial']:.4f}"
-        elif result_2['real'] >= confidence_threshold:
-            label_2 = f"Label: real, Confidence: {result_2['real']:.4f}"
+        if result_2['AI Image'] >= confidence_threshold:
+            label_2 = f"Label: AI Image, Confidence: {result_2['AI Image']:.4f}"
+        elif result_2['Real Image'] >= confidence_threshold:
+            label_2 = f"Label: Real Image, Confidence: {result_2['Real Image']:.4f}"
         else:
             label_2 = "Uncertain Classification"
     except Exception as e:
@@ -94,4 +99,5 @@ gr.Interface(
     inputs=[image, confidence_slider],
     outputs=label,
     title="AI Generated Classification",
+    queue=True  # Enable queuing to handle multiple predictions efficiently
 ).launch()