Files
vllm-omni/vllm_omni/diffusion/models/gr00t/configs/gr00t_n1d7.py
T
2026-06-20 11:45:26 +08:00

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)