yangkun.zhuandwangbolun
0016006974
Signed-off-by: zhuyangkun <zyk_120403228@163.com> Signed-off-by: bl-wang1 <918395762@qq.com> Co-authored-by: wangbolun <918395762@qq.com>
127 lines
4.1 KiB
Python
127 lines
4.1 KiB
Python
from __future__ import annotations
|
|
|
|
import json
|
|
import os
|
|
from dataclasses import dataclass, field
|
|
from dataclasses import fields as dataclass_fields
|
|
from pathlib import Path
|
|
from typing import Any
|
|
|
|
from huggingface_hub.constants import HF_HOME
|
|
|
|
ACTION = "action"
|
|
OBS_STR = "observation"
|
|
OBS_PREFIX = OBS_STR + "."
|
|
OBS_STATE = OBS_STR + ".state"
|
|
OBS_TASK = OBS_STR + ".task"
|
|
OBS_IMAGE = OBS_STR + ".image"
|
|
OBS_IMAGES = OBS_IMAGE + "s"
|
|
OPENPI_ATTENTION_MASK_VALUE = -2.3819763e38
|
|
|
|
DEFAULT_QWEN3_VL_MODEL = os.getenv("INTERNVLA_A1_PROCESSOR_DIR", "Qwen/Qwen3-VL-2B-Instruct")
|
|
DEFAULT_COSMOS_REPO = "nvidia/Cosmos-Tokenizer-CI8x8"
|
|
DEFAULT_COSMOS_DIR = Path(
|
|
os.getenv(
|
|
"INTERNVLA_A1_COSMOS_DIR",
|
|
str(Path(HF_HOME) / "hub" / "Cosmos-Tokenizer-CI8x8"),
|
|
)
|
|
)
|
|
|
|
|
|
@dataclass
|
|
class InternVLAA1Config:
|
|
"""Standalone-compatible InternVLA-A1 config with a few fake-smoke defaults."""
|
|
|
|
type: str = "internvla_a1"
|
|
qwen3_vl_variant: str = "qwen3_vl_28l"
|
|
action_expert_variant: str = "qwen3_28l"
|
|
dtype: str = "bfloat16"
|
|
device: str = "cuda"
|
|
chunk_size: int = 50
|
|
n_action_steps: int = 50
|
|
max_state_dim: int = 32
|
|
max_action_dim: int = 32
|
|
num_inference_steps: int = 10
|
|
time_sampling_beta_alpha: float = 1.5
|
|
time_sampling_beta_beta: float = 1.0
|
|
time_sampling_scale: float = 0.999
|
|
time_sampling_offset: float = 0.001
|
|
min_period: float = 4e-3
|
|
max_period: float = 4.0
|
|
image_resolution: tuple[int, int] = (224, 224)
|
|
empty_cameras: int = 0
|
|
gradient_checkpointing: bool = False
|
|
compile_model: bool = False
|
|
compile_mode: str = "max-autotune"
|
|
enable_regional_compile: bool = False
|
|
regional_compile_dynamic: bool = True
|
|
enable_suffix_static_context_optimization: bool = False
|
|
attn_implementation: str = "eager"
|
|
tokenizer_max_length: int = 48
|
|
freeze_vision_encoder: bool = False
|
|
train_expert_only: bool = False
|
|
train_vlm_only: bool = False
|
|
scale_factor: int = 8
|
|
lambda_gen: float = 0.01
|
|
input_features: dict[str, Any] = field(default_factory=dict)
|
|
output_features: dict[str, Any] = field(default_factory=dict)
|
|
|
|
# Fake/smoke-only knobs used by the lightweight pipeline path.
|
|
vocab_size: int = 256
|
|
hidden_size: int = 128
|
|
intermediate_size: int = 256
|
|
num_attention_heads: int = 4
|
|
num_hidden_layers: int = 2
|
|
num_cameras: int = 3
|
|
image_history: int = 2
|
|
pixel_feature_dim: int = 48
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, checkpoint_dir: str | Path) -> InternVLAA1Config:
|
|
checkpoint_dir = Path(checkpoint_dir)
|
|
with open(checkpoint_dir / "config.json", encoding="utf-8") as f:
|
|
raw = json.load(f)
|
|
return cls.from_model_config(raw)
|
|
|
|
@classmethod
|
|
def from_model_config(cls, model_config: dict[str, Any] | None) -> InternVLAA1Config:
|
|
if not model_config:
|
|
return cls()
|
|
|
|
raw = dict(model_config)
|
|
raw["image_resolution"] = tuple(raw.get("image_resolution", [224, 224]))
|
|
allowed = {item.name for item in dataclass_fields(cls)}
|
|
filtered = {key: value for key, value in raw.items() if key in allowed}
|
|
return cls(**filtered)
|
|
|
|
|
|
@dataclass
|
|
class InternVLAA1TrainMetadata:
|
|
action_mode: str = "delta"
|
|
processor_model_name: str = DEFAULT_QWEN3_VL_MODEL
|
|
|
|
@classmethod
|
|
def from_pretrained(cls, checkpoint_dir: str | Path) -> InternVLAA1TrainMetadata:
|
|
checkpoint_dir = Path(checkpoint_dir)
|
|
path = checkpoint_dir / "train_config.json"
|
|
if not path.exists():
|
|
return cls()
|
|
|
|
with open(path, encoding="utf-8") as f:
|
|
raw = json.load(f)
|
|
|
|
dataset_cfg = raw.get("dataset", {})
|
|
processor_model_name = DEFAULT_QWEN3_VL_MODEL
|
|
for transform in dataset_cfg.get("data_transforms", {}).get("inputs", []):
|
|
if transform.get("type") == "internvla_a1_processor":
|
|
processor_model_name = transform.get(
|
|
"pretrained_model_name_or_path",
|
|
processor_model_name,
|
|
)
|
|
break
|
|
|
|
return cls(
|
|
action_mode=dataset_cfg.get("action_mode", "delta"),
|
|
processor_model_name=processor_model_name,
|
|
)
|