115 lines
4.4 KiB
Python
115 lines
4.4 KiB
Python
# Copyright 2025 SGLang Team
|
|
# 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.
|
|
# ==============================================================================
|
|
# Adapted from https://huggingface.co/nvidia/NVIDIA-Nemotron-Nano-12B-v2-VL-BF16/blob/cb5a65ff10232128389d882d805fa609427544f1/configuration.py
|
|
|
|
from typing import Any
|
|
|
|
from transformers.configuration_utils import PretrainedConfig
|
|
|
|
from sglang.srt.configs.nemotron_h import NemotronHConfig
|
|
from sglang.srt.configs.radio import RadioConfig
|
|
from sglang.srt.multimodal.internvl_utils import IMAGENET_MEAN, IMAGENET_STD
|
|
|
|
|
|
def float_triplet(seq: Any):
|
|
a, b, c = tuple(seq)
|
|
assert (
|
|
isinstance(a, float) and isinstance(b, float) and isinstance(c, float)
|
|
), "expected three floats"
|
|
return a, b, c
|
|
|
|
|
|
class NemotronH_Nano_VL_V2_Config(PretrainedConfig):
|
|
model_type = "NemotronH_Nano_VL_V2"
|
|
is_composition = True
|
|
|
|
def __init__(
|
|
self,
|
|
vision_config=None,
|
|
llm_config=None,
|
|
force_image_size: int = 512,
|
|
patch_size: int = 16,
|
|
downsample_ratio=0.5,
|
|
template=None,
|
|
ps_version="v2",
|
|
image_tag_type="internvl",
|
|
projector_hidden_size=4096,
|
|
vit_hidden_size=1280,
|
|
video_pruning_rate: float = 0.0,
|
|
video_context_token: str = "<video>",
|
|
img_context_token: str = "<image>",
|
|
img_start_token: str = "<img>",
|
|
img_end_token: str = "</img>",
|
|
norm_mean: tuple[float, float, float] | list[float] = IMAGENET_MEAN,
|
|
norm_std: tuple[float, float, float] | list[float] = IMAGENET_STD,
|
|
use_thumbnail: bool = True,
|
|
**kwargs,
|
|
):
|
|
super().__init__(**kwargs)
|
|
|
|
# Handle both cases: when loading from JSON (llm_config is dict) and when called internally by transformers (llm_config; vision_config are None)
|
|
if llm_config is not None:
|
|
self.llm_config = NemotronHConfig(**llm_config)
|
|
assert isinstance(vision_config, dict), "vision_config must be a dictionary"
|
|
self.raw_vision_config = vision_config
|
|
else:
|
|
assert vision_config is None
|
|
self.llm_config = NemotronHConfig()
|
|
self.raw_vision_config = {}
|
|
|
|
# Assign configuration values
|
|
vision_image_size = self.raw_vision_config.get("image_size", force_image_size)
|
|
vision_patch_size = self.raw_vision_config.get("patch_size", patch_size)
|
|
self.image_size = int(
|
|
vision_image_size[0]
|
|
if isinstance(vision_image_size, list)
|
|
else vision_image_size
|
|
)
|
|
self.patch_size = int(
|
|
vision_patch_size[0]
|
|
if isinstance(vision_patch_size, list)
|
|
else vision_patch_size
|
|
)
|
|
|
|
self.downsample_ratio = downsample_ratio
|
|
self.video_context_token = video_context_token
|
|
self.img_context_token = img_context_token
|
|
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
|
|
|
|
self.norm_mean = float_triplet(norm_mean)
|
|
self.norm_std = float_triplet(norm_std)
|
|
self.use_thumbnail = use_thumbnail
|
|
self.img_start_token = img_start_token
|
|
self.img_end_token = img_end_token
|
|
|
|
def create_radio_config(self):
|
|
config = self.raw_vision_config
|
|
model_name = config["args"]["model"]
|
|
reg_tokens = config["args"].get("register_multiple")
|
|
image_size = config.get("preferred_resolution", [224])[0]
|
|
radio_config = RadioConfig(
|
|
patch_size=self.patch_size,
|
|
norm_mean=self.norm_mean,
|
|
norm_std=self.norm_std,
|
|
model_name=model_name,
|
|
reg_tokens=reg_tokens,
|
|
image_size=image_size,
|
|
)
|
|
return radio_config
|