captcha_solver/engines/florence_engine.py

"""Florence-2 engine for OCR and object detection.

Microsoft's Florence-2 is a strong small vision model. We use the 'base'
variant (~1.2GB) which fits on CPU and supports:
  - <OCR>     : plain OCR
  - <CAPTION> : short caption
  - <OD>      : object detection with boxes
  - <CAPTION_TO_PHRASE_GROUNDING> : grounding
"""

from __future__ import annotations

import os
from typing import Optional

from captcha_solver.engines.base import BaseEngine
from captcha_solver.config import get_settings


class FlorenceEngine(BaseEngine):
    name = "florence2"

    def __init__(self) -> None:
        super().__init__()
        self._model = None
        self._processor = None

    def _do_load(self) -> None:
        import torch
        from transformers import AutoModelForCausalLM, AutoProcessor

        s = get_settings()
        os.environ.setdefault("HF_HOME", str(s.cache_dir / "hf"))
        dtype = self._resolve_dtype(s.florence_torch_dtype)

        self._processor = AutoProcessor.from_pretrained(
            s.florence_model, trust_remote_code=True
        )
        self._model = AutoModelForCausalLM.from_pretrained(
            s.florence_model,
            trust_remote_code=True,
            torch_dtype=dtype,
            low_cpu_mem_usage=True,
        ).to(s.florence_device).eval()

    def _do_unload(self) -> None:
        self._model = None
        self._processor = None

    @staticmethod
    def _resolve_dtype(name: str):
        import torch
        return {"float32": torch.float32, "float16": torch.float16, "bfloat16": torch.bfloat16}[name]

    def ocr(self, pil_image, task: str = "<OCR>") -> str:
        """Run an OCR-style task and return the text. Tasks include <OCR>, <OCR_WITH_REGION>."""
        if not self._loaded:
            self.load()
        import torch

        assert self._model is not None and self._processor is not None
        prompt = task
        inputs = self._processor(text=prompt, images=pil_image, return_tensors="pt").to(
            self._model.device
        )
        with torch.no_grad():
            gen = self._model.generate(
                input_ids=inputs["input_ids"],
                pixel_values=inputs["pixel_values"].to(self._model.dtype),
                max_new_tokens=256,
                num_beams=3,
                do_sample=False,
            )
        text = self._processor.batch_decode(gen, skip_special_tokens=False)[0]
        parsed = self._processor.post_process_generation(
            text, task=prompt, image_size=(pil_image.width, pil_image.height)
        )
        value = parsed.get(prompt, "")
        if isinstance(value, dict):
            return " ".join(str(v) for v in value.values())
        return str(value)