vision/engine.py

"""Vision engine — image analysis, OCR, captioning via ZeroGPU."""

import logging
from typing import Optional

from config import config
from models.zerogpu import requires_gpu

logger = logging.getLogger("synapse.vision")


class VisionEngine:
    """Image understanding and document analysis."""

    def __init__(self):
        self._florence_model = None
        self._qwen_model = None

    def _load_florence(self):
        if self._florence_model is not None:
            return
        try:
            import torch
            from transformers import AutoModelForCausalLM, AutoProcessor
            device = "cuda" if torch.cuda.is_available() else "cpu"
            dtype = torch.float16 if device == "cuda" else torch.float32

            model_id = config.vision.florence_model
            self._florence_model = {
                "model": AutoModelForCausalLM.from_pretrained(
                    model_id, torch_dtype=dtype, trust_remote_code=True,
                ).to(device),
                "processor": AutoProcessor.from_pretrained(
                    model_id, trust_remote_code=True,
                ),
                "device": device,
            }
            logger.info("Florence model loaded")
        except Exception as e:
            logger.warning(f"Failed to load Florence: {e}")

    @requires_gpu("vision")
    def describe_image(self, image_path: str,
                       task: str = "caption") -> str:
        """Analyze an image with Florence-2."""
        self._load_florence()
        if self._florence_model is None:
            return self._describe_fallback(image_path, task)

        try:
            from PIL import Image
            import torch

            image = Image.open(image_path).convert("RGB")
            model = self._florence_model["model"]
            processor = self._florence_model["processor"]
            device = self._florence_model["device"]

            task_prompts = {
                "caption": "<CAPTION>",
                "detailed": "<DETAILED_CAPTION>",
                "ocr": "<OCR>",
                "objects": "<OD>",
                "regions": "<REGION_PROPOSAL>",
            }
            prompt = task_prompts.get(task, "<CAPTION>")

            inputs = processor(text=prompt, images=image, return_tensors="pt").to(device)
            with torch.no_grad():
                generated_ids = model.generate(
                    **inputs, max_new_tokens=256, num_beams=3,
                )
            result = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
            return result.strip()
        except Exception as e:
            logger.error(f"Florence vision error: {e}")
            return self._describe_fallback(image_path, task)

    def _describe_fallback(self, image_path: str, task: str) -> str:
        try:
            from PIL import Image
            img = Image.open(image_path)
            size = img.size
            mode = img.mode
            return f"Image: {size[0]}x{size[1]}, mode={mode}. Detailed analysis unavailable — vision model not loaded."
        except Exception:
            return "Image analysis unavailable."

    def caption(self, image_path: str) -> str:
        return self.describe_image(image_path, "caption")

    def detailed_caption(self, image_path: str) -> str:
        return self.describe_image(image_path, "detailed")

    def ocr(self, image_path: str) -> str:
        return self.describe_image(image_path, "ocr")

    def detect_objects(self, image_path: str) -> str:
        return self.describe_image(image_path, "objects")

    def analyze_for_question(self, image_path: str, question: str) -> str:
        self._load_florence()
        if self._florence_model is None:
            return f"Cannot answer '{question}' — vision model not loaded."

        try:
            from PIL import Image
            import torch

            image = Image.open(image_path).convert("RGB")
            model = self._florence_model["model"]
            processor = self._florence_model["processor"]
            device = self._florence_model["device"]

            inputs = processor(text=question, images=image, return_tensors="pt").to(device)
            with torch.no_grad():
                generated_ids = model.generate(
                    **inputs, max_new_tokens=256, num_beams=3,
                )
            result = processor.batch_decode(generated_ids, skip_special_tokens=False)[0]
            return result.strip()
        except Exception as e:
            return f"Error analyzing image: {e}"

    def get_tasks(self) -> list[dict]:
        return [
            {"id": "caption", "name": "Caption", "description": "Brief image description"},
            {"id": "detailed", "name": "Detailed Caption", "description": "Detailed description"},
            {"id": "ocr", "name": "OCR", "description": "Extract text from image"},
            {"id": "objects", "name": "Object Detection", "description": "Detect objects"},
            {"id": "regions", "name": "Regions", "description": "Detect regions"},
        ]


vision_engine = VisionEngine()