import gradio as gr
from transformers import AutoTokenizer, AutoModelForImageClassification
from PIL import Image
import requests
import torch

# Load model from Hugging Face model hub
model_name = "your-username/your-model-name"  # Replace with your model's name on Hugging Face
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

# Define function for image preprocessing and prediction
def process_image(image):
    # Load and preprocess image
    image = Image.open(image)
    inputs = tokenizer(image, return_tensors="pt", padding=True, truncation=True)
    # Make prediction
    outputs = model(**inputs)
    predicted_class = torch.argmax(outputs.logits, dim=1)
    return predicted_class.item()

# Create Gradio interface
inputs = gr.inputs.Image()
output = gr.outputs.Label(num_top_classes=1)
interface = gr.Interface(fn=process_image, inputs=inputs, outputs=output, capture_session=True)

# Deploy the Gradio interface
interface.launch(share=True)