Cogvit

CogViT — Pytorch

Implementation of CogViT, the parameter-efficient vision encoder from GLM-5V-Turbo: Toward a Native Foundation Model for Multimodal Agents, in Pytorch. A single-file, self-contained reference build of the §2.1 vision tower, the MLP adapter that bridges into the language backbone, and both pretraining stages (DMIM distillation and SigLIP contrastive).

The encoder is a fairly standard ViT spine with three modern touches: QK-Norm for stable attention at scale, learned 2D position embeddings that are bicubically interpolated to arbitrary patch grids (the NaFlex path), and a learnable mask token that lets the same encoder serve both stage-1 masked image modeling and stage-2 contrastive pretraining without architectural surgery.

Install

$ pip install -r requirements.txt

Usage

The encoder by itself:

import torch
from cogvit import CogViT, CogViTConfig

cfg = CogViTConfig(
    image_size = 224,
    patch_size = 14,
    embed_dim  = 1024,
    depth      = 24,
    num_heads  = 16,
)

vit = CogViT(cfg)

images = torch.randn(2, 3, 224, 224)
out = vit(images)

out["patch_tokens"]   # (2, 256, 1024) — Hp*Wp patch tokens
out["cls_token"]      # (2, 1024)
out["grid"]           # (16, 16)

Variable resolution (NaFlex)

Images of arbitrary aspect ratio can be packed into a single batch under a fixed token budget. The encoder consumes a per-patch `valid_mask` so attention ignores padding patches.

from cogvit import CogViT, CogViTConfig, naflex_collate

cfg = CogViTConfig(naflex_max_tokens = 1024)
vit = CogViT(cfg)

images = [torch.randn(3, 480, 640), torch.randn(3, 720, 480)]
batch, valid_mask = naflex_collate(images, patch_size=cfg.patch_size, max_tokens=1024)

out = vit(batch, valid_mask=valid_mask)

Bridging i