App v2

app.py

@@ -5,30 +5,29 @@ import gradio as gr
 from transformers import pipeline
 
 
-device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+device = 0 if torch.cuda.is_available() else -1
 
 # Loading in Model
 model_name = "dima806/ai_vs_real_image_detection"
-model = ViTForImageClassification.from_pretrained(model_name).to(device)
-model.to(device)
+pipe = pipeline("image-classification", model=model_name, device = device)
+    
 
 
-#Classification function
+# Classification function
 def classify_image(img: Image.Image):
-    inputs = model(images=img, return_tensors="pt").to(device)
-    results = model(inputs)
+    results = pipe(img)
     top = results[0]
     label = top["label"]
     score = top["score"]
     return f"Prediction: {label} (Confidence: {score:.2f})"
 
-
-
-# Interface
+# Gradio interface
 interface = gr.Interface(
     fn=classify_image,
     inputs=gr.Image(type="pil"),
     outputs="text",
-    title="Real vs AI Image detection",
-    description="Check if your image is Real or AI"
-)
+    title="Real vs AI Image Detection",
+    description="Upload an image to see if it's REAL or AI-generated."
+)
+
+interface.launch()