Add NVFP4 weights, config, tokenizer, and trust_remote_code modeling files

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+tokenizer.json filter=lfs diff=lfs merge=lfs -text

audio_model.py

@@ -0,0 +1,174 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+"""Sound/Audio model components for multimodal integration.
+
+This module provides the SoundEncoder (wrapping Parakeet from HuggingFace transformers)
+and SoundProjection (MLP to project audio embeddings to LLM hidden size).
+
+The Parakeet model in HuggingFace transformers is documented at:
+https://huggingface.co/docs/transformers/en/model_doc/parakeet
+"""
+
+from typing import Optional
+
+import torch
+import torch.nn as nn
+
+from transformers import ParakeetEncoder, ParakeetEncoderConfig
+from transformers.utils import logging
+
+logger = logging.get_logger(__name__)
+
+
+class SquaredReLU(nn.Module):
+    """Squared ReLU activation function."""
+    def forward(self, x):
+        return torch.pow(torch.nn.functional.relu(x), 2)
+
+
+class RMSNorm(nn.Module):
+    def __init__(self, hidden_size, eps=1e-5):
+        super().__init__()
+        self.weight = nn.Parameter(torch.ones(hidden_size))
+        self.eps = eps
+
+    def forward(self, hidden_states):
+        input_dtype = hidden_states.dtype
+        hidden_states = hidden_states.to(torch.float32)
+        variance = hidden_states.pow(2).mean(-1, keepdim=True)
+        hidden_states = hidden_states * torch.rsqrt(variance + self.eps)
+        return (self.weight.to(torch.float32) * hidden_states).to(input_dtype)
+
+
+class SoundProjection(nn.Module):
+    """MLP projection from sound encoder hidden size to LLM hidden size.
+    
+    Architecture: RMSNorm -> linear1 -> SquaredReLU -> linear2
+    
+    This matches the Megatron checkpoint conversion structure:
+    - sound_projection.norm.weight
+    - sound_projection.linear1.weight
+    - sound_projection.linear2.weight
+    - sound_projection.linear1.bias (optional)
+    - sound_projection.linear2.bias (optional)
+    """
+    
+    def __init__(
+        self,
+        sound_hidden_size: int,
+        projection_hidden_size: int,
+        llm_hidden_size: int,
+        bias: bool = True,
+        eps: float = 1e-5,
+    ):
+        super().__init__()
+        self.norm = RMSNorm(sound_hidden_size, eps=eps)
+        self.linear1 = nn.Linear(sound_hidden_size, projection_hidden_size, bias=bias)
+        self.activation = SquaredReLU()
+        self.linear2 = nn.Linear(projection_hidden_size, llm_hidden_size, bias=bias)
+        
+    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
+        """Project sound embeddings to LLM embedding space.
+        
+        Args:
+            hidden_states: Sound encoder output [batch, seq_len, sound_hidden_size]
+        
+        Returns:
+            Projected embeddings [batch, seq_len, llm_hidden_size]
+        """
+        hidden_states = self.norm(hidden_states)
+        hidden_states = self.linear1(hidden_states)
+        hidden_states = self.activation(hidden_states)
+        hidden_states = self.linear2(hidden_states)
+        return hidden_states
+
+
+class SoundEncoder(nn.Module):
+    """Wrapper around the Parakeet encoder from HuggingFace transformers.
+    
+    The Parakeet model is an ASR model with a Fast Conformer encoder.
+    We use only the encoder portion to extract audio embeddings.
+    
+    Checkpoint structure:
+    - sound_encoder.encoder.feature_extractor.* -> Feature extraction (mel spectrogram)
+    - sound_encoder.encoder.pre_encode.* -> Pre-encoding convolutions
+    - sound_encoder.encoder.layers.* -> Conformer layers
+    
+    Reference: https://huggingface.co/docs/transformers/en/model_doc/parakeet
+    """
+    
+    def __init__(self, config=None):
+        super().__init__()
+        
+        if config is not None:
+            # Build from config - handle both dict and config object
+            if hasattr(config, '__dict__'):
+                # It's a config object, extract relevant params for ParakeetConfig
+                config_dict = {
+                    'attention_bias': getattr(config, 'attention_bias', False),
+                    'hidden_size': getattr(config, 'hidden_size', 1024),
+                    'num_attention_heads': getattr(config, 'num_attention_heads', 8),
+                    'num_hidden_layers': getattr(config, 'num_hidden_layers', 24),
+                    'intermediate_size': getattr(config, 'intermediate_size', 4096),
+                    'conv_kernel_size': getattr(config, 'conv_kernel_size', 31),
+                    'convolution_bias': getattr(config, 'convolution_bias', False),
+                    'feat_in': getattr(config, 'feat_in', 80),
+                    'subsampling_factor': getattr(config, 'subsampling_factor', 8),
+                    'subsampling_conv_channels': getattr(config, 'subsampling_conv_channels', 256),
+                    'subsampling_conv_kernel_size': getattr(config, 'subsampling_conv_kernel_size', 3),
+                    'subsampling_conv_stride': getattr(config, 'subsampling_conv_stride', 2),
+                    'num_mel_bins': getattr(config, 'num_mel_bins', 128),
+                    'scale_input': getattr(config, 'scale_input', False),
+                }
+            elif isinstance(config, dict):
+                config_dict = config
+            else:
+                config_dict = {}
+            
+            # Create ParakeetConfig with the extracted parameters
+            parakeet_config = ParakeetEncoderConfig(**config_dict)
+            self.config = parakeet_config
+            self.encoder = ParakeetEncoder(parakeet_config)
+        else:
+            raise ValueError(
+                "config must be provided, "
+                "and ParakeetEncoder must be available in transformers."
+            )
+    
+    def forward(
+        self,
+        input_features: torch.Tensor,
+        attention_mask: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Encode audio features.
+        
+        Args:
+            input_features: Mel spectrogram features [batch, seq_len, feature_dim]
+            attention_mask: Optional attention mask [batch, seq_len]
+            
+        Returns:
+            Audio embeddings [batch, encoded_seq_len, hidden_size]
+        """
+        outputs = self.encoder(
+            input_features=input_features,
+            attention_mask=attention_mask,
+        )
+        # Return the last hidden state
+        return outputs.last_hidden_state
+    
+    @property
+    def hidden_size(self) -> int:
+        """Return the hidden size of the encoder."""
+        return self.config.hidden_size

chat_template.jinja

@@ -0,0 +1,273 @@
+{% macro render_extra_keys(json_dict, handled_keys) %}
+    {%- if json_dict is mapping %}
+        {%- for json_key in json_dict if json_key not in handled_keys %}
+            {%- if json_dict[json_key] is mapping or (json_dict[json_key] is sequence and json_dict[json_key] is not string) %}
+                {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | tojson | safe) ~ '</' ~ json_key ~ '>' }}
+            {%- else %}
+                {{- '\n<' ~ json_key ~ '>' ~ (json_dict[json_key] | string) ~ '</' ~ json_key ~ '>' }}
+            {%- endif %}
+        {%- endfor %}
+    {%- endif %}
+{%- endmacro -%}
+{%- set enable_thinking = enable_thinking if enable_thinking is defined else True %}
+{%- set reasoning_budget = reasoning_budget if reasoning_budget is defined else None %}
+{%- set truncate_history_thinking = truncate_history_thinking if truncate_history_thinking is defined else True %}
+
+{#- Scan messages for VLM thinking toggles to override enable_thinking -#}
+{%- set toggle = namespace(enable=enable_thinking) %}
+{%- for m in messages %}
+    {%- if m['role'] == 'user' or m['role'] == 'system' -%}
+        {%- if m['content'] is string -%}
+            {%- set c = m['content'] %}
+            {%- if '/think' in c.replace('</think>', '') -%}
+                {%- set toggle.enable = true -%}
+            {%- elif '/no_think' in c -%}
+                {%- set toggle.enable = false -%}
+            {%- endif -%}
+        {%- else -%}
+            {%- for part in m['content'] -%}
+                {%- if part['type'] == 'text' -%}
+                    {%- set c = part['text'] %}
+                    {%- if '/think' in c.replace('</think>', '') -%}
+                        {%- set toggle.enable = true -%}
+                    {%- elif '/no_think' in c -%}
+                        {%- set toggle.enable = false -%}
+                    {%- endif -%}
+                {%- endif -%}
+            {%- endfor -%}
+        {%- endif -%}
+    {%- endif -%}
+{%- endfor -%}
+{#- Prepare message iteration similar to LM template -#}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- set loop_messages = messages %}
+{%- for m in loop_messages %}
+    {%- if m["role"] == "user" %}
+        {%- set ns.last_user_idx = loop.index0 %}
+    {%- endif %}
+{%- endfor -%}
+
+{%- if messages[0]["role"] == "system" %}
+    {%- set system_message = messages[0]["content"] %}
+    {%- set loop_messages = messages[1:] %}
+{%- else %}
+    {%- set system_message = "" %}
+    {%- set loop_messages = messages %}
+{%- endif %}
+{%- if not tools is defined %}
+    {%- set tools = [] %}
+{%- endif %}
+{#- Recompute last_user_idx relative to loop_messages after handling system -#}
+{%- set ns = namespace(last_user_idx = -1) %}
+{%- for m in loop_messages %}
+    {%- if m["role"] == "user" %}
+        {%- set ns.last_user_idx = loop.index0 %}
+    {%- endif %}
+{%- endfor -%}
+{#- System preamble with LM formatting, sanitize thinking toggles -#}
+{%- if system_message is defined %}
+    {%- set sys_content = system_message | string %}
+    {%- set sys_content = sys_content.replace('</think>', '<_end_think>').replace('/think', '').replace('/no_think', '').replace('<_end_think>', '</think>') %}
+    {{- "<|im_start|>system\n" + sys_content }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- "<|im_start|>system\n" }}
+    {%- endif %}
+{%- endif %}
+{%- if tools is iterable and tools | length > 0 %}
+    {%- if system_message is defined and system_message | length > 0 %}
+        {{- "\n\n" }}
+    {%- endif %}
+    {{- "# Tools\n\nYou have access to the following functions:\n\n" }}
+    {{- "<tools>" }}
+    {%- for tool in tools %}
+        {%- if tool.function is defined %}
+            {%- set tool = tool.function %}
+        {%- endif %}
+        {{- "\n<function>\n<name>" ~ tool.name ~ "</name>" }}
+        {%- if tool.description is defined %}
+            {{- '\n<description>' ~ (tool.description | trim) ~ '</description>' }}
+        {%- endif %}
+        {{- '\n<parameters>' }}
+        {%- if tool.parameters is defined and tool.parameters is mapping and tool.parameters.properties is defined and tool.parameters.properties is mapping %}
+            {%- for param_name, param_fields in tool.parameters.properties|items %}
+                {{- '\n<parameter>' }}
+                {{- '\n<name>' ~ param_name ~ '</name>' }}
+                {%- if param_fields.type is defined %}
+                    {{- '\n<type>' ~ (param_fields.type | string) ~ '</type>' }}
+                {%- endif %}
+                {%- if param_fields.description is defined %}
+                    {{- '\n<description>' ~ (param_fields.description | trim) ~ '</description>' }}
+                {%- endif %}
+                {%- if param_fields.enum is defined %}
+                    {{- '\n<enum>' ~ (param_fields.enum | tojson | safe) ~ '</enum>' }}
+                {%- endif %}
+                {%- set handled_keys = ['name', 'type', 'description', 'enum'] %}
+                {{- render_extra_keys(param_fields, handled_keys) }}
+                {{- '\n</parameter>' }}
+            {%- endfor %}
+        {%- endif %}
+        {%- set handled_keys = ['type', 'properties', 'required'] %}
+        {{- render_extra_keys(tool.parameters, handled_keys) }}
+        {%- if tool.parameters is defined and tool.parameters.required is defined %}
+            {{- '\n<required>' ~ (tool.parameters.required | tojson | safe) ~ '</required>' }}
+        {%- endif %}
+        {{- '\n</parameters>' }}
+        {%- set handled_keys = ['type', 'name', 'description', 'parameters'] %}
+        {{- render_extra_keys(tool, handled_keys) }}
+        {{- '\n</function>' }}
+    {%- endfor %}
+    {{- "\n</tools>" }}
+
+    {{- '\n\nIf you choose to call a function ONLY reply in the following format with NO suffix:\n\n<tool_call>\n<function=example_function_name>\n<parameter=example_parameter_1>\nvalue_1\n</parameter>\n<parameter=example_parameter_2>\nThis is the value for the second parameter\nthat can span\nmultiple lines\n</parameter>\n</function>\n</tool_call>\n\n<IMPORTANT>\nReminder:\n- Function calls MUST follow the specified format: an inner <function=...></function> block must be nested within <tool_call></tool_call> XML tags\n- Required parameters MUST be specified\n- You may provide optional reasoning for your function call in natural language BEFORE the function call, but NOT after\n- If there is no function call available, answer the question like normal with your current knowledge and do not tell the user about function calls\n</IMPORTANT>' }}
+{%- endif -%}
+{%- if system_message is defined %}
+    {{- '<|im_end|>\n' }}
+{%- else %}
+    {%- if tools is iterable and tools | length > 0 %}
+        {{- '<|im_end|>\n' }}
+    {%- endif %}
+{%- endif -%}
+{#- Iterate conversation -#}
+{%- for message in loop_messages %}
+    {%- if message.role == "assistant" %}
+        {#- Use LM assistant handling -#}
+        {%- if message.reasoning_content is defined and message.reasoning_content is string and message.reasoning_content | trim | length > 0 %}
+            {%- set content = "<think>\n" ~ message.reasoning_content ~ "\n</think>\n" ~ (message.content | default('', true)) %}
+        {%- else %}
+            {%- set content = message.content | default('', true) %}
+            {%- if content is string -%}
+                {%- if '<think>' not in content and '</think>' not in content -%}
+                    {%- set content = "<think></think>" ~ content -%}
+                {%- endif -%}
+            {%- else -%}
+                {%- set content = content -%}
+            {%- endif -%}
+        {%- endif %}
+        {%- if message.tool_calls is defined and message.tool_calls is iterable and message.tool_calls | length > 0 %}
+            {{- '<|im_start|>assistant\n' }}
+                {%- set include_content = not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {%- if content is string and content | trim | length > 0 %}
+                    {%- if include_content %}
+                        {{- (content | trim) ~ '\n' -}}
+                    {%- else %}
+                        {%- set c = (content | string) %}
+                        {%- if '</think>' in c %}
+                            {%- set c = c.split('</think>')[-1] %}
+                        {%- elif '<think>' in c %}
+                            {%- set c = c.split('<think>')[0] %}
+                        {%- endif %}
+                        {%- set c = "<think></think>" ~ c | trim %}
+                        {%- if c | length > 0 %}
+                            {{- c ~ '\n' -}}
+                        {%- endif %}
+                    {%- endif %}
+                {%- else %}
+                    {{- "<think></think>" -}}
+                {%- endif %}
+                {%- for tool_call in message.tool_calls %}
+                    {%- if tool_call.function is defined %}
+                        {%- set tool_call = tool_call.function %}
+                    {%- endif %}
+                    {{- '<tool_call>\n<function=' ~ tool_call.name ~ '>\n' -}}
+                        {%- if tool_call.arguments is defined %}
+                            {%- for args_name, args_value in tool_call.arguments|items %}
+                                {{- '<parameter=' ~ args_name ~ '>\n' -}}
+                                    {%- set args_value = args_value | tojson | safe if args_value is mapping or (args_value is sequence and args_value is not string) else args_value | string %}
+                                {{- args_value ~ '\n</parameter>\n' -}}
+                            {%- endfor %}
+                        {%- endif %}
+                    {{- '</function>\n</tool_call>\n' -}}
+                {%- endfor %}
+                {{- '<|im_end|>\n' }}
+        {%- else %}
+            {%- if not (truncate_history_thinking and loop.index0 < ns.last_user_idx) %}
+                {{- '<|im_start|>assistant\n' ~ (content | default('', true) | string | trim) ~ '<|im_end|>\n' }}
+            {%- else %}
+                {%- set c = (content | default('', true) | string) %}
+                {%- if '<think>' in c and '</think>' in c %}
+                    {%- set c = "<think></think>" ~ c.split('</think>')[-1] %}
+                {%- endif %}
+                {%- set c = c | trim %}
+                {%- if c | length > 0 %}
+                    {{- '<|im_start|>assistant\n' ~ c ~ '<|im_end|>\n' }}
+                {%- else %}
+                    {{- '<|im_start|>assistant\n<|im_end|>\n' }}
+                {%- endif %}
+            {%- endif %}
+        {%- endif %}
+    {%- elif message.role == "user" or message.role == "system" %}
+        {{- '<|im_start|>' + message.role + '\n' }}
+        {#- Build VLM multimodal content when content is a sequence -#}
+        {%- if message.content is string -%}
+            {%- set content = (message.content | string) %}
+        {%- else -%}
+            {%- set text_ns = namespace(val='') -%}
+            {%- set mm_content = '' -%}
+            {%- set counters = namespace(images=0, videos=0, audios=0) -%}
+            {%- for part in message['content'] -%}
+                {%- if part['type'] == 'image' or part['type'] == 'image_url' -%}
+                    {%- set counters.images = counters.images + 1 -%}
+                {%- elif part['type'] == 'video' or part['type'] == 'video_url' -%}
+                    {%- set counters.videos = counters.videos + 1 -%}
+                {%- elif part['type'] == 'audio' or part['type'] == 'audio_url' -%}
+                    {%- set counters.audios = counters.audios + 1 -%}
+                {%- elif part['type'] == 'text' -%}
+                    {%- set text_ns.val = text_ns.val + part['text'] -%}
+                {%- endif -%}
+            {%- endfor -%}
+            {%- if '<image>' in text_ns.val -%}
+                {%- set counters.images = 0 -%}
+            {%- endif -%}
+            {%- if '<video>' in text_ns.val -%}
+                {%- set counters.videos = 0 -%}
+            {%- endif -%}
+            {%- if '<so_embedding>' in text_ns.val -%}
+                {%- set counters.audios = 0 -%}
+            {%- endif -%}
+            {%- if counters.images > 1 -%}
+                {%- set image_tags = namespace(tags=[]) -%}
+                {%- for i in range(counters.images) -%}
+                    {%- set image_tags.tags = image_tags.tags + ['<image ' + (i + 1)|string + '><image>'] -%}
+                {%- endfor -%}
+                {%- set mm_content = ' '.join(image_tags.tags) + '\n' -%}
+            {%- elif counters.images == 1 -%}
+                {%- set mm_content = '<image>\n' -%}
+            {%- endif -%}
+            {%- set mm_content = mm_content + '<video>\n' * counters.videos -%}
+            {%- set mm_content = mm_content + '<so_embedding>\n' * counters.audios -%}
+            {%- set content = mm_content + text_ns.val.lstrip('\n') -%}
+        {%- endif -%}
+        {#- Sanitize thinking toggle directives from user/system content -#}
+        {%- set content = content.replace('</think>', '<_end_think>').replace('/think', '').replace('/no_think', '').replace('<_end_think>', '</think>') -%}
+        {%- set content = content | trim -%}
+        {%- if message.role == "user" and loop.index0 == ns.last_user_idx and reasoning_budget is not none -%}
+            {{- content + '\n\n{thinking token budget: ' + (reasoning_budget | string) + '}' -}}
+        {%- else -%}
+            {{- content -}}
+        {%- endif -%}
+        {{- '<|im_end|>\n' }}
+    {%- elif message.role == "tool" %}
+        {%- if loop.previtem and loop.previtem.role != "tool" %}
+            {{- '<|im_start|>user\n' }}
+        {%- endif %}
+        {{- '<tool_response>\n' }}
+        {{- message.content }}
+        {{- '\n</tool_response>\n' }}
+        {%- if not loop.last and loop.nextitem.role != "tool" %}
+            {{- '<|im_end|>\n' }}
+        {%- elif loop.last %}
+            {{- '<|im_end|>\n' }}
+        {%- endif %}
+    {%- else %}
+        {{- '<|im_start|>' + message.role + '\n' + message.content + '<|im_end|>\n' }}
+    {%- endif %}
+{%- endfor -%}
+{#- Generation prompt using computed thinking toggle -#}
+{%- if add_generation_prompt %}
+    {%- if toggle.enable %}
+        {{- '<|im_start|>assistant\n<think>\n' }}
+    {%- else %}
+        {{- '<|im_start|>assistant\n<think></think>' }}
+    {%- endif %}
+{%- endif %}

configuration.py

@@ -0,0 +1,114 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+from .configuration_nemotron_h import NemotronHConfig
+from .configuration_radio import RADIOConfig
+
+logger = logging.get_logger(__name__)
+
+
+class SoundConfig(PretrainedConfig):
+    """Configuration for the sound/audio model (Parakeet encoder + projection)."""
+    model_type = "parakeet"
+    
+    def __init__(
+        self,
+        # Parakeet encoder config
+        hidden_size: int = 1024,
+        num_attention_heads: int = 8,
+        num_hidden_layers: int = 24,
+        intermediate_size: int = 4096,
+        conv_kernel_size: int = 31,
+        feat_in: int = 80,  # Mel features
+        subsampling_factor: int = 8,
+        # Projection config  
+        projection_hidden_size: int = 20480,
+        projection_bias: bool = True,
+        # Audio processing
+        sampling_rate: int = 16000,
+        **kwargs,
+    ):
+        super().__init__(**kwargs)
+        self.hidden_size = hidden_size
+        self.num_attention_heads = num_attention_heads
+        self.num_hidden_layers = num_hidden_layers
+        self.intermediate_size = intermediate_size
+        self.conv_kernel_size = conv_kernel_size
+        self.feat_in = feat_in
+        self.subsampling_factor = subsampling_factor
+        self.projection_hidden_size = projection_hidden_size
+        self.projection_bias = projection_bias
+        self.sampling_rate = sampling_rate
+
+
+class NemotronH_Nano_Omni_Reasoning_V3_Config(PretrainedConfig):
+    model_type = 'NemotronH_Nano_Omni_Reasoning_V3'
+    is_composition = True
+
+    def __init__(
+        self,
+        vision_config=None,
+        llm_config=None,
+        sound_config=None,
+        force_image_size=None,
+        downsample_ratio=0.5,
+        template=None,
+        ps_version='v1',
+        image_tag_type="internvl",
+        projector_hidden_size=4096,
+        vit_hidden_size=1280,
+        attn_implementation="flash_attention_2",
+        video_pruning_rate: float = 0.0,
+        # Sound/audio settings
+        sound_context_token_id: int = None,
+        sound_context_token: str = "<audio>",
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+
+        if vision_config is not None:
+            self.vision_config = RADIOConfig(**vision_config)
+        else:
+            self.vision_config = RADIOConfig()
+
+        # Handle both cases: when loading from JSON (llm_config is dict) and when called internally by transformers (llm_config is None)
+        if llm_config is not None:
+            self.llm_config = NemotronHConfig(**llm_config)
+        else:
+            self.llm_config = NemotronHConfig()
+
+        # Sound/audio model configuration
+        if sound_config is not None:
+            self.sound_config = SoundConfig(**sound_config)
+        else:
+            self.sound_config = None  # Sound model is optional
+
+        # Assign configuration values
+        self.force_image_size = force_image_size
+        self.downsample_ratio = downsample_ratio
+        self.template = template  # TODO move out of here and into the tokenizer
+        self.ps_version = ps_version  # Pixel shuffle version
+        self.image_tag_type = image_tag_type # TODO: into the tokenizer too?
+        self.projector_hidden_size = projector_hidden_size
+        self.vit_hidden_size = vit_hidden_size
+        self.video_pruning_rate = video_pruning_rate
+        
+        # Sound/audio token settings
+        self.sound_context_token_id = sound_context_token_id
+        self.sound_context_token = sound_context_token
+
+        self._attn_implementation = attn_implementation
+        self.vision_config.use_flash_attn = self._attn_implementation is not None and "flash_attention" in self._attn_implementation
+        self.llm_config._attn_implementation = self._attn_implementation

configuration_nemotron_h.py

@@ -0,0 +1,262 @@
+# coding=utf-8
+# Copyright 2024 AI21 Labs Ltd. and the HuggingFace Inc. team. All rights reserved.
+# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+"""NemotronH model configuration"""
+
+import re
+
+from transformers.configuration_utils import PretrainedConfig
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class NemotronHConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a [`NemotronHModel`]. It is used to instantiate a
+    NemotronH model according to the specified arguments, defining the model architecture. Instantiating a configuration
+    with the defaults will yield a similar configuration to that of the NemotronH-v0.1 model.
+
+    [todo](todo)
+
+    Configuration objects inherit from [`PretrainedConfig`] and can be used to control the model outputs. Read the
+    documentation from [`PretrainedConfig`] for more information.
+
+
+    Args:
+        vocab_size (`int`, *optional*, defaults to 131072):
+            Vocabulary size of the NemotronH model. Defines the number of different tokens that can be represented by the
+            `inputs_ids` passed when calling [`NemotronHModel`]
+        tie_word_embeddings (`bool`, *optional*, defaults to `False`):
+            Whether the model's input and output word embeddings should be tied. Note that this is only relevant if the
+            model has a output word embedding layer.
+        hidden_size (`int`, *optional*, defaults to 4096):
+            Dimension of the hidden representations.
+        intermediate_size (`int`, *optional*, defaults to 21504):
+            Dimension of the MLP representations.
+        num_hidden_layers (`int`, *optional*, defaults to 52):
+            Number of hidden layers in the Transformer encoder.
+        hybrid_override_pattern (`str`, *optional*, defaults to `"M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-"`):
+            The pattern of the hybrid model. The pattern is a string of characters where each character represents M: Mamba2, *: Attention, -: MLP
+        num_attention_heads (`int`, *optional*, defaults to 32):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        head_dim (`int`, *optional*, defaults to 128):
+            Dimension of each attention head.
+        num_key_value_heads (`int`, *optional*, defaults to 8):
+            This is the number of key_value heads that should be used to implement Grouped Query Attention. If
+            `num_key_value_heads=num_attention_heads`, the model will use Multi Head Attention (MHA), if
+            `num_key_value_heads=1` the model will use Multi Query Attention (MQA) otherwise GQA is used.
+        mlp_hidden_act (`str`, *optional*, defaults to "relu2"):
+            The non-linear activation function in the MLP layers.
+        attention_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use bias in attention layers.
+        mlp_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use bias in MLP layers.
+        use_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use bias in the model.
+        initializer_range (`float`, *optional*, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        layer_norm_epsilon (`float`, *optional*, defaults to 1e-5):
+            The epsilon used by the layer normalization layers.
+        residual_in_fp32 (`bool`, *optional*, defaults to `False`):
+            Whether or not residuals should be in `float32`. If set to `False` residuals will keep the same `dtype` as the rest of the model.
+        use_cache (`bool`, *optional*, defaults to `True`):
+            Whether or not the model should return the last key/values attentions (not used by all models). Only
+            relevant if `config.is_decoder=True`.
+        num_logits_to_keep (`int` or `None`, *optional*, defaults to 1):
+            Number of prompt logits to calculate during generation. If `None`, all logits will be calculated. If an
+            integer value, only last `num_logits_to_keep` logits will be calculated.
+        pad_token_id (`int`, *optional*, defaults to 0):
+            The id of the padding token.
+        bos_token_id (`int`, *optional*, defaults to 1):
+            The id of the "beginning-of-sequence" token.
+        eos_token_id (`int`, *optional*, defaults to 2):
+            The id of the "end-of-sequence" token.
+        sliding_window (`int`, *optional*, defaults to None):
+            Sliding window attention window size.
+        max_position_embeddings (`int`, *optional*, defaults to 4096):
+            The maximum sequence length that this model might ever be used with.
+        attention_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the attention probabilities.
+        hidden_dropout (`float`, *optional*, defaults to 0.0):
+            The dropout ratio for the hidden states.
+        use_mamba_kernels (`bool`, *optional*, defaults to `True`):
+            Flag indicating whether or not to use the fast mamba kernels. These are available only if `mamba-ssm` and
+            `causal-conv1d` are installed, and the mamba modules are running on a CUDA device.
+        ssm_state_size (`int`, *optional*, defaults to 128):
+            The dimension of the mamba state space latents.
+        mamba_num_heads (`int`, *optional*, defaults to 128):
+            Number of heads in Mamba layers.
+        mamba_n_groups (`int`, *optional*, defaults to 8):
+            Number of groups in Mamba layers.
+        mamba_head_dim (`int`, *optional*, defaults to 64):
+            Dimension of each Mamba head.
+        mamba_d_conv (`int`, *optional*, defaults to 4):
+            The size of the mamba convolution kernel.
+        mamba_expand (`int`, *optional*, defaults to 2):
+            Expanding factor used to determine the mamba intermediate size.
+        mamba_hidden_act (`str`, *optional*, defaults to "silu"):
+            The non-linear activation function in the Mamba layers.
+        mamba_dt_min (`float`, *optional*, defaults to 0.001):
+            Minimum value for the time step in Mamba.
+        mamba_dt_max (`float`, *optional*, defaults to 0.1):
+            Maximum value for the time step in Mamba.
+        mamba_dt_limit (`tuple`, *optional*, defaults to (0.0, float("inf"))):
+            Limits for the time step in Mamba.
+        mamba_dt_init_floor (`float`, *optional*, defaults to 1e-4):
+            Floor value for time step initialization in Mamba.
+        mamba_conv_bias (`bool`, *optional*, defaults to `True`):
+            Whether to use bias in the convolution layer of the mamba mixer block.
+        mamba_proj_bias (`bool`, *optional*, defaults to `False`):
+            Whether to use bias in the input and output projections of the mamba mixer block.
+        mamba_chunk_size (`int`, *optional*, defaults to 256):
+            Size of chunks for Mamba processing.
+        rescale_prenorm_residual (`bool`, *optional*, defaults to `True`):
+            Whether to rescale the pre-normalization residual connections.
+    """
+
+    model_type = "nemotron_h"
+    keys_to_ignore_at_inference = ["past_key_values"]
+
+    def __init__(
+        self,
+        vocab_size=131072,
+        tie_word_embeddings=False,
+        hidden_size=4096,
+        intermediate_size=21504,
+        num_hidden_layers=52,
+        hybrid_override_pattern="M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M*-M-M-M-M-M-",
+        num_attention_heads=32,
+        head_dim=128,
+        num_key_value_heads=8,  # nemo: num_query_groups
+        mlp_hidden_act="relu2",
+        attention_bias=False,
+        mlp_bias=False,
+        use_bias=False,
+        initializer_range=0.02, # nemo: init_method_std
+        layer_norm_epsilon=1e-5, # nemo: layernorm_epsilon
+        residual_in_fp32=False,  #  Megatron Core default value
+        use_cache=True,
+        num_logits_to_keep=1,
+        pad_token_id=0,
+        bos_token_id=1,
+        eos_token_id=2,
+        sliding_window=None,
+        max_position_embeddings=4096,
+        attention_dropout=0.0,
+        hidden_dropout=0.0, # * ADDED
+        use_mamba_kernels=True,
+        ssm_state_size=128, # mamba_state_size
+        mamba_num_heads=128,
+        mamba_n_groups=8,  # nemo: mamba_ssm_ngroups = num_heads
+        mamba_head_dim=64,
+        mamba_d_conv=4,
+        mamba_expand=2,
+        mamba_hidden_act="silu",
+        mamba_dt_min=0.001,
+        mamba_dt_max=0.1,
+        mamba_dt_limit=(0.0, float("inf")),
+        mamba_dt_init_floor=1e-4,
+        mamba_conv_bias=True,
+        mamba_proj_bias=False,
+        mamba_chunk_size=128,
+        rescale_prenorm_residual=True,
+        n_routed_experts=8,
+        n_shared_experts=1,
+        moe_intermediate_size=7688,
+        moe_shared_expert_intermediate_size=7688,
+        num_experts_per_tok=2,
+        routed_scaling_factor=1.0,
+        n_group=1,
+        topk_group=1,
+        norm_topk_prob=True,
+        **kwargs,
+    ):
+        self.vocab_size = vocab_size
+        self.tie_word_embeddings = tie_word_embeddings
+        self.hidden_size = hidden_size
+        self.intermediate_size = intermediate_size
+        self.num_hidden_layers = num_hidden_layers
+        self.hybrid_override_pattern = hybrid_override_pattern
+        self.num_attention_heads = num_attention_heads
+        self.head_dim = head_dim
+        self.sliding_window = sliding_window
+        self.max_position_embeddings = max_position_embeddings
+        self.attention_dropout = attention_dropout
+        self.hidden_dropout = hidden_dropout
+
+        # Validate hybrid_override_pattern
+        # M: Mamba2, *: Attention, -: MLP, E: MoE
+        assert len(self.hybrid_override_pattern) == self.num_hidden_layers, "hybrid_override_pattern must have the same length as num_hidden_layers"
+        assert re.match(r"^[*\-ME]+$", self.hybrid_override_pattern), "hybrid_override_pattern must only contain characters 'M', '*', '-', or 'E'"
+
+        # for backward compatibility
+        if num_key_value_heads is None:
+            num_key_value_heads = num_attention_heads
+
+        self.num_key_value_heads = num_key_value_heads
+        self.mlp_hidden_act = mlp_hidden_act
+        self.attention_bias = attention_bias
+        self.mlp_bias = mlp_bias
+        self.use_bias = use_bias
+        self.initializer_range = initializer_range
+        self.layer_norm_epsilon = layer_norm_epsilon
+        self.residual_in_fp32 = residual_in_fp32
+
+        self.use_cache = use_cache
+        self.num_logits_to_keep = num_logits_to_keep
+
+        self.use_mamba_kernels = use_mamba_kernels
+        self.n_groups = mamba_n_groups
+        self.mamba_head_dim = mamba_head_dim
+        self.ssm_state_size = ssm_state_size
+        self.mamba_num_heads = mamba_num_heads
+        self.conv_kernel = mamba_d_conv
+        self.expand = mamba_expand
+        self.mamba_hidden_act = mamba_hidden_act
+        self.time_step_min = mamba_dt_min
+        self.time_step_max = mamba_dt_max
+        self.time_step_limit = mamba_dt_limit
+        self.time_step_floor = mamba_dt_init_floor
+        self.use_conv_bias = mamba_conv_bias
+        self.mamba_proj_bias = mamba_proj_bias
+        self.chunk_size = mamba_chunk_size
+        self.rescale_prenorm_residual = rescale_prenorm_residual
+        self.n_routed_experts = n_routed_experts
+        self.n_shared_experts = n_shared_experts
+        self.moe_intermediate_size = moe_intermediate_size
+        self.moe_shared_expert_intermediate_size = moe_shared_expert_intermediate_size
+        self.num_experts_per_tok = num_experts_per_tok
+        self.routed_scaling_factor = routed_scaling_factor
+        self.n_group = n_group
+        self.topk_group = topk_group
+        self.norm_topk_prob = norm_topk_prob
+
+        super().__init__(
+            pad_token_id=pad_token_id,
+            bos_token_id=bos_token_id,
+            eos_token_id=eos_token_id,
+            tie_word_embeddings=tie_word_embeddings,
+            **kwargs,
+        )
+
+    @property
+    def layers_block_type(self):
+        return [
+            "mamba" if self.hybrid_override_pattern[i] == "M" else
+            "attention" if self.hybrid_override_pattern[i] == "*" else
+            "mlp" if self.hybrid_override_pattern[i] == "-" else "moe"
+            for i in range(self.num_hidden_layers)]

configuration_radio.py

@@ -0,0 +1,152 @@
+# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.
+#
+# NVIDIA CORPORATION and its licensors retain all intellectual property
+# and proprietary rights in and to this software, related documentation
+# and any modifications thereto.  Any use, reproduction, disclosure or
+# distribution of this software and related documentation without an express
+# license agreement from NVIDIA CORPORATION is strictly prohibited.
+
+from dataclasses import dataclass
+from typing import Optional, NamedTuple, Union, List, Dict
+
+from transformers import PretrainedConfig
+
+
+class Resolution(NamedTuple):
+    height: int
+    width: int
+
+
+@dataclass
+class RadioResource:
+    url: str
+    patch_size: int
+    max_resolution: int
+    preferred_resolution: Resolution
+    vitdet_num_windowed: Optional[int] = None
+    vitdet_num_global: Optional[int] = None
+
+
+RESOURCE_MAP = {
+    # RADIOv2.5
+    "radio_v2.5-b": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio-v2.5-b_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-l": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio-v2.5-l_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-h": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-h.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-h-norm": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-h-norm.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=4,
+    ),
+    "radio_v2.5-g": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.5-g.pth.tar?download=true",
+        patch_size=14,
+        max_resolution=1792,
+        preferred_resolution=(896, 896),
+        vitdet_num_global=8,
+    ),
+    # RADIO
+    "radio_v2.1": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.1_bf16.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(432, 432),
+        vitdet_num_windowed=5,
+    ),
+    "radio_v2": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v2.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(432, 432),
+        vitdet_num_windowed=5,
+    ),
+    "radio_v1": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/radio_v1.pth.tar?download=true",
+        patch_size=14,
+        max_resolution=1050,
+        preferred_resolution=Resolution(378, 378),
+    ),
+    # E-RADIO
+    "e-radio_v2": RadioResource(
+        "https://huggingface.co/nvidia/RADIO/resolve/main/eradio_v2.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
+    # C-RADIO
+    "c-radio_v2.5-g": RadioResource(
+        "https://huggingface.co/nvidia/C-RADIOv2-g/resolve/main/c-radio_v2-g_half.pth.tar",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=(768, 768),
+        vitdet_num_global=8,
+    ),
+    "c-radio_v3-l": RadioResource(
+        # NOTE: Currently, this model cannot be loaded via TorchHub. Instead, use the transformers API at https://huggingface.co/nvidia/C-RADIOv3-L
+        # and accept the license terms.
+        "https://huggingface.co/nvidia/C-RADIOv3-L/resolve/main/c-radio-v3_l_half.pth.tar?download=true",
+        patch_size=16,
+        max_resolution=2048,
+        preferred_resolution=Resolution(512, 512),
+    ),
+}
+
+DEFAULT_VERSION = "radio_v2.5-h"
+
+
+class RADIOConfig(PretrainedConfig):
+    """Pretrained Hugging Face configuration for RADIO models."""
+
+    def __init__(
+        self,
+        args: Optional[dict] = None,
+        version: Optional[str] = DEFAULT_VERSION,
+        patch_size: Optional[int] = None,
+        max_resolution: Optional[int] = None,
+        preferred_resolution: Optional[Resolution] = None,
+        adaptor_names: Union[str, List[str]] = None,
+        adaptor_configs: Dict[str, Dict[str, int]] = None,
+        vitdet_window_size: Optional[int] = None,
+        feature_normalizer_config: Optional[dict] = None,
+        inter_feature_normalizer_config: Optional[dict] = None,
+        **kwargs,
+    ):
+        self.args = args
+        for field in ["dtype", "amp_dtype"]:
+            if self.args is not None and field in self.args:
+                # Convert to a string in order to make it serializable.
+                # For example for torch.float32 we will store "float32",
+                # for "bfloat16" we will store "bfloat16".
+                self.args[field] = str(args[field]).split(".")[-1]
+        self.version = version
+        resource = RESOURCE_MAP[version]
+        self.patch_size = patch_size or resource.patch_size
+        self.max_resolution = max_resolution or resource.max_resolution
+        self.preferred_resolution = (
+            preferred_resolution or resource.preferred_resolution
+        )
+        self.adaptor_names = adaptor_names
+        self.adaptor_configs = adaptor_configs
+        self.vitdet_window_size = vitdet_window_size
+        self.feature_normalizer_config = feature_normalizer_config
+        self.inter_feature_normalizer_config = inter_feature_normalizer_config
+        super().__init__(**kwargs)

evs.py

@@ -0,0 +1,73 @@
+import torch
+from typing import Tuple
+
+class EfficientVideoSampling:
+    @staticmethod
+    def compute_retention_mask(
+        *,
+        video_embeds: torch.FloatTensor,
+        thw: torch.LongTensor,
+        spatial_merge_size: int,
+        q: float,
+    ):
+        """
+        Computes the retention mask for video embeddings based on the grid dimensions.
+
+        Args:
+            video_embeds (`torch.FloatTensor` of shape `(T * H * W, hidden_size)`):
+                The video embeddings to compute the retention mask for.
+            thw (`torch.LongTensor` of shape `(3)`):
+                The temporal, height and width of feature shape of each video in LLM.
+            spatial_merge_size (`int`): The spatial merge size of the video embeddings.
+                If embeddings will be downsampled *later*, this should be the downsampling factor.
+            q: (`float`): Pruning rate factor, indicating number of tokens to prune (remove)
+
+        Returns:
+            `torch.Tensor`: The retention mask for the video embeddings (T * H * W).
+                1 for tokens to keep, 0 for tokens to prune.
+        """
+        T, H, W = thw
+
+        # video_embeds = einops.rearrange(
+        #     video_embeds,
+        #     "(T H W) C -> T H W C",
+        #     T=T,
+        #     H=H // spatial_merge_size,
+        #     W=W // spatial_merge_size,
+        # )
+        # Use reshape instead of einops to avoid graph breaks
+        video_embeds = video_embeds.reshape(
+            T, H // spatial_merge_size, W // spatial_merge_size, video_embeds.size(-1)
+        )
+
+        # Core EVS
+        similarity = torch.nn.functional.cosine_similarity(
+            video_embeds[1:, ...], video_embeds[:-1, ...], dim=-1
+        )
+        dissimilarity = 1 - similarity
+
+        # Always ensure we include all tokens from the first frame
+        dissimilarity = torch.cat(
+            [255 * torch.ones_like(video_embeds[:1, :, :, 0]), dissimilarity], dim=0
+        )
+        dissimilarity_flat = dissimilarity.view(-1)
+
+        min_num_tokens = (H // spatial_merge_size) * (W // spatial_merge_size)  # a single frame
+        evs_num_tokens = int(T * min_num_tokens * (1 - q))
+        num_tokens_to_keep = max(min_num_tokens, evs_num_tokens)
+
+        order = torch.argsort(dissimilarity_flat,
+                              dim=-1,
+                              descending=True,
+                              stable=True)
+        topk_indices = order[:num_tokens_to_keep]
+
+        retention_mask = torch.zeros_like(dissimilarity_flat, dtype=torch.bool)
+        retention_mask[topk_indices] = True
+        retention_mask = retention_mask.reshape(dissimilarity.size())
+
+        # print(
+        #     f"Computed retention mask of shape {retention_mask.shape=} with sparsity {retention_mask.float().mean().item():.4f} for {q=}",
+        # )
+        mask = retention_mask.view(-1)  # "T H W -> (T H W)"
+        return mask

generation_config.json

@@ -0,0 +1,14 @@
+{
+  "_from_model_config": true,
+  "bos_token_id": 1,
+  "eos_token_id": [2, 11],
+  "pad_token_id": 0,
+  "do_sample": true,
+  "temperature": 0.6,
+  "top_p": 0.95,
+  "max_new_tokens": 16384,
+  "reasoning_budget": 16384,
+  "reasoning_grace": 512,
+  "repetition_penalty": 1.0,
+  "transformers_version": "4.55.4"
+}

image_processing.py

@@ -0,0 +1,148 @@
+from typing import List, Optional, Union, Any, Dict
+
+from PIL import Image
+import torch
+from transformers.image_processing_base import BatchFeature
+from transformers.image_processing_utils_fast import BaseImageProcessorFast, divide_to_patches
+from transformers.image_utils import (make_list_of_images, get_image_size,
+                                      get_image_type, ImageInput, ImageType, ChannelDimension)
+from transformers.utils import TensorType
+import torchvision.transforms as T
+
+
+
+class NemotronH_Nano_Omni_Reasoning_V3ImageProcessor(BaseImageProcessorFast):
+    model_input_names = ["pixel_values"]
+
+    def __init__(self, image_size=512, max_num_tiles=12, use_thumbnail=True, norm_mean=None, norm_std=None, do_rescale=True, patch_size=16, downsample_ratio=0.5, **kwargs):
+        super().__init__(**kwargs)
+        self.image_size = image_size
+        self.max_num_tiles = max_num_tiles
+        self.use_thumbnail = use_thumbnail
+        self.norm_mean = norm_mean
+        self.norm_std = norm_std
+        self.do_rescale = do_rescale
+        self.num_image_token = int((image_size // patch_size) ** 2 * (downsample_ratio ** 2))
+
+    def _process_image(
+        self,
+        image: ImageInput,
+        **kwargs,
+    ) -> torch.Tensor:
+        image_type = get_image_type(image)
+        if image_type == ImageType.PIL:
+            if image.mode != 'RGB':
+                image = image.convert('RGB')
+            image = T.ToTensor()(image)
+        return image
+
+    def _preprocess(
+        self,
+        images: List[torch.Tensor],
+        image_size: int = None,
+        max_num_tiles: int = None,
+        use_thumbnail: bool = None,
+        do_rescale: bool = None,
+        return_tensors: Optional[Union[str, TensorType]] = None,
+        **kwargs,
+    ) -> List[torch.Tensor]:
+        image_size = image_size if image_size is not None else self.image_size
+        max_num_tiles = max_num_tiles if max_num_tiles is not None else self.max_num_tiles
+        use_thumbnail = use_thumbnail if use_thumbnail is not None else self.use_thumbnail
+        do_rescale = do_rescale if do_rescale is not None else self.do_rescale
+
+        images = make_list_of_images(images)
+
+        all_patches = []
+        num_patches = []
+        for image in images:
+            patches = dynamic_preprocess(image, image_size, max_num_tiles, use_thumbnail)
+            all_patches.extend(patches)
+            num_patches.append(len(patches))
+
+        pixel_values = torch.stack(all_patches, dim=0)
+        norm_mean = torch.Tensor(self.norm_mean).view(1, 3, 1, 1)
+        norm_std = torch.Tensor(self.norm_std).view(1, 3, 1, 1)
+        pixel_values = (pixel_values - norm_mean) / norm_std
+        return BatchFeature(data={"pixel_values": pixel_values, "num_patches": num_patches}, tensor_type=return_tensors)
+
+
+def get_internvl_target_ratios(
+    min_num: int,
+    max_num: int,
+) -> list[tuple[int, int]]:
+    target_ratios = {(i, j)
+                     for n in range(min_num, max_num + 1)
+                     for i in range(1, n + 1)
+                     for j in range(1, n + 1) if min_num <= i * j <= max_num}
+    return sorted(target_ratios, key=lambda x: x[0] * x[1])
+
+
+# From https://github.com/OpenGVLab/InternVL/blob/c62fa4f7c850165d7386bdc48ac6bc5a6fab0864/internvl_chat/internvl/train/dataset.py#L685
+# Copyright (c) 2023 OpenGVLab.
+def find_closest_aspect_ratio(
+    aspect_ratio: float,
+    target_ratios: list[tuple[int, int]],
+    width: int,
+    height: int,
+    image_size: int,
+) -> tuple[int, int]:
+    best_ratio_diff = float("inf")
+    best_ratio = (1, 1)
+    area = width * height
+    for ratio in target_ratios:
+        target_aspect_ratio = ratio[0] / ratio[1]
+        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
+        if ratio_diff < best_ratio_diff:
+            best_ratio_diff = ratio_diff
+            best_ratio = ratio
+        elif ratio_diff == best_ratio_diff:
+            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
+                best_ratio = ratio
+    return best_ratio
+
+
+def calculate_targets(
+    orig_width: int,
+    orig_height: int,
+    target_ratios: list[tuple[int, int]],
+    image_size: int,
+) -> tuple[int, int, int]:
+    aspect_ratio = orig_width / orig_height
+
+    # find the closest aspect ratio to the target
+    target_aspect_ratio = find_closest_aspect_ratio(
+        aspect_ratio,
+        target_ratios,
+        width=orig_width,
+        height=orig_height,
+        image_size=image_size,
+    )
+
+    # calculate the target width and height
+    target_width = image_size * target_aspect_ratio[0]
+    target_height = image_size * target_aspect_ratio[1]
+    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
+
+    return blocks, target_width, target_height
+
+
+def dynamic_preprocess(image, image_size=512, max_num_tiles=12, use_thumbnail=True):
+    orig_height, orig_width = get_image_size(image, channel_dim=ChannelDimension.FIRST)
+    target_ratios = get_internvl_target_ratios(1, max_num_tiles)
+
+    blocks, target_width, target_height = calculate_targets(
+        orig_width,
+        orig_height,
+        target_ratios,
+        image_size
+    )
+    # resize the image
+    resized_img = T.Resize((target_height, target_width), interpolation=T.InterpolationMode.BICUBIC)(image)
+    patches = divide_to_patches(resized_img, image_size)
+    assert len(patches) == blocks
+    if use_thumbnail and len(patches) != 1:
+        thumbnail_img = T.Resize((image_size, image_size), interpolation=T.InterpolationMode.BICUBIC)(image)
+        patches.append(thumbnail_img)
+
+    return patches

model-lm-00000.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c9f3af499b4c0cb30895da4f03b15e4e2c830bb3611a32992174fc5f9476af78
+size 743427376

model-lm-00001.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f5c0d5e27ba33f0c3200c811c2f5692410e53e0a50a13d20079f4ca2c642f1df
+size 731121536

model-lm-00002.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c3d70aaf2778ddba74c1a823576aed897638477200eb8e973e427cc836bca696
+size 38784120

model-lm-00003.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:88b0fa0fc6e5bc9ed1c1d91c4c3d070c8775d69a805510f9e3467a213fe79384
+size 731121536

model-lm-00004.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:98e7a7cde6e6bf799dc565bb993c81265010d494fbd7227d051ab7238dc66438
+size 38784120

model-lm-00005.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:0c04f4870be9538becb6cf2dc4031af8fc4ed0f43ef8a3cda2e6b9069e60b5da
+size 46798952

model-lm-00006.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:942d096467f659714b848ece369444f44da4efd72ca17d310f1b5a4e683163ec
+size 731121536

model-lm-00007.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ccc5b4f2529c6a20dbf0a47d306d48e0b646e1e550f76e4b7a3ee38dfc16f2a2
+size 38784120

model-lm-00008.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:4a1ee8c710d27a99ad395180af1cdb82be9813dc945fb574b7ea13f993ffbb6b
+size 731121536

model-lm-00009.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bf875675df08b717b9924de8048d10c09cfb0a56e8a11d1bc90fc1cfff32c859
+size 38784120

model-lm-00010.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2fd6c9931ccfed80f00d89cb97705b9178880401fa9c541e6ae0e729007890fb
+size 731122312

model-lm-00011.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9dafdb7c89d53ca0565ca14030512174c31e03fa2a075352d5195300eb791c43
+size 38784128

model-lm-00012.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cc61bd883be853844d6e741e25b223c2db164422d91af890a97724d868b8565b
+size 46798960

model-lm-00013.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b650001085bcc3be68ffd47dffff62c5f60f49ef415e08c4025d84daa05ee2d7
+size 731122312

model-lm-00014.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b35036638fb9b459c24924e6a1033a3443f2f6efb329cae0fe2d6a8939844603
+size 38784128

model-lm-00015.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:69ce75e9ca97da3f7c3f51c0d236cccd5929fdc11fc2c0c2fc38588bb5d97327
+size 731122312

model-lm-00016.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1a2dbf2ed2e8fe5a7a867e030f1f4f4e911be8226f8ce5c7d0a8502bb0e3bfb4
+size 38784128

model-lm-00017.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:00ab5e1cb48989aa56a1f5f29abade96c4a45eceb712b9b0327e360276f174a8
+size 731122312

model-lm-00018.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ddd42ea3fad64779baae013f34b25031cd3a9fd32d8fb931341fdcb7b15afd1d
+size 38784128

model-lm-00019.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8655c90c2ada6e8326840ddd8b39c851d22b18f7139fd7164b8e94766cc7acef
+size 46798960

model-lm-00020.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ba66fa32a4ecae79030cc7c3cdf4f0ecab7beee07d03c0c2b3b7b2a9bb6b8ef0
+size 731122312

model-lm-00021.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:db89967ea1bdf29ef784b3b3bc00f0ab33c848f969cf1eb00d3012e45372d267
+size 38784128

model-lm-00022.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:354e6f50d3f915c5107772600aaa64d99e3be8c9e34968f344ca8b7ca8ad191e
+size 731122312

model-lm-00023.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f7e631303ba302cefb512a37e0b81a8b8d923d3081b0ef59904846693be248a2
+size 38784128

model-lm-00024.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:72f723726d2b2b6c7d63fbde81b25e56c9d04fed6bd4ce5ca180158ddadb83b1
+size 731122312

model-lm-00025.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:50803a63db51b753509904c621c095b5ef176443b656b8fc55bc092e0a69db0e
+size 38784128

model-lm-00026.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:bef8ac9c2b2f43d0ed4165b7020e64272e6c9224458a11e0f6961b00f48f7f24
+size 46798960

model-lm-00027.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:46bf7ae3d22136da150d348abc8c867ce15f2e75a4773b16d24070d3deea30a7
+size 731122312

model-lm-00028.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9c4f689109a653295146049d2a013d343659f255468ea469f0f3609f799f179f
+size 38784128

model-lm-00029.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:30452ed27255dbb8a50f39b4b9bac10c9fb32a801dc88812b5d082a285acc3c1
+size 731122312

model-lm-00030.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:611862ea391e5ce73045a0c65527aa7fec4b4d01ea6da55c732dadcb4c5d973a
+size 38784128

model-lm-00031.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:c4764c09380c89550a83cb4774af2310bf4a75870400eff3f1282ee732c8843c
+size 731122312

model-lm-00032.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9bef7fce4690b7d515ee083c4778c83ef5d9953f330899fa894170ecc7a01260
+size 38784128

model-lm-00033.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:9b70e6e16d4fe6a7d7f89d941be647946d6a68feb55827cbdf2c93a43ab0e534
+size 46798960

model-lm-00034.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:92508517b0b47b87c0e10a6a5cb28925b65f25ca128dbedc39a91a4a4633067f
+size 731122312

model-lm-00035.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8c022a3fa8a5e0536e20e11ffb9825ad0d878d4674c350601cec74fd0bae91d4
+size 38784128

model-lm-00036.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3e522cf75166ef83747153ff01c3bf17b1d6799d41fc12fe559901143d0f0331
+size 731122312

model-lm-00037.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:86e46624d7f7c72e4a88b1cdad70de7dc0fa1295473251089f3d1f812ddcc369
+size 38784128

model-lm-00038.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:79fc589b86210b4a85599bea126d5da3f23376af5418537d2bccd7edce78d20b
+size 731122312