| """Offline initialization helpers for DiffusionVLA backbones. |
| |
| Diffusion-action checkpoints contain the trained vision, projector, LLM, and action decoder weights, but upstream model |
| construction still tries to download TIMM DINO/SigLIP pretrained weights and the LLaMA config/tokenizer before loading the |
| checkpoint. This patch prevents network access during construction and lets the local checkpoint provide the actual |
| weights. |
| """ |
| from __future__ import annotations |
|
|
| import os |
| from pathlib import Path |
|
|
| _MODEL_DIR = Path(__file__).resolve().parent |
| _TOKENIZER_DIR = Path(os.environ.get('VLA_TOKENIZER_DIR', _MODEL_DIR)) |
| _LLAMA_CONFIG_PATH = Path(os.environ.get('VLA_LLAMA_CONFIG_PATH', _MODEL_DIR / 'llama2_7b_config.json')) |
| _LLAMA2_7B_ID = 'meta-llama/Llama-2-7b-hf' |
|
|
|
|
| def patch_offline_backbone_loading() -> None: |
| """Patch TIMM and Transformers loaders so the diffusion-action policy can instantiate fully offline.""" |
| import timm |
| from transformers import AutoConfig, AutoTokenizer, LlamaConfig, LlamaForCausalLM |
|
|
| if not getattr(timm.create_model, '_local_vla_offline_patch', False): |
| original_create_model = timm.create_model |
|
|
| def create_model_offline(model_name, *args, **kwargs): |
| |
| if kwargs.get('pretrained', False) and ('dinov2' in str(model_name) or 'siglip' in str(model_name)): |
| kwargs = dict(kwargs) |
| kwargs['pretrained'] = False |
| return original_create_model(model_name, *args, **kwargs) |
|
|
| create_model_offline._local_vla_offline_patch = True |
| timm.create_model = create_model_offline |
|
|
| def local_llama_config() -> LlamaConfig: |
| if _LLAMA_CONFIG_PATH.exists(): |
| return LlamaConfig.from_json_file(str(_LLAMA_CONFIG_PATH)) |
| return LlamaConfig( |
| vocab_size=32000, |
| hidden_size=4096, |
| intermediate_size=11008, |
| num_hidden_layers=32, |
| num_attention_heads=32, |
| num_key_value_heads=32, |
| max_position_embeddings=4096, |
| rms_norm_eps=1e-5, |
| rope_theta=10000.0, |
| pad_token_id=0, |
| bos_token_id=1, |
| eos_token_id=2, |
| tie_word_embeddings=False, |
| ) |
|
|
| if not getattr(AutoConfig.from_pretrained, '_local_vla_offline_patch', False): |
| original_config_from_pretrained = AutoConfig.from_pretrained |
|
|
| def config_from_pretrained_offline(model_name_or_path, *args, **kwargs): |
| if str(model_name_or_path) == _LLAMA2_7B_ID: |
| return local_llama_config() |
| return original_config_from_pretrained(model_name_or_path, *args, **kwargs) |
|
|
| config_from_pretrained_offline._local_vla_offline_patch = True |
| AutoConfig.from_pretrained = config_from_pretrained_offline |
|
|
| if not getattr(LlamaConfig.from_pretrained, '_local_vla_offline_patch', False): |
| original_llama_config_from_pretrained = LlamaConfig.from_pretrained |
|
|
| def llama_config_from_pretrained_offline(model_name_or_path, *args, **kwargs): |
| if str(model_name_or_path) == _LLAMA2_7B_ID: |
| return local_llama_config() |
| return original_llama_config_from_pretrained(model_name_or_path, *args, **kwargs) |
|
|
| llama_config_from_pretrained_offline._local_vla_offline_patch = True |
| LlamaConfig.from_pretrained = llama_config_from_pretrained_offline |
|
|
| if not getattr(LlamaForCausalLM.from_pretrained, '_local_vla_offline_patch', False): |
| original_llama_from_pretrained = LlamaForCausalLM.from_pretrained |
|
|
| def llama_from_pretrained_offline(model_name_or_path, *args, **kwargs): |
| if str(model_name_or_path) == _LLAMA2_7B_ID: |
| return LlamaForCausalLM(local_llama_config()) |
| return original_llama_from_pretrained(model_name_or_path, *args, **kwargs) |
|
|
| llama_from_pretrained_offline._local_vla_offline_patch = True |
| LlamaForCausalLM.from_pretrained = llama_from_pretrained_offline |
|
|
| if not getattr(AutoTokenizer.from_pretrained, '_local_vla_offline_patch', False): |
| original_tokenizer_from_pretrained = AutoTokenizer.from_pretrained |
|
|
| def tokenizer_from_pretrained_offline(model_name_or_path, *args, **kwargs): |
| if str(model_name_or_path) == _LLAMA2_7B_ID: |
| kwargs = dict(kwargs) |
| kwargs.pop('token', None) |
| return original_tokenizer_from_pretrained(str(_TOKENIZER_DIR), *args, **kwargs) |
| return original_tokenizer_from_pretrained(model_name_or_path, *args, **kwargs) |
|
|
| tokenizer_from_pretrained_offline._local_vla_offline_patch = True |
| AutoTokenizer.from_pretrained = tokenizer_from_pretrained_offline |
|
|
