from fastapi import FastAPI, UploadFile
from transformers import pipeline
from PIL import Image
import io
import uvicorn

app = FastAPI()
pipe = pipeline("image-classification", model="google/vit-base-patch16-224")

@app.post("/predict")
async def predict(file: UploadFile):
    img = Image.open(io.BytesIO(await file.read()))
    preds = pipe(img)
    return preds

if __name__ == "__main__":
    uvicorn.run(app, host="0.0.0.0", port=7860)