| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from transformers.configuration_utils import PreTrainedConfig |
| 5 | from transformers.models.siglip2.configuration_siglip2 import Siglip2VisionConfig |
| 6 | |
| 7 | from .nemotron_h import NemotronHConfig |
| 8 | |
| 9 | |
| 10 | def _normalize_mrope_parameters(config: NemotronHConfig) -> None: |
| 11 | """Normalize the checkpoint's mRoPE configuration for vLLM.""" |
| 12 | rope_parameters = dict(getattr(config, "rope_parameters", None) or {}) |
| 13 | mrope_section = rope_parameters.get("mrope_section") |
| 14 | if mrope_section is None: |
| 15 | return |
| 16 | |
| 17 | rope_parameters["mrope_section"] = list(mrope_section) |
| 18 | rope_parameters.setdefault("mrope_interleaved", True) |
| 19 | config.rope_parameters = rope_parameters |
| 20 | |
| 21 | |
| 22 | class Cosmos3EdgeTextConfig(NemotronHConfig): |
| 23 | """Dense Cosmos3 Edge text config backed by Nemotron-H layers.""" |
| 24 | |
| 25 | model_type = "cosmos3_edge_text" |
| 26 | has_no_defaults_at_init = True |
| 27 | ignore_keys_at_rope_validation = {"mrope_section", "mrope_interleaved"} |
| 28 | |
| 29 | def __init__( |
| 30 | self, |
| 31 | num_hidden_layers: int = 28, |
| 32 | hidden_act: str = "relu2", |
| 33 | rms_norm_eps: float = 1e-5, |
| 34 | **kwargs, |
| 35 | ) -> None: |
| 36 | super().__init__( |
| 37 | num_hidden_layers=2 * num_hidden_layers, |
| 38 | hybrid_override_pattern="*-" * num_hidden_layers, |
| 39 | mlp_hidden_act=hidden_act, |
| 40 | layer_norm_epsilon=rms_norm_eps, |
| 41 | **kwargs, |
| 42 | ) |
| 43 | |
| 44 | def to_dict(self) -> dict: |
| 45 | config_dict = super().to_dict() |
| 46 | config_dict["num_hidden_layers"] = self.num_hidden_layers // 2 |
| 47 | config_dict["hidden_act"] = config_dict.pop("mlp_hidden_act") |
| 48 | config_dict["rms_norm_eps"] = config_dict.pop("layer_norm_epsilon") |
| 49 | config_dict.pop("hybrid_override_pattern", None) |
| 50 | return config_dict |
| 51 | |
| 52 | |
| 53 | class Cosmos3EdgeProjectorConfig(PreTrainedConfig): |
| 54 | model_type = "cosmos3_edge_projector" |
| 55 | |
| 56 | def __init__( |
| 57 | self, |
| 58 | input_hidden_size: int = 1024, |
| 59 | use_postshuffle_norm: bool = False, |
| 60 | spatial_merge_size: int = 2, |
| 61 | merger_intermediate_size: int = 11520, |
| 62 | out_hidden_size: int = 1024, |
| 63 | **kwargs, |
| 64 | ) -> None: |
| 65 | super().__init__(**kwargs) |
| 66 | self.input_hidden_size = input_hidden_size |
| 67 | self.use_postshuffle_norm = use_postshuffle_norm |
| 68 | self.spatial_merge_size = spatial_merge_size |
| 69 | self.merger_intermediate_size = merger_intermediate_size |
| 70 | self.out_hidden_size = out_hidden_size |
| 71 | |
| 72 | |
| 73 | class Cosmos3EdgeVisionConfig(Siglip2VisionConfig): |
| 74 | model_type = "cosmos3_edge_vision" |
| 75 | |
| 76 | |
| 77 | class Cosmos3EdgeConfig(PreTrainedConfig): |
| 78 | model_type = "cosmos3_edge" |
| 79 | keys_to_ignore_at_inference = ["past_key_values"] |
| 80 | has_no_defaults_at_init = True |
| 81 | sub_configs = { |
| 82 | "vision_config": Cosmos3EdgeVisionConfig, |
| 83 | "projector_config": Cosmos3EdgeProjectorConfig, |
| 84 | "text_config": Cosmos3EdgeTextConfig, |
| 85 | } |
| 86 | |
| 87 | def __init__( |
| 88 | self, |
| 89 | text_config: Cosmos3EdgeTextConfig | dict | None = None, |
| 90 | vision_config: Cosmos3EdgeVisionConfig | dict | None = None, |
| 91 | projector_config: Cosmos3EdgeProjectorConfig | dict | None = None, |
| 92 | image_token_id: int = 19, |
| 93 | video_token_id: int = 18, |
| 94 | vision_start_token_id: int = 20, |
| 95 | vision_end_token_id: int = 21, |
| 96 | tie_word_embeddings: bool = False, |
| 97 | **kwargs, |
| 98 | ) -> None: |
| 99 | if text_config is None: |
| 100 | text_config = {} |
| 101 | if vision_config is None: |
| 102 | vision_config = {} |
| 103 | if projector_config is None: |
| 104 | projector_config = {} |
| 105 | |
| 106 | self.text_config = ( |
| 107 | Cosmos3EdgeTextConfig(**text_config) |
| 108 | if isinstance(text_config, dict) |
| 109 | else text_config |
| 110 | ) |
| 111 | _normalize_mrope_parameters(self.text_config) |
| 112 | |
| 113 | self.vision_config = ( |
| 114 | Cosmos3EdgeVisionConfig(**vision_config) |
| 115 | if isinstance(vision_config, dict) |
| 116 | else vision_config |
| 117 | ) |
| 118 | self.projector_config = ( |
| 119 | Cosmos3EdgeProjectorConfig(**projector_config) |
| 120 | if isinstance(projector_config, dict) |
| 121 | else projector_config |
| 122 | ) |
| 123 | |
| 124 | # Qwen3-VL processing and M-RoPE read these attributes from the |
| 125 | # vision config, while this checkpoint stores them in projector config. |
| 126 | self.vision_config.spatial_merge_size = self.projector_config.spatial_merge_size |
| 127 | self.vision_config.temporal_patch_size = 1 |
| 128 | self.vision_config.out_hidden_size = self.projector_config.out_hidden_size |
| 129 | |
| 130 | self.image_token_id = image_token_id |
| 131 | self.video_token_id = video_token_id |
| 132 | self.vision_start_token_id = vision_start_token_id |
| 133 | self.vision_end_token_id = vision_end_token_id |
| 134 | |
| 135 | super().__init__( |
| 136 | tie_word_embeddings=tie_word_embeddings, |
| 137 | **kwargs, |
| 138 | ) |
| 139 | |
| 140 | |
| 141 | __all__ = [ |
| 142 | "Cosmos3EdgeConfig", |
| 143 | "Cosmos3EdgeProjectorConfig", |
| 144 | "Cosmos3EdgeTextConfig", |
| 145 | "Cosmos3EdgeVisionConfig", |
| 146 | ] |