cortex_ai/vision/native.py

"""Native image-embedding fusion path for CORTEX-compatible language cores.

This module performs the real tensor path IMAGE -> encoder -> projector ->
inputs_embeds -> language core. It is intentionally model-agnostic: the
existing 1.65T DeepSeek checkpoint is not modified, and a core is only
accepted when its runtime exposes the required inputs_embeds contract.
"""

from collections.abc import Mapping
from dataclasses import dataclass
from typing import Any, Callable

import torch
from torch import Tensor, nn

from .core import PrefixFusion, VisionConfig, VisionProjector


@dataclass(frozen=True)
class NativeFusionResult:
    """Output of a native fusion forward pass."""

    outputs: Any
    image_token_count: int
    text_token_count: int


class NativeVisionFusion(nn.Module):
    """Fuse projected image tokens into a language core via ``inputs_embeds``."""

    def __init__(
        self,
        vision_encoder: Callable[..., Any],
        config: VisionConfig,
        core: nn.Module,
        image_processor: Callable[..., Any] | None = None,
    ) -> None:
        super().__init__()
        self.vision_encoder = vision_encoder
        self.projector = VisionProjector(config)
        self.core = core
        self.image_processor = image_processor
        self.fusion = PrefixFusion()

    def encode_image(self, image: Any) -> Tensor:
        model_input = (
            self.image_processor(image, return_tensors="pt")
            if self.image_processor is not None
            else image
        )
        # Transformers processors return BatchFeature, a UserDict rather than a
        # dict subclass, so a plain isinstance(..., dict) check would pass the
        # processor output positionally into the encoder's input_ids parameter.
        if isinstance(model_input, Mapping):
            encoded = self.vision_encoder(**model_input)
        else:
            encoded = self.vision_encoder(model_input)
        return self._project(encoded)

    def encode_pixels(self, pixel_values: Tensor, **encoder_kwargs: Any) -> Tensor:
        """Project already-preprocessed ``pixel_values`` into the core's hidden space.

        Training loops preprocess whole batches once, so they use this instead of
        ``encode_image`` to avoid re-running the image processor on every step. The encoder is
        called by keyword here, which suits Transformers vision towers.
        """
        return self._project(self.vision_encoder(pixel_values=pixel_values, **encoder_kwargs))

    def _project(self, encoded: Any) -> Tensor:
        if hasattr(encoded, "last_hidden_state"):
            encoded = encoded.last_hidden_state
        if not isinstance(encoded, Tensor) or encoded.ndim != 3:
            raise ValueError("vision_encoder must return [batch, image_tokens, vision_hidden_size]")
        return self.projector(encoded)

    def build_inputs_embeds(self, input_ids: Tensor, image: Any) -> tuple[Tensor, Tensor]:
        if input_ids.ndim != 2:
            raise ValueError("input_ids must have shape [batch, text_tokens]")
        image_embeds = self.encode_image(image)
        token_embedding = getattr(self.core, "get_input_embeddings", None)
        if token_embedding is None or not callable(token_embedding):
            raise TypeError("core must expose get_input_embeddings() for native fusion")
        text_embeds = token_embedding()(input_ids)
        fused = self.fusion.fuse(text_embeds, image_embeds)
        attention_mask = torch.ones(fused.shape[:2], dtype=torch.long, device=fused.device)
        return fused, attention_mask

    def forward(self, input_ids: Tensor, image: Any, **kwargs: Any) -> NativeFusionResult:
        inputs_embeds, attention_mask = self.build_inputs_embeds(input_ids, image)
        try:
            outputs = self.core(inputs_embeds=inputs_embeds, attention_mask=attention_mask, **kwargs)
        except TypeError as exc:
            raise TypeError(
                "core does not accept inputs_embeds; native fusion cannot be claimed for this core"
            ) from exc
        return NativeFusionResult(outputs, inputs_embeds.shape[1] - input_ids.shape[1], input_ids.shape[1])


def build_native_fusion(
    vision_encoder: Callable[..., Any],
    core: nn.Module,
    vision_hidden_size: int,
    cortex_hidden_size: int,
    image_processor: Callable[..., Any] | None = None,
    projector_type: str = "linear",
) -> NativeVisionFusion:
    """Construct the native path around externally supplied encoder and core modules."""
    config = VisionConfig(
        vision_hidden_size=vision_hidden_size,
        cortex_hidden_size=cortex_hidden_size,
        projector_type=projector_type,
    )
    return NativeVisionFusion(vision_encoder, config, core, image_processor)