app.py

#                                                                                                                 
# gradio demo                                                                                                     
#                                                                                                                 
import gradio as gr
from transformers import ViTFeatureExtractor, ViTModel

feature_extractor = ViTFeatureExtractor.from_pretrained('google/vit-base-patch16-224-in21k')
model = ViTModel.from_pretrained('akahana/vit-base-cats-vs-dogs')

title = "Sandbox"
description = "Place to try various models"

def classify(image):
    inputs = feature_extractor(images=image, return_tensors="pt")
    outputs = model(**inputs)
    return "dog"
    
demo = gr.Interface(fn=classify, inputs="image", outputs="text")

demo.launch()