src/server.py

from __future__ import annotations

import os
from typing import Any

from fastapi import FastAPI, File, Form, HTTPException, UploadFile
from fastapi.middleware.cors import CORSMiddleware

from .model import MODEL_ID, DEFAULT_MAX_TOKENS, generate_caption, load_model

app = FastAPI(
    title="img3txt - Florence-2 API",
    description="Simple image-to-text endpoint powered by Florence-2-base.",
    version="1.0.0",
)

app.add_middleware(
    CORSMiddleware,
    allow_origins=["*"],
    allow_credentials=True,
    allow_methods=["*"],
    allow_headers=["*"],
)


@app.on_event("startup")
def warmup_model() -> None:
    load_model()




@app.get("/health")
def health() -> dict[str, Any]:
    return {"status": "ok", "model": MODEL_ID}


@app.post("/predict")
async def predict(
    file: UploadFile = File(...),
    text: str | None = Form(None),
    max_tokens: int = Form(DEFAULT_MAX_TOKENS),
) -> dict[str, Any]:
    image_bytes = await file.read()
    if not image_bytes:
        raise HTTPException(status_code=400, detail="Empty file uploaded")

    try:
        result = generate_caption(image_bytes, text, max_tokens)
    except ValueError as exc:
        raise HTTPException(status_code=400, detail=str(exc)) from exc
    return {"result": result}


if __name__ == "__main__":
    import uvicorn

    port = int(os.getenv("PORT", "7860"))
    uvicorn.run("src.server:app", host="0.0.0.0", port=port)