Signed-off-by: Zhengyuan Su <su.zhengyuan@u.nus.edu> Co-authored-by: Claude Opus 4.8 (1M context) <noreply@anthropic.com>
110 lines
4.2 KiB
Python
110 lines
4.2 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
from transformers import PretrainedConfig
|
|
|
|
from . import register_model_config
|
|
|
|
|
|
def _default_diffusion_model_cfg() -> dict:
|
|
return {
|
|
"positional_embeddings": None,
|
|
"num_layers": 16,
|
|
"num_attention_heads": 32,
|
|
"attention_head_dim": 48,
|
|
"norm_type": "ada_norm",
|
|
"output_dim": 1024,
|
|
"interleave_self_attention": True,
|
|
}
|
|
|
|
|
|
class Gr00tN1d7Config(PretrainedConfig):
|
|
model_type = "Gr00tN1d7"
|
|
|
|
def __init__(
|
|
self,
|
|
model_dtype: str = "bfloat16",
|
|
model_name: str = "nvidia/Cosmos-Reason2-2B",
|
|
backbone_model_type: str = "qwen",
|
|
model_revision: str | None = None,
|
|
backbone_embedding_dim: int = 2048,
|
|
select_layer: int = 12,
|
|
reproject_vision: bool = False,
|
|
use_flash_attention: bool = True,
|
|
load_bf16: bool = False,
|
|
image_crop_size: tuple | None = (230, 230),
|
|
image_target_size: tuple | None = (256, 256),
|
|
shortest_image_edge: int | None = None,
|
|
crop_fraction: float | None = None,
|
|
random_rotation_angle: int | None = None,
|
|
color_jitter_params: dict | None = None,
|
|
formalize_language: bool = True,
|
|
apply_sincos_state_encoding: bool = False,
|
|
use_percentiles: bool = True,
|
|
use_relative_action: bool = False,
|
|
max_state_dim: int = 132,
|
|
max_action_dim: int = 132,
|
|
action_horizon: int = 40,
|
|
hidden_size: int = 1024,
|
|
input_embedding_dim: int = 1536,
|
|
state_history_length: int = 1,
|
|
add_pos_embed: bool = True,
|
|
attn_dropout: float = 0.2,
|
|
use_vlln: bool = True,
|
|
max_seq_len: int = 1024,
|
|
use_alternate_vl_dit: bool = True,
|
|
attend_text_every_n_blocks: int = 2,
|
|
diffusion_model_cfg: dict | None = None,
|
|
vl_self_attention_cfg: dict | None = None,
|
|
num_inference_timesteps: int = 4,
|
|
num_timestep_buckets: int = 1000,
|
|
exclude_state: bool = False,
|
|
use_mean_std: bool = False,
|
|
max_num_embodiments: int = 32,
|
|
**kwargs,
|
|
):
|
|
super().__init__(**kwargs)
|
|
self.model_dtype = model_dtype
|
|
self.model_name = model_name
|
|
self.backbone_model_type = backbone_model_type
|
|
self.model_revision = model_revision
|
|
self.backbone_embedding_dim = backbone_embedding_dim
|
|
self.select_layer = select_layer
|
|
self.reproject_vision = reproject_vision
|
|
self.use_flash_attention = use_flash_attention
|
|
self.load_bf16 = load_bf16
|
|
self.image_crop_size = image_crop_size
|
|
self.image_target_size = image_target_size
|
|
self.shortest_image_edge = shortest_image_edge
|
|
self.crop_fraction = crop_fraction
|
|
self.random_rotation_angle = random_rotation_angle
|
|
self.color_jitter_params = color_jitter_params
|
|
self.formalize_language = formalize_language
|
|
self.apply_sincos_state_encoding = apply_sincos_state_encoding
|
|
self.use_percentiles = use_percentiles
|
|
self.use_relative_action = use_relative_action
|
|
self.max_state_dim = max_state_dim
|
|
self.max_action_dim = max_action_dim
|
|
self.action_horizon = action_horizon
|
|
self.hidden_size = hidden_size
|
|
self.input_embedding_dim = input_embedding_dim
|
|
self.state_history_length = state_history_length
|
|
self.add_pos_embed = add_pos_embed
|
|
self.attn_dropout = attn_dropout
|
|
self.use_vlln = use_vlln
|
|
self.max_seq_len = max_seq_len
|
|
self.use_alternate_vl_dit = use_alternate_vl_dit
|
|
self.attend_text_every_n_blocks = attend_text_every_n_blocks
|
|
self.diffusion_model_cfg = (
|
|
diffusion_model_cfg if diffusion_model_cfg is not None else _default_diffusion_model_cfg()
|
|
)
|
|
self.vl_self_attention_cfg = vl_self_attention_cfg
|
|
self.num_inference_timesteps = num_inference_timesteps
|
|
self.num_timestep_buckets = num_timestep_buckets
|
|
self.exclude_state = exclude_state
|
|
self.use_mean_std = use_mean_std
|
|
self.max_num_embodiments = max_num_embodiments
|
|
|
|
|
|
register_model_config("Gr00tN1d7", Gr00tN1d7Config)
|