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)