File size: 4,591 Bytes
0190e6b | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 | # SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
# SPDX-License-Identifier: Apache-2.0
#
# 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.
"""Helpers for Alpamayo 2 Super inference examples."""
import collections.abc
from typing import Any
import torch
from transformers import AutoProcessor, AutoTokenizer
from alpamayo2_super.chat_template.conversation import build_conversation
from alpamayo2_super.config import Alpamayo2SuperConfig, resolve_checkpoint_name_or_path
def create_messages(
data: dict[str, Any],
model_config: Alpamayo2SuperConfig,
) -> list[dict[str, Any]]:
"""Create generation-mode chat messages for CoT + trajectory inference."""
messages = build_conversation(
data=data,
num_tokens_per_history_traj=model_config.tokens_per_history_traj,
num_tokens_per_future_traj=model_config.tokens_per_future_traj,
components_order=["image", "traj_history", "prompt"],
components_prompt=["cot", "traj_future"],
generation_mode=True,
include_camera_ids=model_config.include_camera_ids,
camera_ids=data["camera_indices"],
include_frame_nums=model_config.frame_label == "frame_num",
)
if messages[-1]["role"] == "assistant" and not messages[-1]["content"]:
messages = messages[:-1]
return messages
def get_processor(
tokenizer: AutoTokenizer,
model_config: Alpamayo2SuperConfig,
) -> AutoProcessor:
"""Load the checkpoint processor and attach the Alpamayo tokenizer."""
name_or_path = resolve_checkpoint_name_or_path(model_config)
if name_or_path is None:
raise ValueError("model_config must include a checkpoint path or vlm_name_or_path")
processor_kwargs = {"fix_mistral_regex": True}
if model_config.min_pixels is not None:
processor_kwargs["min_pixels"] = model_config.min_pixels
if model_config.max_pixels is not None:
processor_kwargs["max_pixels"] = model_config.max_pixels
processor = AutoProcessor.from_pretrained(name_or_path, **processor_kwargs)
processor.tokenizer = tokenizer
return processor
def prepare_model_inputs(
data: dict[str, Any],
model_config: Alpamayo2SuperConfig,
tokenizer: AutoTokenizer,
) -> dict[str, Any]:
"""Tokenize one PhysicalAI-AV sample for ``Alpamayo2Super.sample_trajectories_from_data``."""
processor = get_processor(tokenizer, model_config)
messages = create_messages(data, model_config)
has_assistant_content = messages[-1]["role"] == "assistant" and bool(messages[-1]["content"])
text = processor.apply_chat_template(
messages,
tokenize=False,
add_generation_prompt=not has_assistant_content,
add_vision_id=False,
continue_final_message=has_assistant_content,
)
images = data["image_frames"].flatten(0, 1)
images = (images.float() / 255.0) if images.dtype == torch.uint8 else images.float()
tokenized_data = dict(
processor(
text=text,
images=images,
videos=None,
padding=False,
return_tensors="pt",
do_rescale=False,
)
)
if tokenized_data["input_ids"].shape[0] != 1:
raise ValueError("prepare_model_inputs expects one sample at a time")
return {
"tokenized_data": tokenized_data,
"ego_history_xyz": data["ego_history_xyz"],
"ego_history_rot": data["ego_history_rot"],
}
def to_device(
data: Any,
device: str | torch.device | None = None,
dtype: torch.dtype | None = None,
) -> Any:
"""Recursively move tensors to ``device`` and optional ``dtype``."""
if isinstance(data, torch.Tensor):
return data.to(device=device, dtype=dtype)
if isinstance(data, collections.abc.Mapping):
return {key: to_device(value, device=device, dtype=dtype) for key, value in data.items()}
if isinstance(data, collections.abc.Sequence) and not isinstance(data, (str, bytes)):
return [to_device(elem, device=device, dtype=dtype) for elem in data]
return data
|