from fastapi import FastAPI, File, UploadFile, Form
from fastapi.responses import JSONResponse
from transformers import pipeline
from io import BytesIO
from PIL import Image

app = FastAPI()

# Load pipeline once
pipe = pipeline("image-text-to-text", model="google/medgemma-27b-it")

@app.get("/")
def root():
    return {"message": "MedGemma Image-to-Text API running"}

@app.post("/analyze/")
async def analyze_image(prompt: str = Form(...), file: UploadFile = File(...)):
    try:
        image_data = await file.read()
        image = Image.open(BytesIO(image_data)).convert("RGB")

        response = pipe([
            {
                "role": "user",
                "content": [
                    {"type": "image", "image": image},
                    {"type": "text", "text": prompt}
                ]
            }
        ])

        return {"result": response[0]['generated_text']}

    except Exception as e:
        return JSONResponse(status_code=500, content={"error": str(e)})