Support Transformers 4 and 5 in Cosmos-Embed1-336p remote code

README.md

@@ -206,21 +206,45 @@ One can optionally install Transformer Engine for faster inference:
 pip install --no-build-isolation transformer_engine[pytorch]
 ```
 
+The example below also uses `decord` for video decoding and `wget` to fetch a
+sample clip. Install `decord` separately and ensure `wget` is available.
+
+### Transformers compatibility
+
+The remote model code has been validated with Transformers 4.57.6 and 5.17.0.
+Transformers 4 loads positional embeddings through a state-dict pre-hook;
+Transformers 5 uses a registered weight conversion operation instead. This
+matters when loading the 336p and 448p checkpoints: they store positional
+embeddings on a 224p patch grid, which must be interpolated to the model's
+target resolution. Do not use `ignore_mismatched_sizes=True` to work around a
+shape mismatch; that reinitializes the positional embeddings.
+
+Community validation with the same checkpoints and fixed inputs found exact
+text/video-vector parity with the original Transformers 4.57.6 code on CPU at
+all three resolutions. Transformers 5.17.0 was also compared against that
+baseline on CPU, Apple MPS, and NVIDIA CUDA. These comparisons establish
+numerical parity for the tested inputs; they do not change the supported
+platform statement above or establish a performance or retrieval-accuracy
+guarantee.
+
 ## Example inference
 
-A code snippet for video and text inference is shown below.
-For a step-by-step guide, please refer to the Juypter notebook [here](https://huggingface.co/nvidia/Cosmos-Embed1-336p/blob/main/examples/example.ipynb).
-```
+This example keeps model weights in FP32 and runs on CUDA. For a step-by-step
+guide, see the [example notebook](https://huggingface.co/nvidia/Cosmos-Embed1-336p/blob/main/examples/example.ipynb).
+
+```python
 import decord
 import numpy as np
 import torch
 from transformers import AutoProcessor, AutoModel
 import subprocess
-import io
 
-# load model and pre-processor
-model = AutoModel.from_pretrained("nvidia/Cosmos-Embed1-336p", trust_remote_code=True).to("cuda", dtype=torch.bfloat16)
-preprocess = AutoProcessor.from_pretrained("nvidia/Cosmos-Embed1-336p", trust_remote_code=True)
+model_id = "nvidia/Cosmos-Embed1-336p"
+device = "cuda"
+model = AutoModel.from_pretrained(
+    model_id, trust_remote_code=True, dtype=torch.float32,
+).to(device).eval()
+preprocess = AutoProcessor.from_pretrained(model_id, trust_remote_code=True)
 
 # load mock data
 video_url = "https://upload.wikimedia.org/wikipedia/commons/3/3d/Branko_Paukovic%2C_javelin_throw.webm"
@@ -238,15 +262,14 @@ captions = [
     "a man throwing a javelin with both hands", # distractor
 ]
 
-# video and text processing
-video_inputs = preprocess(videos=batch).to("cuda", dtype=torch.bfloat16)
-video_out = model.get_video_embeddings(**video_inputs)
-text_inputs = preprocess(text=captions).to("cuda", dtype=torch.bfloat16)
-text_out = model.get_text_embeddings(**text_inputs)
-
-# ranking and argmax
-probs = (torch.softmax(model.logit_scale.exp() * video_out.visual_proj @ text_out.text_proj.T, dim=-1))[0]
-print(captions[probs.argmax()])
+# Video and text embeddings are L2-normalized by the model.
+video_inputs = preprocess(videos=batch).to(device)
+text_inputs = preprocess(text=captions).to(device)
+with torch.inference_mode():
+    video_out = model.get_video_embeddings(**video_inputs)
+    text_out = model.get_text_embeddings(**text_inputs)
+    scores = (video_out.visual_proj @ text_out.text_proj.T)[0]
+print(captions[scores.argmax().item()])
 ```
 
 # Training and Evaluation

examples/test_transformers_compat.py

@@ -0,0 +1,147 @@
+# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Weight-free compatibility tests for Transformers 4 and 5.
+
+Load standalone remote-code modules without model weights. V5-specific
+conversion tests are skipped when the conversion API is unavailable.
+"""
+
+import importlib.util
+import sys
+from pathlib import Path
+
+import pytest
+import torch
+
+ROOT = Path(__file__).resolve().parent.parent
+
+
+def _load(name: str):
+    spec = importlib.util.spec_from_file_location(name, ROOT / f"{name}.py")
+    module = importlib.util.module_from_spec(spec)
+    sys.modules[name] = module
+    spec.loader.exec_module(module)
+    return module
+
+
+modeling_vit = _load("modeling_vit")
+modeling_qformer = _load("modeling_qformer")
+
+
+def test_qformer_initialization_and_attention_masks() -> None:
+    from transformers import BertConfig
+
+    config = BertConfig(
+        hidden_size=16, num_hidden_layers=1, num_attention_heads=2,
+        intermediate_size=32, vocab_size=64,
+    )
+    model = modeling_qformer.BertModel(config).eval()
+    if modeling_vit._HAS_TRANSFORMERS5_CONVERSION_API:
+        assert hasattr(model, "all_tied_weights_keys")
+    assert model.get_head_mask(None, 1) == [None]
+    mask = model.invert_attention_mask(torch.tensor([[1, 0]]))
+    assert mask.shape == (1, 1, 1, 2)
+    assert mask[0, 0, 0, 0] == 0
+    assert mask[0, 0, 0, 1] < 0
+
+
+def test_pruning_helper_is_vendored() -> None:
+    heads, index = modeling_qformer.find_pruneable_heads_and_indices([0, 1], 4, 2, set())
+    assert heads == {0, 1}
+    assert index.tolist() == [4, 5, 6, 7]
+
+
+def test_bicubic_interpolation_matches_reference() -> None:
+    pos_embed = torch.arange(1 * 257 * 8, dtype=torch.float32).reshape(1, 257, 8)
+    out = modeling_vit._bicubic_interpolate_pos_embed(pos_embed, 1, 24, 24)
+    assert out.shape == (1, 577, 8)
+
+    # Independent reference: same math as the historical interpolate_pos_embed.
+    extra = pos_embed[:, :1]
+    tokens = pos_embed[:, 1:].reshape(-1, 16, 16, 8).permute(0, 3, 1, 2)
+    tokens = torch.nn.functional.interpolate(tokens, size=(24, 24), mode="bicubic", align_corners=False)
+    tokens = tokens.permute(0, 2, 3, 1).flatten(1, 2)
+    reference = torch.cat((extra, tokens), dim=1)
+    assert torch.equal(out, reference)
+    assert torch.equal(out[:, :1], pos_embed[:, :1])  # cls token untouched
+
+
+def test_bicubic_interpolation_is_noop_on_matching_grid() -> None:
+    pos_embed = torch.rand(1, 577, 8)
+    out = modeling_vit._bicubic_interpolate_pos_embed(pos_embed, 1, 24, 24)
+    assert torch.equal(out, pos_embed)
+
+
+def test_interpolation_rejects_non_square_source_and_resizes_equal_area_rectangle() -> None:
+    with pytest.raises(ValueError, match="square patch grid"):
+        modeling_vit._bicubic_interpolate_pos_embed(torch.rand(1, 258, 8), 1, 24, 24)
+    source = torch.arange(1 * 257 * 8, dtype=torch.float32).reshape(1, 257, 8)
+    resized = modeling_vit._bicubic_interpolate_pos_embed(source, 1, 8, 32)
+    assert resized.shape == source.shape
+    assert not torch.equal(resized[:, 1:], source[:, 1:])
+    state_dict = {"visual_encoder.pos_embed": source.clone()}
+    modeling_vit.interpolate_pos_embed(
+        "visual_encoder.pos_embed", 256, 257, state_dict, target_h=8, target_w=32
+    )
+    assert torch.equal(state_dict["visual_encoder.pos_embed"], resized)
+
+
+def test_interpolate_pos_embed_writes_only_on_mismatch() -> None:
+    state_dict = {"visual_encoder.pos_embed": torch.rand(1, 257, 8)}
+    modeling_vit.interpolate_pos_embed(
+        "visual_encoder.pos_embed", 256, 257, state_dict, target_h=16, target_w=16
+    )
+    assert state_dict["visual_encoder.pos_embed"].shape == (1, 257, 8)  # no-op
+
+
+@pytest.mark.skipif(
+    not modeling_vit._HAS_TRANSFORMERS5_CONVERSION_API,
+    reason="Transformers 5 conversion API not available",
+)
+def test_v5_conversion_mapping_registered() -> None:
+    from transformers.conversion_mapping import get_checkpoint_conversion_mapping
+
+    modeling_vit.register_transformers5_pos_embed_conversion()
+    converters = get_checkpoint_conversion_mapping("cosmos-embed1")
+    assert converters and len(converters) == 1
+    assert converters[0].source_patterns == ["visual_encoder.pos_embed"]
+
+
+@pytest.mark.skipif(
+    not modeling_vit._HAS_TRANSFORMERS5_CONVERSION_API,
+    reason="Transformers 5 conversion API not available",
+)
+def test_v5_operation_interpolates_257_to_577() -> None:
+    class _PatchEmbed:
+        patch_shape = (24, 24)
+
+    class _Visual:
+        patch_embed = _PatchEmbed()
+
+    class _Model:
+        visual_encoder = _Visual()
+
+    op = modeling_vit.InterpolatePosEmbed()
+    tensor = torch.arange(1 * 257 * 4, dtype=torch.bfloat16).reshape(1, 257, 4)
+    out = op.convert(
+        {"visual_encoder.pos_embed": [tensor]},
+        ["visual_encoder.pos_embed"],
+        ["visual_encoder.pos_embed"],
+        model=_Model(),
+    )
+    result = out["visual_encoder.pos_embed"][0]
+    assert result.shape == (1, 577, 4)
+    assert result.dtype == torch.bfloat16

modeling_embed1.py

@@ -28,7 +28,18 @@ from .configuration_embed1 import CosmosEmbed1Config
 from .modeling_outputs import TextEmbedderOutput, TextVideoEmbedderOutput, VideoEmbedderOutput
 from .modeling_qformer import BertLMHeadModel, load_qformer
 from .modeling_utils import EncodingFactory, rank0_first
-from .modeling_vit import EvaViTG
+from .modeling_vit import EvaViTG, register_transformers5_pos_embed_conversion
+
+# Register the Transformers>=5 pos_embed interpolation converter. This is a
+# no-op under Transformers<5, where the EvaViTG load-state-dict pre-hook does
+# the interpolation. Transformers 5 bypasses that pre-hook entirely.
+register_transformers5_pos_embed_conversion()
+
+# ImageNet normalization constants, created at import time so they exist on a
+# real device even when the model is later constructed under Transformers 5's
+# "init empty weights" (meta device) context.
+_NORMALIZATION_MEAN = torch.tensor([0.485, 0.456, 0.406]).view(1, 1, 3, 1, 1)
+_NORMALIZATION_STD = torch.tensor([0.229, 0.224, 0.225]).view(1, 1, 3, 1, 1)
 
 
 class CosmosEmbed1(PreTrainedModel):
@@ -51,17 +62,15 @@ class CosmosEmbed1(PreTrainedModel):
         self.transformer_engine = config.transformer_engine
         self.use_fp8 = config.use_fp8
 
-        # visual encoder initialization
-        self.register_buffer(
-            "normalization_mean",
-            torch.tensor([0.485, 0.456, 0.406]).view(1, 1, 3, 1, 1),
-            persistent=False,
-        )
-        self.register_buffer(
-            "normalization_std",
-            torch.tensor([0.229, 0.224, 0.225]).view(1, 1, 3, 1, 1),
-            persistent=False,
-        )
+        # visual encoder normalization constants. These reference module-level
+        # CPU tensors created at import time, not tensors allocated here:
+        # Transformers >= 5 constructs the model under an "init empty weights"
+        # context, so any tensor allocated in __init__ lands on the meta device
+        # with no data. Plain attributes are also not touched by the loader
+        # (unlike non-persistent buffers, which v5 replaces with uninitialized
+        # storage); get_video_embeddings moves them to the input per call.
+        self.normalization_mean = _NORMALIZATION_MEAN
+        self.normalization_std = _NORMALIZATION_STD
         self.visual_encoder = EvaViTG(
             img_size=self.resolution,
             transformer_engine=self.transformer_engine,
@@ -97,6 +106,13 @@ class CosmosEmbed1(PreTrainedModel):
         self.logit_scale = nn.Parameter(torch.tensor(math.log(10.0)))
         self.logit_bias = nn.Parameter(torch.tensor(-10.0))
 
+        # Transformers >= 5 requires models to run post_init explicitly; it
+        # populates tied-weight/parallelism bookkeeping (e.g.
+        # ``all_tied_weights_keys``) that the dynamic loader relies on. It is
+        # also the correct call under Transformers < 5 and replaces the
+        # historical ``init_weights()`` usage in the Q-Former classes.
+        self.post_init()
+
     @property
     def hidden_dim(self) -> int:
         return self.visual_encoder.embed_dim
@@ -131,7 +147,9 @@ class CosmosEmbed1(PreTrainedModel):
         return TextVideoEmbedderOutput(**video_output, **text_output)
 
     def get_video_embeddings(self, videos: torch.Tensor) -> VideoEmbedderOutput:
-        videos = (videos - self.normalization_mean) / self.normalization_std
+        mean = self.normalization_mean.to(device=videos.device, dtype=videos.dtype)
+        std = self.normalization_std.to(device=videos.device, dtype=videos.dtype)
+        videos = (videos - mean) / std
         batch_size, num_frames, _, H, W = videos.shape
         frame_batch = rearrange(videos, "b t c h w -> (b t) c h w")
 

modeling_qformer.py

@@ -44,7 +44,6 @@ from transformers.modeling_utils import (
 )
 from transformers.pytorch_utils import (
     apply_chunking_to_forward,
-    find_pruneable_heads_and_indices,
     prune_linear_layer,
 )
 from transformers.models.bert.configuration_bert import BertConfig
@@ -52,6 +51,24 @@ from transformers.models.bert.configuration_bert import BertConfig
 logger = getLogger(__file__)
 
 
+def find_pruneable_heads_and_indices(heads, n_heads, head_size, already_pruned_heads):
+    """Vendored copy of the Transformers<5 helper removed in Transformers 5.
+
+    ``BertAttention.prune_heads`` calls this on the pruning path only (never
+    during inference), but the module imported it unconditionally, which broke
+    the import under Transformers 5. Defining it locally keeps the remote code
+    importable; the historical implementation is preserved verbatim.
+    """
+    mask = torch.ones(n_heads, head_size)
+    heads = set(heads) - already_pruned_heads
+    for head in heads:
+        head = head - sum(1 if h < head else 0 for h in already_pruned_heads)
+        mask[head] = 0
+    mask = mask.view(-1).contiguous().eq(1)
+    index = torch.arange(len(mask))[mask].long()
+    return heads, index
+
+
 class BertEmbeddings(nn.Module):
     """Construct the embeddings from word and position embeddings."""
 
@@ -615,6 +632,40 @@ class BertPreTrainedModel(PreTrainedModel):
         if isinstance(module, nn.Linear) and module.bias is not None:
             module.bias.data.zero_()
 
+    def invert_attention_mask(self, encoder_attention_mask: Tensor) -> Tensor:
+        """Vendored from Transformers<5 ``ModuleUtilsMixin`` (removed in v5)."""
+        if encoder_attention_mask.dim() == 3:
+            encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
+        if encoder_attention_mask.dim() == 2:
+            encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
+        # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
+        # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
+        # /transformer/transformer_layers.py#L270
+        encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
+        encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * torch.finfo(self.dtype).min
+        return encoder_extended_attention_mask
+
+    def get_head_mask(self, head_mask, num_hidden_layers: int, is_attention_chunked: bool = False):
+        """Vendored from Transformers<5 ``ModuleUtilsMixin`` (removed in v5)."""
+        if head_mask is not None:
+            head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
+            if is_attention_chunked is True:
+                head_mask = head_mask.unsqueeze(-1)
+        else:
+            head_mask = [None] * num_hidden_layers
+        return head_mask
+
+    def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers):
+        """Vendored from Transformers<5 ``ModuleUtilsMixin`` (removed in v5)."""
+        if head_mask.dim() == 1:
+            head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+            head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1)
+        elif head_mask.dim() == 2:
+            head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
+        assert head_mask.dim() == 5, f"head_mask.dim != 5, instead {head_mask.dim()}"
+        head_mask = head_mask.to(dtype=self.dtype)
+        return head_mask
+
 
 class BertModel(BertPreTrainedModel):
     """
@@ -636,7 +687,7 @@ class BertModel(BertPreTrainedModel):
 
         self.pooler = BertPooler(config) if add_pooling_layer else None
 
-        self.init_weights()
+        self.post_init()
 
     def get_input_embeddings(self):
         return self.embeddings.word_embeddings
@@ -886,7 +937,7 @@ class BertLMHeadModel(BertPreTrainedModel, GenerationMixin):
         self.bert = BertModel(config, add_pooling_layer=False)
         self.cls = BertOnlyMLMHead(config)
 
-        self.init_weights()
+        self.post_init()
 
     def get_output_embeddings(self):
         return self.cls.predictions.decoder

modeling_vit.py

@@ -48,6 +48,23 @@ try:
 except ImportError:
     pass
 
+# Transformers >= 5 uses a new dynamic weight loader that no longer calls
+# ``nn.Module.load_state_dict``. The load-state-dict pre-hook registered by
+# ``EvaViTG`` therefore never runs and ``pos_embed`` is not interpolated. When
+# the new conversion API is available we register an equivalent
+# ``WeightConverter`` operation; under Transformers < 5 the imports fail and the
+# pre-hook keeps doing the work.
+try:  # pragma: no cover - branch is selected by the installed version
+    from transformers.conversion_mapping import register_checkpoint_conversion_mapping
+    from transformers.core_model_loading import ConversionOps, WeightConverter
+
+    _HAS_TRANSFORMERS5_CONVERSION_API = True
+except ImportError:
+    register_checkpoint_conversion_mapping = None
+    ConversionOps = None
+    WeightConverter = None
+    _HAS_TRANSFORMERS5_CONVERSION_API = False
+
 
 def drop_path(x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True):
     """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
@@ -565,6 +582,32 @@ class VisionTransformer(nn.Module):
             return len(self.blocks)
 
 
+def _bicubic_interpolate_pos_embed(pos_embed, num_extra_tokens, target_h, target_w):
+    """Bicubically resize position tokens to ``(target_h, target_w)``.
+
+    Extra tokens (e.g. the cls token) are preserved unchanged. Shared by the
+    Transformers<5 pre-hook and the Transformers>=5 conversion operation so
+    both loading paths are numerically identical.
+    """
+    if pos_embed.ndim != 3 or num_extra_tokens < 0 or target_h <= 0 or target_w <= 0:
+        raise ValueError("pos_embed requires (batch, tokens, channels) and a positive target grid")
+    source_tokens = pos_embed.shape[1] - num_extra_tokens
+    orig_size = math.isqrt(max(source_tokens, 0))
+    if source_tokens <= 0 or orig_size * orig_size != source_tokens:
+        raise ValueError("pos_embed source tokens must form a square patch grid")
+    x = pos_embed.float()
+    if orig_size == target_h == target_w:
+        return x
+    extra_tokens = x[:, :num_extra_tokens]
+    pos_tokens = x[:, num_extra_tokens:]
+    pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, x.shape[-1]).permute(0, 3, 1, 2)
+    pos_tokens = torch.nn.functional.interpolate(
+        pos_tokens, size=(target_h, target_w), mode="bicubic", align_corners=False
+    )
+    pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
+    return torch.cat((extra_tokens, pos_tokens), dim=1)
+
+
 def interpolate_pos_embed(
     pos_embed_key: str,
     num_patches: int,
@@ -574,11 +617,10 @@ def interpolate_pos_embed(
     target_w: int = None,
 ) -> None:
     if pos_embed_key in checkpoint_model:
-        pos_embed_checkpoint = checkpoint_model[pos_embed_key].float()
-        embedding_size = pos_embed_checkpoint.shape[-1]
+        pos_embed_checkpoint = checkpoint_model[pos_embed_key]
         num_extra_tokens = patch_embed_shape - num_patches
         # height (== width) for the checkpoint position embedding
-        orig_size = int((pos_embed_checkpoint.shape[-2] - num_extra_tokens) ** 0.5)
+        orig_size = math.isqrt(max(pos_embed_checkpoint.shape[-2] - num_extra_tokens, 0))
 
         # If target dimensions are provided, use them; otherwise assume square
         if target_h is not None and target_w is not None:
@@ -589,18 +631,11 @@ def interpolate_pos_embed(
             new_h, new_w = new_size, new_size
 
         # class_token and dist_token are kept unchanged
-        if orig_size * orig_size != new_h * new_w:
+        if orig_size != new_h or orig_size != new_w:
             logger.info("Positional interpolation from %dx%d to %dx%d" % (orig_size, orig_size, new_h, new_w))
-            extra_tokens = pos_embed_checkpoint[:, :num_extra_tokens]
-            # only the position tokens are interpolated
-            pos_tokens = pos_embed_checkpoint[:, num_extra_tokens:]
-            pos_tokens = pos_tokens.reshape(-1, orig_size, orig_size, embedding_size).permute(0, 3, 1, 2)
-            pos_tokens = torch.nn.functional.interpolate(
-                pos_tokens, size=(new_h, new_w), mode="bicubic", align_corners=False
+            checkpoint_model[pos_embed_key] = _bicubic_interpolate_pos_embed(
+                pos_embed_checkpoint, num_extra_tokens, new_h, new_w
             )
-            pos_tokens = pos_tokens.permute(0, 2, 3, 1).flatten(1, 2)
-            new_pos_embed = torch.cat((extra_tokens, pos_tokens), dim=1)
-            checkpoint_model[pos_embed_key] = new_pos_embed
 
 
 class PositionalEmbeddingHook:
@@ -619,6 +654,80 @@ class PositionalEmbeddingHook:
         )
 
 
+if _HAS_TRANSFORMERS5_CONVERSION_API:
+
+    class _PosEmbedInterpolationNoOp(ConversionOps):
+        """Reverse-op placeholder: bicubic interpolation is not inverted."""
+
+        def convert(self, input_dict, source_patterns=None, target_patterns=None, **kwargs):
+            return input_dict
+
+        @property
+        def reverse_op(self):
+            return self
+
+    class InterpolatePosEmbed(ConversionOps):
+        """Bicubic ``pos_embed`` interpolation for the Transformers>=5 loader.
+
+        Transformers 5 assigns weights with ``setattr`` and never calls
+        ``nn.Module.load_state_dict``, so ``PositionalEmbeddingHook`` is dead
+        code there. The loader passes the instantiated model to ``convert``,
+        from which the target patch grid is read. The math is shared with the
+        pre-hook path via :func:`_bicubic_interpolate_pos_embed`.
+        """
+
+        def __init__(self, num_extra_tokens: int = 1):
+            self.num_extra_tokens = num_extra_tokens
+
+        @torch.no_grad()
+        def convert(self, input_dict, source_patterns, target_patterns, **kwargs):
+            tensor = next(iter(input_dict.values()))[0]
+            patch_h, patch_w = _model_patch_shape(kwargs.get("model"))
+            interpolated = _bicubic_interpolate_pos_embed(
+                tensor, self.num_extra_tokens, int(patch_h), int(patch_w)
+            )
+            return {target_patterns[0]: [interpolated.to(tensor.dtype)]}
+
+        @property
+        def reverse_op(self):
+            return _PosEmbedInterpolationNoOp()
+
+    def _model_patch_shape(model):
+        visual = getattr(model, "visual_encoder", None)
+        patch_embed = getattr(visual, "patch_embed", None)
+        shape = getattr(patch_embed, "patch_shape", None)
+        if shape is None:
+            raise RuntimeError(
+                "InterpolatePosEmbed: model has no visual_encoder.patch_embed.patch_shape"
+            )
+        return shape[0], shape[1]
+
+    def register_transformers5_pos_embed_conversion() -> None:
+        """Register the pos_embed converter for ``model_type='cosmos-embed1'``.
+
+        ``get_model_conversion_mapping`` skips custom-code submodules unless
+        their class name or model_type is user-registered, so this call is what
+        makes the converter apply to the remote model. Safe to call repeatedly.
+        """
+        register_checkpoint_conversion_mapping(
+            "cosmos-embed1",
+            [
+                WeightConverter(
+                    source_patterns=["visual_encoder.pos_embed"],
+                    target_patterns=["visual_encoder.pos_embed"],
+                    operations=[InterpolatePosEmbed()],
+                )
+            ],
+            overwrite=True,
+        )
+
+else:
+
+    def register_transformers5_pos_embed_conversion() -> None:
+        """No-op under Transformers<5; the load-state-dict pre-hook handles it."""
+        return None
+
+
 class EvaViTG(VisionTransformer):
     def __init__(
         self,