import gradio as gr
import torch
from transformers import AutoFeatureExtractor, AutoModelForImageClassification

def aiornot(img):    
    labels = ["Real", "AI"]
    feature_extractor = AutoFeatureExtractor.from_pretrained("Nahrawy/AIorNot")
    model = AutoModelForImageClassification.from_pretrained("Nahrawy/AIorNot")
    
    input = feature_extractor(image, return_tensors="pt")
    with torch.no_grad():
      outputs = model(**input)
      logits = outputs.logits
    prediction = logits.argmax(-1).item()
    label = labels[prediction] 
    return label

with gr.Blocks() as app:
    with gr.Row():
        inp = gr.Image()
        btn = gr.Button()
    outp = gr.Textbox()
    btn.click(aiornot,inp,outp)
app.launch()