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)})