Abdoulaye coumbassa
openhands
feat(vision): trainable stage-1 vision projector with frozen encoder and LM
ba3bbc7 | """Modular vision adaptation interfaces for CORTEX AI. | |
| The module intentionally does not ship vision weights. A real encoder must be | |
| provided by the caller (for example, a Transformers vision model loaded from a | |
| separate approved checkpoint). The adapter can therefore be tested with a | |
| small local encoder while remaining honest about deployment readiness. | |
| """ | |
| from collections.abc import Mapping | |
| from dataclasses import dataclass | |
| from typing import Any, Protocol, Sequence | |
| import torch | |
| from torch import Tensor, nn | |
| class VisionConfig: | |
| """Configuration for the image-to-language adaptation layer.""" | |
| vision_hidden_size: int | |
| cortex_hidden_size: int = 7168 | |
| projector_type: str = "linear" | |
| image_token_count: int | None = None | |
| status: str = "experimental / in development" | |
| class MultimodalInput: | |
| """A language prompt plus projected image tokens for a future fusion path.""" | |
| input_ids: Tensor | None | |
| image_embeddings: Tensor | |
| attention_mask: Tensor | None = None | |
| metadata: dict[str, Any] | None = None | |
| class VisionEncoder(Protocol): | |
| """Minimal callable contract for an image encoder.""" | |
| def __call__(self, images: Any) -> Tensor: ... | |
| class ImageProcessor(Protocol): | |
| """Minimal preprocessing contract, compatible with HF processors.""" | |
| def __call__(self, images: Any, **kwargs: Any) -> Any: ... | |
| class VisionProjector(nn.Module): | |
| """Map encoder features into CORTEX hidden space. | |
| ``projector_type='linear'`` is a single linear map. ``projector_type='mlp'`` is the | |
| LLaVA-style two-layer gated projection, whose intermediate width equals the target | |
| hidden size. Only this module is trained in the first experimental stage. | |
| """ | |
| def __init__(self, config: VisionConfig) -> None: | |
| super().__init__() | |
| if config.projector_type == "linear": | |
| self.projection: nn.Module = nn.Linear( | |
| config.vision_hidden_size, config.cortex_hidden_size | |
| ) | |
| elif config.projector_type == "mlp": | |
| self.projection = nn.Sequential( | |
| nn.Linear(config.vision_hidden_size, config.cortex_hidden_size), | |
| nn.GELU(), | |
| nn.Linear(config.cortex_hidden_size, config.cortex_hidden_size), | |
| ) | |
| else: | |
| raise ValueError( | |
| f"Unknown projector_type={config.projector_type!r}; " | |
| "expected 'linear' or 'mlp'" | |
| ) | |
| self.config = config | |
| def forward(self, features: Tensor) -> Tensor: | |
| if features.ndim != 3: | |
| raise ValueError("vision features must have shape [batch, tokens, hidden]") | |
| return self.projection(features) | |
| class VisionAdapter(nn.Module): | |
| """Image -> encoder -> projector, without modifying the CORTEX core.""" | |
| def __init__(self, encoder: VisionEncoder, projector: VisionProjector, | |
| processor: ImageProcessor | None = None) -> None: | |
| super().__init__() | |
| self.encoder = encoder | |
| self.projector = projector | |
| self.processor = processor | |
| def encode_images(self, images: Any) -> Tensor: | |
| model_inputs = self.processor(images, return_tensors="pt") if self.processor else images | |
| # Same BatchFeature caveat as native.py: HF processors return a UserDict, which is | |
| # not a dict subclass, so it must be tested as a Mapping before unpacking by keyword. | |
| if isinstance(model_inputs, Mapping): | |
| features = self.encoder(**model_inputs) | |
| else: | |
| features = self.encoder(model_inputs) | |
| if hasattr(features, "last_hidden_state"): | |
| features = features.last_hidden_state | |
| if not isinstance(features, Tensor): | |
| raise TypeError("vision encoder must return a Tensor or an object with last_hidden_state") | |
| return self.projector(features) | |
| def prepare_multimodal_input(self, images: Any, input_ids: Tensor | None = None, | |
| attention_mask: Tensor | None = None) -> MultimodalInput: | |
| projected = self.encode_images(images) | |
| return MultimodalInput(input_ids, projected, attention_mask, | |
| metadata={"vision_status": "experimental / in development"}) | |
| class MultimodalFusion(Protocol): | |
| """Future contract for inserting projected image tokens into the LM context.""" | |
| def fuse(self, text_embeddings: Tensor, image_embeddings: Tensor) -> Tensor: ... | |
| class PrefixFusion(nn.Module): | |
| """Reference fusion strategy: prepend image tokens to text embeddings.""" | |
| def fuse(self, text_embeddings: Tensor, image_embeddings: Tensor) -> Tensor: | |
| if text_embeddings.ndim != 3 or image_embeddings.ndim != 3: | |
| raise ValueError("text and image embeddings must be rank-3 tensors") | |
| if text_embeddings.shape[0] != image_embeddings.shape[0]: | |
| raise ValueError("text and image batches must match") | |
| return torch.cat((image_embeddings, text_embeddings), dim=1) | |
| def build_vision_adapter(config: VisionConfig, encoder: VisionEncoder, | |
| processor: ImageProcessor | None = None) -> VisionAdapter: | |
| """Build the adaptation layer around an externally supplied encoder.""" | |
| return VisionAdapter(encoder, VisionProjector(config), processor) | |