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

# Load the model and processor
model_name = "DeathDaDev/Materializer"
processor = AutoProcessor.from_pretrained(model_name)
model = AutoModelForImageClassification.from_pretrained(model_name)

# Define the prediction function
def classify_image(image):
    # Preprocess the image
    inputs = processor(images=image, return_tensors="pt")
    
    # Perform inference
    with torch.no_grad():
        logits = model(**inputs).logits
    
    # Get the predicted class
    predicted_class_idx = logits.argmax(-1).item()
    return model.config.id2label[predicted_class_idx]

# Create the Gradio interface
iface = gr.Interface(
    fn=classify_image,
    inputs=gr.inputs.Image(type="pil"),
    outputs=gr.outputs.Label(num_top_classes=3),
    title="Image Classification with Materializer",
    description="Upload an image to classify it using the Materializer model."
)

# Launch the interface
iface.launch()