from fastapi import FastAPI, File, UploadFile
from fastapi.responses import JSONResponse
from PIL import Image
from transformers import Qwen2VLForConditionalGeneration, AutoProcessor
from qwen_vl_utils import process_vision_info
import torch
import io
import os

app = FastAPI()

model_name = "NAMAA-Space/Qari-OCR-0.2.2.1-VL-2B-Instruct"

# ✅ CPU ONLY
model = Qwen2VLForConditionalGeneration.from_pretrained(
    model_name,
    device_map="cpu",
    torch_dtype=torch.float32
)

processor = AutoProcessor.from_pretrained(model_name)


@app.get("/")
def home():
    return {"message": "OCR API Running"}


@app.post("/ocr")
async def ocr_endpoint(file: UploadFile = File(...)):

    contents = await file.read()
    image = Image.open(io.BytesIO(contents))

    src = "temp_image.png"
    image.save(src)

    prompt = "Extract all text accurately."

    messages = [
        {
            "role": "user",
            "content": [
                {"type": "image", "image": f"file://{src}"},
                {"type": "text", "text": prompt},
            ],
        }
    ]

    text = processor.apply_chat_template(
        messages,
        tokenize=False,
        add_generation_prompt=True
    )

    image_inputs, video_inputs = process_vision_info(messages)

    inputs = processor(
        text=[text],
        images=image_inputs,
        videos=video_inputs,
        padding=True,
        return_tensors="pt"
    )

    # ❌ removed: inputs.to("cuda")

    with torch.no_grad():
        generated_ids = model.generate(
            **inputs,
            max_new_tokens=500
        )

    generated_ids_trimmed = [
        out_ids[len(in_ids):]
        for in_ids, out_ids in zip(inputs.input_ids, generated_ids)
    ]

    output_text = processor.batch_decode(
        generated_ids_trimmed,
        skip_special_tokens=True
    )[0]

    os.remove(src)

    return JSONResponse(content={"text": output_text})