34 lines
870 B
Python
34 lines
870 B
Python
import os
|
|
from dotenv import load_dotenv
|
|
from fastapi import FastAPI, File, UploadFile, Form
|
|
from fastapi.responses import JSONResponse
|
|
from PIL import Image
|
|
import moondream as md
|
|
import io
|
|
|
|
# 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:
|
|
# Read the uploaded image
|
|
image_bytes = await image.read()
|
|
img = Image.open(io.BytesIO(image_bytes))
|
|
|
|
# Generate caption
|
|
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)})
|