File size: 1,416 Bytes
641b32e
 
 
 
 
 
 
 
afd6ed3
641b32e
 
 
afd6ed3
641b32e
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
afd6ed3
641b32e
 
afd6ed3
 
641b32e
 
 
 
 
 
 
 
49f8ccd
 
 
 
afd6ed3
641b32e
 
 
 
 
 
 
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
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)