How to Add a New Custom Vision Encoder to Eagle's Architecture

Adding a custom vision encoder to Eagle requires implementing a tower wrapper class that exposes a standard interface, registering it in the VISION_TOWER_REGISTRY, and updating the CLI configuration to recognize the new encoder name.

Eagle, developed by NVIDIA's NVlabs, employs a modular multimodal design that isolates vision encoding behind lightweight "tower" wrappers. To add a new custom vision encoder to Eagle's architecture, you must implement a wrapper class that standardizes the interface between your pretrained vision model and Eagle's multimodal projector. This approach allows you to integrate any vision-only backbone—whether a custom ViT, CLIP variant, or convolutional network—without modifying the core language model or projection layers.

Step 1 – Implement the Vision Tower Wrapper

Eagle's architecture expects vision encoders to be wrapped in a tower class that handles lazy loading, preprocessing, and API uniformity. The reference implementation, SAMVisionTower, resides in Eagle/eagle/model/multimodal_encoder/sam_encoder.py (lines 76-106). Your custom wrapper must inherit from nn.Module and implement specific methods and properties to integrate seamlessly with the multimodal projector in Eagle/eagle/model/multimodal_projector/builder.py.

The wrapper must lazy-load the pretrained model via a load_model() method, store a processor for image transformations, and expose a forward method returning a tensor of shape (batch, tokens, dim).

# file: Eagle/eagle/model/multimodal_encoder/my_encoder.py

import torch
from torch import nn
from transformers import AutoProcessor, AutoModel

class MyVisionTower(nn.Module):
    def __init__(self, vision_tower, args):
        super().__init__()
        self.args = args
        self.is_loaded = False
        self.vision_tower_name = vision_tower
        self.pixel_shuffle = getattr(args, "add_pixel_shuffle", False)

    def load_model(self):
        if self.is_loaded:
            return
        # Load processor & model

        self.processor = AutoProcessor.from_pretrained(self.vision_tower_name)
        model = AutoModel.from_pretrained(self.vision_tower_name).vision_encoder
        self.vision_tower = model
        if self.args.freeze_vision:
            self.vision_tower.requires_grad_(False)
        self.is_loaded = True

    def forward(self, images):
        # Accept either a list of PIL images or a batch tensor

        if isinstance(images, list):
            inputs = self.processor(images, return_tensors="pt")
            images = inputs["pixel_values"].to(self.device)
        else:
            images = images.to(self.device)

        feats = self.vision_tower(images).last_hidden_state  # (B, N, C)

        # Optional pixel-shuffle handling

        if self.pixel_shuffle:
            b, n, c = feats.shape
            h = w = int(n ** 0.5)
            feats = feats.transpose(1, 2).reshape(b, c, h, w)
            feats = nn.functional.pixel_unshuffle(feats, 2)
        return feats

    @property
    def dummy_feature(self):
        return torch.zeros(1, self.hidden_size, device=self.device, dtype=self.dtype)

    @property
    def dtype(self):
        return next(self.vision_tower.parameters()).dtype

    @property
    def device(self):
        return next(self.vision_tower.parameters()).device

    @property
    def config(self):
        return self.vision_tower.config

    @property
    def hidden_size(self):
        return self.vision_tower.config.hidden_size

    @property
    def num_patches(self):
        return self.config.num_patches

Required Properties and Methods

Your wrapper must expose the following properties to maintain compatibility with Eagle's projector and trainer:

• dummy_feature: Returns a zero tensor of shape (1, hidden_size) for placeholder operations.
• dtype and device: Infer data type and device from the loaded model parameters.
• config: Exposes the underlying model configuration.
• hidden_size: Returns the feature dimension (adjust if you add a projection neck).
• num_patches: Indicates the number of visual tokens output per image.

Handling Pixel Shuffle

If your encoder produces high-resolution feature maps that require downsampling before projection, implement the pixel_shuffle logic shown in the example above. This reshapes the sequence output into a spatial grid and applies pixel_unshuffle to reduce spatial dimensions while increasing channel depth, matching the pattern used in SAMVisionTower.

Step 2 – Register the Encoder in the Vision Tower Registry

Eagle selects vision towers dynamically using a registry pattern. Locate the VISION_TOWER_REGISTRY dictionary in Eagle/eagle/model/multimodal_encoder/__init__.py and add an entry mapping your encoder's string identifier to your wrapper class.

# file: Eagle/eagle/model/multimodal_encoder/__init__.py

from .sam_encoder import SAMVisionTower
from .my_encoder import MyVisionTower

VISION_TOWER_REGISTRY = {
    "sam-vit-large": SAMVisionTower,
    "my-custom-encoder": MyVisionTower,   # new entry

}

The model builder imports this registry to instantiate the correct tower based on configuration. Ensure your string key is unique and descriptive.

Step 3 – Update Configuration and CLI Arguments

The training script Eagle/train.py parses the --vision_tower argument to determine which encoder to load. Update the argument parser to include your new encoder name in help text and validation logic.

parser.add_argument(
    "--vision_tower",
    type=str,
    default="sam-vit-large",
    help="Name of the vision encoder tower (e.g., 'sam-vit-large' or 'my-custom-encoder')",
)

If the codebase includes validation assertions (e.g., assert args.vision_tower in VISION_TOWER_REGISTRY), ensure they check against the imported registry to avoid hardcoded lists.

Verification and Integration Testing

After implementation, verify that your encoder integrates correctly with Eagle's multimodal stack. Instantiate a minimal model with --vision_tower my-custom-encoder and pass a dummy image tensor through the tower to confirm the output shape matches the projector's expected input dimensions. For reference on attention mechanisms and patch handling, consult the MoonViT implementation in Embodied/eaglevl/model/moon_vit/modeling_vit.py (lines 82-106).

Summary

• Implement a tower wrapper in Eagle/eagle/model/multimodal_encoder/ that provides load_model, forward, and standard properties (hidden_size, num_patches, dummy_feature).
• Register the class in VISION_TOWER_REGISTRY within Eagle/eagle/model/multimodal_encoder/__init__.py using a unique string key.
• Update CLI arguments in Eagle/train.py to recognize the new encoder name and ensure validation logic references the registry.
• Maintain interface compatibility to ensure seamless integration with Eagle/eagle/model/multimodal_projector/builder.py without modifying core multimodal logic.

Frequently Asked Questions

What interface must a custom vision encoder implement for Eagle?

Your encoder must expose load_model() for lazy initialization, forward(images) returning a tensor of shape (batch, tokens, dim), and properties including hidden_size, num_patches, dummy_feature, dtype, device, and config. These properties allow Eagle's projector and trainer to query model capabilities without accessing internal implementation details.

Where does Eagle instantiate the vision tower during model initialization?

The tower instantiation occurs through the registry lookup in the encoder builder. When the model configuration specifies a vision_tower string, Eagle queries VISION_TOWER_REGISTRY to retrieve the corresponding class (such as SAMVisionTower or your custom wrapper), then instantiates it with the model name and training arguments.

Can I use a vision encoder with different output dimensions than the default SAM tower?

Yes, provided you correctly set hidden_size and num_patches to reflect your encoder's output dimensions. Ensure the multimodal projector in builder.py can accommodate the token count, or enable pixel_shuffle in your wrapper to reduce spatial dimensions before projection.

How do I freeze the vision encoder weights during training?

In your tower's load_model() method, check self.args.freeze_vision and call self.vision_tower.requires_grad_(False) when loading the pretrained weights. This pattern matches the implementation in the reference SAMVisionTower and prevents gradient updates to the vision backbone while training the projection and language layers.

{
  "mcpServers": {
    "instagit": {
      "command": "npx",
      "args": ["-y", "instagit@latest"]
    }
  }
}