from pathlib import Path

import numpy as np
import openvino as ov
import openvino_genai as ov_genai
from PIL import Image


MODEL_DIR = Path("models/medgemma-1.5-4b-it-int4-ov")
IMAGE_PATH = Path("images/test.png")
DEVICE = "CPU"  # 有 Intel 核显/独显也可以试 "GPU"


def main() -> None:
    if not MODEL_DIR.exists():
        raise FileNotFoundError(f"Model directory not found: {MODEL_DIR.resolve()}")

    if not IMAGE_PATH.exists():
        raise FileNotFoundError(f"Image not found: {IMAGE_PATH.resolve()}")

    image = Image.open(IMAGE_PATH).convert("RGB")
    image_tensor = ov.Tensor(np.array(image)[None])

    pipe = ov_genai.VLMPipeline(str(MODEL_DIR), DEVICE)

    prompt = "<start_of_image> Describe this medical image in a concise way."

    result = pipe.generate(
        prompt,
        images=[image_tensor],
        max_new_tokens=256,
    )

    if hasattr(result, "texts"):
        print(result.texts[0])
    else:
        print(result)


if __name__ == "__main__":
    main()