Spaces:
Sleeping
Sleeping
File size: 2,417 Bytes
6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c e1b76ea 6180f4c | 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 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import logging
import threading
from contextlib import asynccontextmanager
import uvicorn
from fastapi import FastAPI, HTTPException
from pydantic import BaseModel
try:
from inference import generate_text, generate_with_image, health_check, load_model
except ModuleNotFoundError:
from .inference import generate_text, generate_with_image, health_check, load_model
logger = logging.getLogger(__name__)
def _preload_model() -> None:
"""Load the model in the background at startup so the first inference request
doesn't pay the full cold-start cost."""
try:
load_model()
logger.info("Model preloaded at startup")
except Exception:
logger.exception("Background model preload failed — will retry on first request")
@asynccontextmanager
async def lifespan(app: FastAPI):
threading.Thread(target=_preload_model, daemon=True, name="model-preloader").start()
yield
app = FastAPI(lifespan=lifespan)
class GenerateRequest(BaseModel):
prompt: str
model: str = "e4b"
max_tokens: int = 512
class VisionRequest(BaseModel):
prompt: str
image_base64: str
model: str = "e4b"
max_tokens: int = 512
@app.get("/")
def root():
return {"message": "KnowLedge Inference API is running"}
@app.get("/api/health")
def health():
# health_check() is now non-blocking — just reports current load state
return health_check()
@app.post("/api/generate")
def generate(req: GenerateRequest):
model_name = "e2b" if "e2b" in req.model.lower() else "e4b"
try:
response = generate_text(model_name, req.prompt, req.max_tokens)
return {"response": response}
except RuntimeError as exc:
raise HTTPException(status_code=503, detail=str(exc))
except Exception as exc:
logger.exception("Text generation error")
raise HTTPException(status_code=500, detail=str(exc))
@app.post("/api/generate_vision")
def generate_vision(req: VisionRequest):
try:
response = generate_with_image(req.prompt, req.image_base64, req.max_tokens)
return {"response": response}
except RuntimeError as exc:
raise HTTPException(status_code=503, detail=str(exc))
except Exception as exc:
logger.exception("Vision generation error")
raise HTTPException(status_code=500, detail=str(exc))
if __name__ == "__main__":
uvicorn.run(app, host="0.0.0.0", port=7860)
|