import torch
from PIL import Image
from transformers import BlipProcessor, BlipForConditionalGeneration

# Repo model đã upload lên HF Hub (theo ảnh bạn gửi)
MODEL_REPO = "demo-at-wsld-8/blip-at-wsld-fashion200k"

device = "cuda" if torch.cuda.is_available() else "cpu"

_processor = None
_model = None


def load_model():
    """
    Load processor + model 1 lần duy nhất (cache lại), dùng cho mọi request sau đó.
    model.safetensors trong repo đã là bản distilled cuối cùng -> from_pretrained
    là đủ, không cần đụng tới student_at_wsld.pth.
    """
    global _processor, _model
    if _model is None:
        _model = BlipForConditionalGeneration.from_pretrained(MODEL_REPO).to(device).eval()
        _processor = BlipProcessor.from_pretrained(MODEL_REPO)
    return _processor, _model


def generate_caption(image_path):
    processor, model = load_model()
    image = Image.open(image_path).convert("RGB")
    inputs = processor(images=image, return_tensors="pt").to(device)

    with torch.no_grad():
        out = model.generate(
            **inputs,
            num_beams=1,
            max_new_tokens=80,
            min_new_tokens=6,
            no_repeat_ngram_size=3,
            repetition_penalty=1.3,
        )

    caption = processor.batch_decode(out, skip_special_tokens=True)[0]
    return caption