File size: 1,394 Bytes
5e3465d 4b9df52 5e3465d a4ae355 5e3465d 4b9df52 5e3465d a4ae355 5e3465d 4b9df52 a4ae355 4b9df52 5e3465d 4b9df52 |
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 |
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
from concurrent.futures import ThreadPoolExecutor
import uuid, time
import wai_service
from typing import Optional
app = FastAPI()
EXEC = ThreadPoolExecutor(max_workers=1)
JOBS = {} # job_id ➜ (start_ts, Future)
class InferenceRequest(BaseModel):
word: str
optimized_letter: str
font: str = "KaushanScript-Regular"
seed: int = 0
token: Optional[str] = None
class Config:
extra = "allow"
@app.post("/generate")
def enqueue(req: InferenceRequest):
"""Start a job and return its UUID immediately."""
if req.token is None:
raise HTTPException(422, "field 'token' is required")
payload = req.dict(exclude_none=True)
job_id = str(uuid.uuid4())
fut = EXEC.submit(wai_service.handler, dict(req.__dict__))
JOBS[job_id] = (time.time(), fut)
return {"job_id": job_id}
@app.get("/result/{job_id}")
def get_result(job_id: str):
"""Poll for job completion."""
if job_id not in JOBS:
raise HTTPException(404, "job_id not found")
start, fut = JOBS[job_id]
if not fut.done():
return {"status": "running", "elapsed": int(time.time() - start)}
try:
img = fut.result()
return {"status": "finished", "image_base64": img}
finally:
del JOBS[job_id] # housekeeping
|