request_id
parent
cb1294bd4d
commit
a2b6730b88
|
@ -0,0 +1,57 @@
|
|||
import os
|
||||
import io
|
||||
from dotenv import load_dotenv
|
||||
from fastapi import FastAPI, File, UploadFile, Form, Query
|
||||
from fastapi.responses import JSONResponse, StreamingResponse
|
||||
from PIL import Image
|
||||
import moondream as md
|
||||
|
||||
# Load environment variables
|
||||
load_dotenv()
|
||||
api_key = os.getenv("MOON_DREAM_KEY")
|
||||
|
||||
# Initialize Moondream model
|
||||
model = md.vl(api_key=api_key)
|
||||
|
||||
# FastAPI app
|
||||
app = FastAPI()
|
||||
|
||||
@app.post("/caption")
|
||||
async def generate_caption(
|
||||
image: UploadFile = File(...),
|
||||
length: str = Form("short")
|
||||
):
|
||||
try:
|
||||
image_bytes = await image.read()
|
||||
img = Image.open(io.BytesIO(image_bytes))
|
||||
response = model.caption(img, length=length)
|
||||
return JSONResponse(content={"caption": response["caption"]})
|
||||
except Exception as e:
|
||||
return JSONResponse(status_code=500, content={"error": str(e)})
|
||||
|
||||
@app.post("/query")
|
||||
async def query_image(
|
||||
image: UploadFile = File(...),
|
||||
question: str = Form(...),
|
||||
stream: bool = Form(False)
|
||||
):
|
||||
try:
|
||||
image_bytes = await image.read()
|
||||
img = Image.open(io.BytesIO(image_bytes))
|
||||
|
||||
if stream:
|
||||
def generate():
|
||||
result = model.query(img, question, stream=True)
|
||||
for chunk in result["chunk"]:
|
||||
yield chunk
|
||||
|
||||
return StreamingResponse(generate(), media_type="text/plain")
|
||||
else:
|
||||
result = model.query(img, question)
|
||||
return JSONResponse(content={
|
||||
"answer": result["answer"],
|
||||
"request_id": result["request_id"]
|
||||
})
|
||||
|
||||
except Exception as e:
|
||||
return JSONResponse(status_code=500, content={"error": str(e)})
|
Loading…
Reference in New Issue