Spaces:
Sleeping
Sleeping
File size: 1,791 Bytes
03767fe 91cb1e2 30b7070 03767fe 30b7070 91cb1e2 30b7070 91cb1e2 03767fe 30b7070 7c4ca02 30b7070 03767fe 91cb1e2 30b7070 91cb1e2 03767fe 30b7070 7c4ca02 30b7070 aae7362 7c4ca02 30b7070 aae7362 7c4ca02 30b7070 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 |
from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse, HTMLResponse
from transformers import pipeline
from PIL import Image
import io, os, traceback
# Make sure Hugging Face cache is writable
os.environ["HF_HOME"] = "/app/cache"
os.environ["TRANSFORMERS_CACHE"] = "/app/cache"
os.environ["HF_HUB_CACHE"] = "/app/cache/hub"
app = FastAPI()
# Load SmolVLM with the pipeline API
pipe = pipeline(
"image-to-text",
model="HuggingFaceTB/SmolVLM-256M-Instruct",
device=-1 # CPU for free tier
)
@app.get("/")
def home():
return {
"message": "API is running. Use POST /predict with an image, or visit /upload to test in browser."
}
@app.get("/upload", response_class=HTMLResponse)
def upload_form():
return """
<html>
<body>
<h2>Upload an ID Image</h2>
<form action="/predict" enctype="multipart/form-data" method="post">
<input name="file" type="file">
<input type="submit" value="Upload">
</form>
</body>
</html>
"""
@app.post("/predict")
async def predict_gender(file: UploadFile = File(...)):
try:
# Read uploaded image
image_bytes = await file.read()
image = Image.open(io.BytesIO(image_bytes)).convert("RGB")
# Instruction for the model
prompt = "Is the person on this ID male or female?"
# Run model (pipeline handles image + prompt via generate_kwargs)
result = pipe(image, generate_kwargs={"max_new_tokens": 32, "prompt": prompt})
# Extract model output
answer = result[0]["generated_text"].strip()
return JSONResponse({"gender": answer})
except Exception as e:
traceback.print_exc()
return JSONResponse({"error": str(e)}, status_code=500)
|