from transformers import ViTForImageClassification, AutoFeatureExtractor
from PIL import Image
import torch

model = ViTForImageClassification.from_pretrained("SahilSha/SahAi")
feature_extractor = AutoFeatureExtractor.from_pretrained("google/vit-base-patch16-224-in21k")

image = Image.open("test_image.jpg").convert("RGB")
inputs = feature_extractor(images=image, return_tensors="pt")

with torch.no_grad():
    outputs = model(**inputs)
    prediction = outputs.logits.argmax(-1).item()

print("Real" if prediction == 0 else "Fake")