From 0016006974b331d36bdaea581abc35f169fde24e Mon Sep 17 00:00:00 2001 From: "yangkun.zhu" <34429807+Greyman-Seu@users.noreply.github.com> Date: Mon, 27 Apr 2026 17:28:07 +0800 Subject: [PATCH] [Model] Add InternVLA-A1 offline inference support (#2737) Signed-off-by: zhuyangkun Signed-off-by: bl-wang1 <918395762@qq.com> Co-authored-by: wangbolun <918395762@qq.com> --- docs/.nav.yml | 1 + docs/models/supported_models.md | 1 + .../offline_inference/internvla_a1.md | 205 ++++ .../offline_inference/internvla_a1/README.md | 30 + .../internvla_a1/collect_results.sh | 220 ++++ .../offline_inference/internvla_a1/end2end.py | 327 ++++++ .../internvla_a1/internvla_a1_common.py | 471 +++++++++ .../offline_inference/internvla_a1/run.sh | 23 + .../offline_inference/test_internvla_a1.py | 90 ++ .../diffusion/models/internvla_a1/__init__.py | 25 + .../models/internvla_a1/adapter_qwen3_vl.py | 252 +++++ .../diffusion/models/internvla_a1/config.py | 126 +++ .../models/internvla_a1/cosmos_ci_torch.py | 416 ++++++++ .../models/internvla_a1/model_cosmos.py | 97 ++ .../models/internvla_a1/model_internvla_a1.py | 962 ++++++++++++++++++ .../internvla_a1/pipeline_internvla_a1.py | 257 +++++ vllm_omni/diffusion/registry.py | 6 + 17 files changed, 3509 insertions(+) create mode 100644 docs/user_guide/examples/offline_inference/internvla_a1.md create mode 100644 examples/offline_inference/internvla_a1/README.md create mode 100644 examples/offline_inference/internvla_a1/collect_results.sh create mode 100755 examples/offline_inference/internvla_a1/end2end.py create mode 100644 examples/offline_inference/internvla_a1/internvla_a1_common.py create mode 100755 examples/offline_inference/internvla_a1/run.sh create mode 100644 tests/e2e/offline_inference/test_internvla_a1.py create mode 100644 vllm_omni/diffusion/models/internvla_a1/__init__.py create mode 100644 vllm_omni/diffusion/models/internvla_a1/adapter_qwen3_vl.py create mode 100644 vllm_omni/diffusion/models/internvla_a1/config.py create mode 100644 vllm_omni/diffusion/models/internvla_a1/cosmos_ci_torch.py create mode 100644 vllm_omni/diffusion/models/internvla_a1/model_cosmos.py create mode 100644 vllm_omni/diffusion/models/internvla_a1/model_internvla_a1.py create mode 100644 vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py diff --git a/docs/.nav.yml b/docs/.nav.yml index b44e8e6b5..ddcef5d7a 100644 --- a/docs/.nav.yml +++ b/docs/.nav.yml @@ -22,6 +22,7 @@ nav: - HunyuanImage-3.0 Image-to-Text Inference: user_guide/examples/offline_inference/hunyuan_image3.md - Image-To-Image: user_guide/examples/offline_inference/image_to_image.md - Image-To-Video: user_guide/examples/offline_inference/image_to_video.md + - InternVLA-A1: user_guide/examples/offline_inference/internvla_a1.md - MammothModa2-Preview: user_guide/examples/offline_inference/mammothmodal2_preview.md - MiMo-Audio Offline Inference: user_guide/examples/offline_inference/mimo_audio.md - Qwen2.5-Omni: user_guide/examples/offline_inference/qwen2_5_omni.md diff --git a/docs/models/supported_models.md b/docs/models/supported_models.md index feea969e5..b81ff5b99 100644 --- a/docs/models/supported_models.md +++ b/docs/models/supported_models.md @@ -20,6 +20,7 @@ th { | `Qwen3OmniMoeForConditionalGeneration` | Qwen3-Omni | `Qwen/Qwen3-Omni-30B-A3B-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `Qwen2_5OmniForConditionalGeneration` | Qwen2.5-Omni | `Qwen/Qwen2.5-Omni-7B`, `Qwen/Qwen2.5-Omni-3B` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `BagelForConditionalGeneration` | BAGEL (DiT-only) | `ByteDance-Seed/BAGEL-7B-MoT` | ✅︎ | ✅︎ | | ✅︎ | +| `InternVLAA1Pipeline` | InternVLA-A1 | `InternRobotics/InternVLA-A1-3B` | ✅︎ | ✅︎ | | | | `HunyuanImage3ForCausalMM` | HunyuanImage3.0 (DiT-only) | `tencent/HunyuanImage-3.0`, `tencent/HunyuanImage-3.0-Instruct` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `QwenImagePipeline` | Qwen-Image | `Qwen/Qwen-Image` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | | `QwenImagePipeline` | Qwen-Image-2512 | `Qwen/Qwen-Image-2512` | ✅︎ | ✅︎ | ✅︎ | ✅︎ | diff --git a/docs/user_guide/examples/offline_inference/internvla_a1.md b/docs/user_guide/examples/offline_inference/internvla_a1.md new file mode 100644 index 000000000..83906a037 --- /dev/null +++ b/docs/user_guide/examples/offline_inference/internvla_a1.md @@ -0,0 +1,205 @@ +# InternVLA-A1 + +Source . + + +## Setup +This example is adapted from: https://github.com/InternRobotics/InternVLA-A1/blob/master/tests/policies/internvla_a1_3b/open_loop_genie1_real.ipynb + +This example runs the single-path vLLM offline inference workflow for InternVLA-A1 open-loop action prediction. + +Before running the script, export the required local paths: + +```bash +# hf Jia-Zeng/InternVLA-A1-3B-FineTuned-Place_Markpen +export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen +# hf download InternRobotics/InternData-A1 real_lerobotv30/genie1/Genie1-Place_Markpen.tar.gz --repo-type dataset --local-dir /path/to/Genie1-Place_Markpen +export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen +export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct +# hf tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors +export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor +``` + +The shell entrypoint also accepts these variables as defaults, so you can keep the command line itself short. + +`INTERNVLA_A1_COSMOS_DIR` is expected to contain: + +- `encoder.safetensors` +- `decoder.safetensors` + +Reference Hugging Face repo: `tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors` + +## Run examples + +Get into the example folder: + +```bash +cd examples/offline_inference/internvla_a1 +``` + +### Run one sample + +This runs a single sample through the maintained vLLM registry path. + +```bash +bash run.sh --num-samples 1 --num-episodes 0 +``` + +### Run open-loop evaluation against GT + +This runs one episode, writes evaluation metrics, and saves prediction-vs-GT plots when `matplotlib` is available. + +```bash +bash run.sh --num-episodes 1 +``` + +Outputs are written under: + +```bash +outputs/internvla_a1/vllm_infer/ +``` + +Typical files: + +- `summary.json`: top-level run summary +- `registry/log.json`: per-episode GT comparison metrics +- `registry/plots/*.jpg`: prediction-vs-GT figures + +### Eager float32 validation + +Use eager attention and float32 when you want the most stable numerical comparison baseline. + +```bash +bash run.sh \ + --num-episodes 1 \ + --attn-implementation eager \ + --dtype float32 +``` + +### Optional runtime switches + +- `--dtype {bfloat16,float32}`: choose inference dtype +- `--attn-implementation {eager,sdpa}`: switch attention backend +- `--enable-regional-compile`: enable regional `torch.compile` +- `--enable-warmup`: run pipeline warmup in initialization +- `--skip-plots`: skip plot generation even if `matplotlib` is installed + +### Collect results and performance logs + +Use the helper below when you want to collect one-sample output, forward latency, GT comparison, plots, timing logs, and optional pytest output into a single result directory: + +```bash +bash collect_results.sh +``` + +The script writes a timestamped directory under: + +```bash +outputs/internvla_a1/collected_results/ +``` + +Typical result files: + +- `env_summary.txt`: environment and path summary +- `sample_run.log`: one-sample inference log +- `forward_benchmark/forward_latency.json`: isolated `pipeline.forward` latency summary +- `eval_run.log`: open-loop GT evaluation log +- `eval_outputs/summary.json`: top-level output summary +- `eval_outputs/registry/log.json`: GT comparison metrics +- `eval_outputs/registry/plots/*.jpg`: prediction-vs-GT figures +- `*_time.txt`: timing output +- `*_gpu.csv`: sampled GPU usage + +### Benchmark forward latency + +Use the dedicated mode below when you want the isolated `pipeline.forward` latency instead of end-to-end script time: + +```bash +python end2end.py \ + --model-dir "$INTERNVLA_A1_MODEL_DIR" \ + --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \ + --benchmark-forward \ + --dtype bfloat16 \ + --attn-implementation eager \ + --warmup-iters 3 \ + --benchmark-iters 10 \ + --output-dir outputs/internvla_a1/forward_benchmark +``` + +The output JSON is written to `outputs/internvla_a1/forward_benchmark/forward_latency.json` and contains: + +- `cold_start_ms` +- `warmup_summary` +- `benchmark_summary` +- `benchmark_samples_ms` + +### Reference results + +Reference run collected on `1x NVIDIA H200`, `bfloat16`, `eager`: + +- one-sample end-to-end run: `38s` +- one-episode GT evaluation run: `45s` +- `average_mse = 1.7173260857816786e-05` +- `average_mae = 0.0011860118247568607` +- `average_mse_joint = 7.42028441891307e-06` +- `average_mae_joint = 0.0010777723509818316` +- `average_mse_gripper = 8.544408774469048e-05` +- `average_mae_gripper = 0.0019436875591054559` + +The reference plot filename is: + +- `eval_outputs/registry/plots/vllm_registry_open_loop_ep0.jpg` + +## Measure latency and VRAM + +You can collect end-to-end runtime and peak host memory with: + +```bash +/usr/bin/time -v bash run.sh --num-episodes 1 +``` + +For GPU memory snapshots, use: + +```bash +nvidia-smi --query-gpu=name,memory.total,memory.used --format=csv +``` + +Run `nvidia-smi` before and during inference to record the peak device memory for your PR report. + +## FAQ + +If `matplotlib` is missing, evaluation still runs and only plot generation is skipped. To enable plots: + +```bash +pip install matplotlib +``` + +## Example materials + +- `end2end.py`: main offline inference and GT evaluation entrypoint +- `run.sh`: shell wrapper with local path env vars +- `collect_results.sh`: helper to gather result summaries and performance logs into one directory +- `internvla_a1_common.py`: dataset, evaluation, and plotting helpers + +## Embedded source listings + +??? abstract "collect_results.sh" + ``````sh + --8<-- "examples/offline_inference/internvla_a1/collect_results.sh" + `````` +??? abstract "end2end.py" + ``````py + --8<-- "examples/offline_inference/internvla_a1/end2end.py" + `````` +??? abstract "internvla_a1_common.py" + ``````py + --8<-- "examples/offline_inference/internvla_a1/internvla_a1_common.py" + `````` +??? abstract "run.sh" + ``````sh + --8<-- "examples/offline_inference/internvla_a1/run.sh" + `````` +??? abstract "standalone/outputs/qwena1/standalone_infer/log.json" + ``````json + --8<-- "examples/offline_inference/internvla_a1/standalone/outputs/qwena1/standalone_infer/log.json" + `````` diff --git a/examples/offline_inference/internvla_a1/README.md b/examples/offline_inference/internvla_a1/README.md new file mode 100644 index 000000000..bb7fd3abb --- /dev/null +++ b/examples/offline_inference/internvla_a1/README.md @@ -0,0 +1,30 @@ +# InternVLA-A1 + +Full usage and result-reporting guidance lives in [docs/user_guide/examples/offline_inference/internvla_a1.md](../../../docs/user_guide/examples/offline_inference/internvla_a1.md). + +Quick start: + +```bash +export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen +export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen +export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct +# hf tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors +export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor + +bash run.sh --num-samples 1 --num-episodes 0 +bash run.sh --num-episodes 1 +bash collect_results.sh +``` + +Expected files under `INTERNVLA_A1_COSMOS_DIR`: + +- `encoder.safetensors` +- `decoder.safetensors` + +Reference Hugging Face repo: `tenstep/Cosmos-Tokenizer-CI8x8-SafeTensors` + +Key entrypoints: + +- `run.sh`: wrapper for offline inference and GT evaluation +- `collect_results.sh`: collect sample output, latency, metrics, plots, and logs +- `end2end.py`: underlying Python entrypoint diff --git a/examples/offline_inference/internvla_a1/collect_results.sh b/examples/offline_inference/internvla_a1/collect_results.sh new file mode 100644 index 000000000..6f4556a34 --- /dev/null +++ b/examples/offline_inference/internvla_a1/collect_results.sh @@ -0,0 +1,220 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$ROOT_DIR/../../.." && pwd)" +TIMESTAMP="$(date +%Y%m%d_%H%M%S)" +RESULT_ROOT="${INTERNVLA_A1_RESULT_DIR:-$REPO_ROOT/outputs/internvla_a1/collected_results/$TIMESTAMP}" + +: "${INTERNVLA_A1_MODEL_DIR:?Please export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen}" +: "${INTERNVLA_A1_DATASET_DIR:?Please export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen}" +: "${INTERNVLA_A1_PROCESSOR_DIR:?Please export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct}" +: "${INTERNVLA_A1_COSMOS_DIR:?Please export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor}" + +EVAL_OUTPUT_DIR="$RESULT_ROOT/eval_outputs" +GPU_MONITOR_INTERVAL="${GPU_MONITOR_INTERVAL:-1}" + +mkdir -p "$RESULT_ROOT" + +export INTERNVLA_A1_MODEL_DIR +export INTERNVLA_A1_DATASET_DIR +export INTERNVLA_A1_PROCESSOR_DIR +export INTERNVLA_A1_COSMOS_DIR + +write_env_summary() { + cat >"$RESULT_ROOT/env_summary.txt" <&1 || true) +model_dir=$INTERNVLA_A1_MODEL_DIR +dataset_dir=$INTERNVLA_A1_DATASET_DIR +processor_dir=$INTERNVLA_A1_PROCESSOR_DIR +cosmos_dir=$INTERNVLA_A1_COSMOS_DIR +pwd=$(pwd) +EOF +} + +capture_gpu_snapshot() { + local output_file="$1" + if command -v nvidia-smi >/dev/null 2>&1; then + nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv >"$output_file" 2>&1 || true + else + echo "nvidia-smi not found" >"$output_file" + fi +} + +start_gpu_monitor() { + local output_file="$1" + if ! command -v nvidia-smi >/dev/null 2>&1; then + echo "nvidia-smi not found" >"$output_file" + return 1 + fi + + ( + echo "timestamp,index,name,memory.total [MiB],memory.used [MiB],utilization.gpu [%]" + while true; do + local now + now="$(date '+%Y-%m-%d %H:%M:%S')" + nvidia-smi --query-gpu=index,name,memory.total,memory.used,utilization.gpu --format=csv,noheader,nounits \ + | awk -F',' -v ts="$now" '{gsub(/^[ \t]+|[ \t]+$/, "", $0); print ts "," $0}' + sleep "$GPU_MONITOR_INTERVAL" + done + ) >"$output_file" 2>/dev/null & + echo $! +} + +stop_gpu_monitor() { + local monitor_pid="${1:-}" + if [[ -n "$monitor_pid" ]] && kill -0 "$monitor_pid" 2>/dev/null; then + kill "$monitor_pid" 2>/dev/null || true + wait "$monitor_pid" 2>/dev/null || true + fi +} + +copy_eval_outputs() { + local source_dir="$1" + local target_dir="$2" + local source_real="" + local target_real="" + source_real="$(realpath "$source_dir" 2>/dev/null || echo "$source_dir")" + target_real="$(realpath "$target_dir" 2>/dev/null || echo "$target_dir")" + if [[ "$source_real" == "$target_real" ]]; then + return 0 + fi + mkdir -p "$target_dir" + if [[ -d "$source_dir" ]]; then + cp -r "$source_dir"/. "$target_dir"/ + fi +} + +run_with_artifacts() { + local name="$1" + shift + + local log_file="$RESULT_ROOT/${name}.log" + local time_file="$RESULT_ROOT/${name}_time.txt" + local gpu_file="$RESULT_ROOT/${name}_gpu.csv" + local status_file="$RESULT_ROOT/${name}_status.txt" + local monitor_pid="" + + monitor_pid="$(start_gpu_monitor "$gpu_file" || true)" + set +e + local exit_code=0 + local start_ts=0 + local end_ts=0 + start_ts="$(date +%s)" + if [[ -x /usr/bin/time ]]; then + /usr/bin/time -v -o "$time_file" "$@" >"$log_file" 2>&1 + exit_code=$? + if [[ ! -s "$time_file" ]]; then + end_ts="$(date +%s)" + { + echo "timing_mode=usr_bin_time_empty_fallback" + echo "start_ts=$start_ts" + echo "end_ts=$end_ts" + echo "elapsed_seconds=$((end_ts - start_ts))" + } >"$time_file" + fi + else + "$@" >"$log_file" 2>&1 + exit_code=$? + end_ts="$(date +%s)" + { + echo "timing_mode=shell_date" + echo "start_ts=$start_ts" + echo "end_ts=$end_ts" + echo "elapsed_seconds=$((end_ts - start_ts))" + } >"$time_file" + fi + set -e + stop_gpu_monitor "$monitor_pid" + + echo "exit_code=$exit_code" >"$status_file" + if [[ $exit_code -ne 0 ]]; then + echo "[error] command failed for $name, see $log_file" + return $exit_code + fi +} + +write_manifest() { + cat >"$RESULT_ROOT/README.txt" <"$RESULT_ROOT/${name}.log" + echo "timing_mode=skipped" >"$RESULT_ROOT/${name}_time.txt" + echo "exit_code=0" >"$RESULT_ROOT/${name}_status.txt" + echo "skipped_reason=$reason" >>"$RESULT_ROOT/${name}_status.txt" + if [[ ! -f "$RESULT_ROOT/${name}_gpu.csv" ]]; then + echo "skipped,$reason" >"$RESULT_ROOT/${name}_gpu.csv" + fi +} + +write_env_summary +write_manifest +capture_gpu_snapshot "$RESULT_ROOT/gpu_info_before.csv" + +run_with_artifacts \ + "sample_run" \ + bash "$ROOT_DIR/run.sh" \ + --output-dir "$RESULT_ROOT/sample_outputs" \ + --num-samples 1 \ + --num-episodes 0 + +run_with_artifacts \ + "forward_benchmark" \ + python "$ROOT_DIR/end2end.py" \ + --model-dir "$INTERNVLA_A1_MODEL_DIR" \ + --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \ + --benchmark-forward \ + --dtype bfloat16 \ + --attn-implementation eager \ + --warmup-iters 3 \ + --benchmark-iters 10 \ + --output-dir "$RESULT_ROOT/forward_benchmark" + +run_with_artifacts \ + "eval_run" \ + bash "$ROOT_DIR/run.sh" \ + --output-dir "$EVAL_OUTPUT_DIR" \ + --num-episodes 1 + +copy_eval_outputs "$EVAL_OUTPUT_DIR" "$RESULT_ROOT/eval_outputs" + +if python - <<'PY' >/dev/null 2>&1 +import importlib.util +raise SystemExit(0 if importlib.util.find_spec("pytest") else 1) +PY +then + run_with_artifacts \ + "pytest_e2e" \ + python -m pytest -sv tests/e2e/offline_inference/test_internvla_a1.py -m advanced_model +else + write_skip_artifact "pytest_e2e" "pytest is not installed in the current python environment" +fi + +capture_gpu_snapshot "$RESULT_ROOT/gpu_info_after.csv" + +echo "Results written to: $RESULT_ROOT" diff --git a/examples/offline_inference/internvla_a1/end2end.py b/examples/offline_inference/internvla_a1/end2end.py new file mode 100755 index 000000000..a7799d7ce --- /dev/null +++ b/examples/offline_inference/internvla_a1/end2end.py @@ -0,0 +1,327 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import json +import os +import statistics +import sys +import time +from pathlib import Path + +import torch + +_WORKSPACE_ROOT = Path(__file__).resolve().parent +_REPO_ROOT = _WORKSPACE_ROOT.parents[2] +if str(_REPO_ROOT) not in sys.path: + # Keep the shared helper import clean by doing script-only path bootstrap here. + sys.path.insert(0, str(_REPO_ROOT)) + +from internvla_a1_common import ( # noqa: E402 + A2DOpenLoopDataset, + collate_open_loop_samples, + make_shared_noise, + run_open_loop_evaluation, + select_indices, + tensor_dtype, + tensor_sha256, +) + +from vllm_omni.diffusion.data import OmniDiffusionConfig # noqa: E402 +from vllm_omni.diffusion.models.internvla_a1 import ( # noqa: E402 + InternVLAA1Config, + InternVLAA1TrainMetadata, +) +from vllm_omni.diffusion.models.internvla_a1.config import OBS_STATE # noqa: E402 +from vllm_omni.diffusion.registry import initialize_model # noqa: E402 +from vllm_omni.diffusion.request import OmniDiffusionRequest # noqa: E402 +from vllm_omni.inputs.data import OmniDiffusionSamplingParams # noqa: E402 + + +def _required_path_arg(env_name: str, cli_value: str | None) -> str: + value = cli_value or os.getenv(env_name) + if not value: + raise ValueError(f"Missing required path: set --{env_name.lower().replace('_', '-')} or {env_name}.") + return value + + +def parse_args() -> argparse.Namespace: + parser = argparse.ArgumentParser(description="Run vLLM InternVLA-A1 inference on a few samples.") + parser.add_argument("--model-dir") + parser.add_argument("--dataset-dir") + parser.add_argument("--num-samples", type=int, default=1) + parser.add_argument("--num-episodes", type=int, default=0) + parser.add_argument("--seed", type=int, default=1234) + parser.add_argument("--dtype", choices=["bfloat16", "float32"], default="bfloat16") + parser.add_argument("--device", default="cuda") + parser.add_argument("--compile-model", action="store_true") + parser.add_argument("--attn-implementation", default="eager") + parser.add_argument("--enable-regional-compile", action="store_true") + parser.add_argument("--enable-warmup", action="store_true") + parser.add_argument("--strict-load", action="store_true") + parser.add_argument("--output-dir", default="outputs/internvla_a1/vllm_infer") + parser.add_argument("--skip-plots", action="store_true") + parser.add_argument("--benchmark-forward", action="store_true") + parser.add_argument("--warmup-iters", type=int, default=3) + parser.add_argument("--benchmark-iters", type=int, default=10) + return parser.parse_args() + + +def build_od_config(args: argparse.Namespace, processor_model_name: str) -> OmniDiffusionConfig: + return OmniDiffusionConfig( + model=str(Path(args.model_dir).resolve()), + model_class_name="InternVLAA1Pipeline", + dtype=tensor_dtype(args.dtype), + custom_pipeline_args={ + "device": args.device, + "dtype": args.dtype, + "compile_model": args.compile_model, + "attn_implementation": args.attn_implementation, + "enable_regional_compile": args.enable_regional_compile, + "enable_warmup": args.enable_warmup, + "strict_load": args.strict_load, + "processor_model_name": processor_model_name, + }, + ) + + +def build_dataset(args: argparse.Namespace) -> tuple[A2DOpenLoopDataset, InternVLAA1Config, InternVLAA1TrainMetadata]: + model_dir = Path(args.model_dir) + config = InternVLAA1Config.from_pretrained(model_dir) + config.device = args.device + config.dtype = args.dtype + config.compile_model = args.compile_model + config.attn_implementation = args.attn_implementation + config.enable_regional_compile = args.enable_regional_compile + + train_meta = InternVLAA1TrainMetadata.from_pretrained(model_dir) + with open(model_dir / "stats.json", encoding="utf-8") as f: + train_stats = json.load(f)["a2d"] + + dataset = A2DOpenLoopDataset( + args.dataset_dir, + config=config, + train_stats=train_stats, + ) + return dataset, config, train_meta + + +def run_one_path( + pipeline, + dataset: A2DOpenLoopDataset, + config: InternVLAA1Config, + args: argparse.Namespace, + indices: list[int], +) -> list[dict[str, object]]: + results: list[dict[str, object]] = [] + for index in indices: + sample = dataset.get_sample(index) + batch_inputs, _ = collate_open_loop_samples([sample], device=args.device, dtype=tensor_dtype(args.dtype)) + noise = make_shared_noise( + args.seed, + index, + ( + batch_inputs[OBS_STATE].shape[0], + config.chunk_size, + config.max_action_dim, + ), + args.device, + ) + pred = run_pipeline_forward(pipeline, batch_inputs, noise) + pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu() + results.append( + { + "path": "registry", + "index": index, + "episode_index": sample.episode_index, + "task": sample.task, + "seed": args.seed, + "shape": list(pred.shape), + "mean": float(pred.mean().item()), + "std": float(pred.std().item()), + "action_sha256": tensor_sha256(pred), + "first_action_prefix": pred[0, 0, :8].tolist(), + } + ) + return results + + +def run_pipeline_forward( + pipeline, + batch_inputs: dict[str, torch.Tensor], + noise: torch.Tensor, +) -> torch.Tensor: + output = pipeline.forward( + OmniDiffusionRequest( + prompts=[""], + sampling_params=OmniDiffusionSamplingParams( + extra_args={ + "batch_inputs": batch_inputs, + "noise": noise, + "decode_image": False, + } + ), + ) + ) + if output.error: + raise RuntimeError(output.error) + if output.output is None: + raise RuntimeError("InternVLAA1Pipeline.forward returned no output tensor.") + return output.output + + +def _synchronize(device: str) -> None: + if device.startswith("cuda") and torch.cuda.is_available(): + torch.cuda.synchronize() + + +def _latency_summary(values_ms: list[float]) -> dict[str, float]: + sorted_values = sorted(values_ms) + p50_index = min(len(sorted_values) - 1, round(0.50 * (len(sorted_values) - 1))) + p90_index = min(len(sorted_values) - 1, round(0.90 * (len(sorted_values) - 1))) + return { + "mean_ms": float(statistics.mean(sorted_values)), + "stdev_ms": float(statistics.pstdev(sorted_values)) if len(sorted_values) > 1 else 0.0, + "min_ms": float(sorted_values[0]), + "max_ms": float(sorted_values[-1]), + "p50_ms": float(sorted_values[p50_index]), + "p90_ms": float(sorted_values[p90_index]), + } + + +def benchmark_forward( + pipeline, + dataset: A2DOpenLoopDataset, + config: InternVLAA1Config, + args: argparse.Namespace, + index: int, + output_dir: Path, +) -> dict[str, object]: + sample = dataset.get_sample(index) + batch_inputs, _ = collate_open_loop_samples([sample], device=args.device, dtype=tensor_dtype(args.dtype)) + noise = make_shared_noise( + args.seed, + index, + ( + batch_inputs[OBS_STATE].shape[0], + config.chunk_size, + config.max_action_dim, + ), + args.device, + ) + + _synchronize(args.device) + cold_start_begin = time.perf_counter() + pred = run_pipeline_forward(pipeline, batch_inputs, noise) + _synchronize(args.device) + cold_start_ms = (time.perf_counter() - cold_start_begin) * 1000.0 + + warmup_ms: list[float] = [] + for _ in range(args.warmup_iters): + _synchronize(args.device) + begin = time.perf_counter() + _ = run_pipeline_forward(pipeline, batch_inputs, noise) + _synchronize(args.device) + warmup_ms.append((time.perf_counter() - begin) * 1000.0) + + benchmark_ms: list[float] = [] + for _ in range(args.benchmark_iters): + _synchronize(args.device) + begin = time.perf_counter() + _ = run_pipeline_forward(pipeline, batch_inputs, noise) + _synchronize(args.device) + benchmark_ms.append((time.perf_counter() - begin) * 1000.0) + + pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu() + summary = { + "mode": "forward_latency", + "model_dir": str(Path(args.model_dir).resolve()), + "dataset_dir": str(Path(args.dataset_dir).resolve()), + "sample_index": index, + "episode_index": sample.episode_index, + "task": sample.task, + "device": args.device, + "dtype": args.dtype, + "attn_implementation": args.attn_implementation, + "enable_regional_compile": args.enable_regional_compile, + "warmup_iters": args.warmup_iters, + "benchmark_iters": args.benchmark_iters, + "output_shape": list(pred.shape), + "cold_start_ms": cold_start_ms, + "warmup_summary": _latency_summary(warmup_ms) if warmup_ms else {}, + "benchmark_summary": _latency_summary(benchmark_ms) if benchmark_ms else {}, + "benchmark_samples_ms": benchmark_ms, + } + with open(output_dir / "forward_latency.json", "w", encoding="utf-8") as f: + json.dump(summary, f, ensure_ascii=False, indent=2) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + return summary + + +def main() -> None: + args = parse_args() + args.model_dir = _required_path_arg("INTERNVLA_A1_MODEL_DIR", args.model_dir) + args.dataset_dir = _required_path_arg("INTERNVLA_A1_DATASET_DIR", args.dataset_dir) + output_dir = Path(args.output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + dataset, config, train_meta = build_dataset(args) + indices = select_indices(dataset, args.num_samples) + od_config = build_od_config(args, train_meta.processor_model_name) + + eval_summaries: dict[str, object] = {} + pipeline = initialize_model(od_config) + if args.benchmark_forward: + benchmark_forward( + pipeline=pipeline, + dataset=dataset, + config=config, + args=args, + index=indices[0], + output_dir=output_dir, + ) + return + results = run_one_path( + pipeline=pipeline, + dataset=dataset, + config=config, + args=args, + indices=indices, + ) + if args.num_episodes > 0: + eval_summaries["registry"] = run_open_loop_evaluation( + mode="vllm_registry", + policy=pipeline, + config=config, + dataset=dataset, + train_meta=train_meta, + collate_samples=collate_open_loop_samples, + run_sample_actions=lambda policy, batch_inputs, noise: run_pipeline_forward(policy, batch_inputs, noise), + num_episodes=args.num_episodes, + seed=args.seed, + device=args.device, + dtype=tensor_dtype(args.dtype), + output_dir=output_dir / "registry", + skip_plots=args.skip_plots, + ) + + summary = { + "mode": "registry", + "model_dir": str(Path(args.model_dir).resolve()), + "dataset_dir": str(Path(args.dataset_dir).resolve()), + "device": args.device, + "dtype": args.dtype, + "attn_implementation": args.attn_implementation, + "enable_regional_compile": args.enable_regional_compile, + "seed": args.seed, + "indices": indices, + "results": results, + "output_dir": str(output_dir.resolve()), + "eval_summaries": eval_summaries, + } + with open(output_dir / "summary.json", "w", encoding="utf-8") as f: + json.dump(summary, f, ensure_ascii=False, indent=2) + print(json.dumps(summary, ensure_ascii=False, indent=2)) + + +if __name__ == "__main__": + main() diff --git a/examples/offline_inference/internvla_a1/internvla_a1_common.py b/examples/offline_inference/internvla_a1/internvla_a1_common.py new file mode 100644 index 000000000..806e14538 --- /dev/null +++ b/examples/offline_inference/internvla_a1/internvla_a1_common.py @@ -0,0 +1,471 @@ +from __future__ import annotations + +import hashlib +import json +import math +from collections.abc import Callable +from dataclasses import dataclass +from pathlib import Path +from typing import Any + +import numpy as np +import pyarrow.parquet as pq +import torch +import torch.nn.functional as F +import torchvision + +from vllm_omni.diffusion.models.internvla_a1.config import ( # noqa: E402 + OBS_IMAGES, + OBS_STATE, + OBS_TASK, + InternVLAA1Config, +) + + +def tensor_dtype(name: str) -> torch.dtype: + return torch.bfloat16 if name == "bfloat16" else torch.float32 + + +def select_indices(dataset, num_samples: int) -> list[int]: + indices: list[int] = [] + for _, episode_indices in dataset.episode_start_indices(max_episodes=num_samples): + if episode_indices: + indices.append(episode_indices[0]) + return indices[:num_samples] + + +def make_shared_noise(seed: int, sample_index: int, shape: tuple[int, ...], device: str) -> torch.Tensor: + generator = torch.Generator(device="cpu") + generator.manual_seed(seed + sample_index) + noise = torch.randn(shape, generator=generator, dtype=torch.float32) + return noise.to(device=device) + + +def tensor_sha256(tensor: torch.Tensor) -> str: + array = tensor.detach().contiguous().cpu().numpy() + return hashlib.sha256(array.tobytes()).hexdigest() + + +def _clamp_index(index: int, start: int, end: int) -> int: + return max(start, min(end - 1, index)) + + +def _load_parquet_rows(path: Path) -> list[dict[str, Any]]: + return pq.read_table(path).to_pylist() + + +def _load_json(path: Path) -> dict[str, Any]: + with open(path, encoding="utf-8") as f: + return json.load(f) + + +def _stack_stats(stats: dict[str, Any], keys: list[str]) -> dict[str, torch.Tensor]: + result = {} + for stat_name in ("mean", "std"): + values = [] + for key in keys: + values.extend(stats[key][stat_name]) + result[stat_name] = torch.tensor(values, dtype=torch.float32) + return result + + +def normalize_vector(values: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor: + denom = torch.where(stats["std"] == 0, torch.ones_like(stats["std"]), stats["std"]) + return (values - stats["mean"]) / denom + + +def unnormalize_vector(values: torch.Tensor, stats: dict[str, torch.Tensor]) -> torch.Tensor: + return values * stats["std"] + stats["mean"] + + +class TorchvisionVideoReaderCache: + def __init__(self, backend: str = "pyav") -> None: + self.backend = backend + self._readers: dict[str, Any] = {} + torchvision.set_video_backend(backend) + + def get(self, path: str) -> Any: + reader = self._readers.get(path) + if reader is None: + reader = torchvision.io.VideoReader(path, "video") + self._readers[path] = reader + return reader + + def decode_frames(self, path: str, timestamps: list[float], tolerance_s: float = 1e-4) -> torch.Tensor: + reader = self.get(path) + first_ts = min(timestamps) + last_ts = max(timestamps) + reader.seek(first_ts, keyframes_only=self.backend == "pyav") + + loaded_frames: list[torch.Tensor] = [] + loaded_ts: list[float] = [] + for frame in reader: + current_ts = float(frame["pts"]) + loaded_frames.append(frame["data"]) + loaded_ts.append(current_ts) + if current_ts >= last_ts: + break + + query_ts = torch.tensor(timestamps, dtype=torch.float32) + loaded_ts_tensor = torch.tensor(loaded_ts, dtype=torch.float32) + distances = torch.cdist(query_ts[:, None], loaded_ts_tensor[:, None], p=1) + min_dist, argmin = distances.min(dim=1) + if not torch.all(min_dist < tolerance_s): + raise RuntimeError( + f"Video timestamps are outside tolerance: query={query_ts.tolist()} " + f"loaded={loaded_ts_tensor.tolist()} path={path}" + ) + return torch.stack([loaded_frames[i] for i in argmin]).float() / 255.0 + + +@dataclass +class A2DOpenLoopSample: + index: int + episode_index: int + task: str + state_raw: torch.Tensor + action_raw: torch.Tensor + inputs: dict[str, torch.Tensor] + + +class A2DOpenLoopDataset: + image_keys = [ + "observation.images.head", + "observation.images.hand_left", + "observation.images.hand_right", + ] + state_keys = [ + "observation.states.joint.position", + "observation.states.effector.position", + ] + action_keys = [ + "actions.joint.position", + "actions.effector.position", + ] + + def __init__( + self, + dataset_root: str | Path, + *, + config: InternVLAA1Config, + train_stats: dict[str, Any], + image_offsets: tuple[int, int] = (-15, 0), + tolerance_s: float = 1e-4, + ) -> None: + self.root = Path(dataset_root) + self.config = config + self.info = _load_json(self.root / "meta" / "info.json") + self.dataset_stats = _load_json(self.root / "meta" / "stats.json") + self.data_rows = _load_parquet_rows(self.root / "data" / "chunk-000" / "file-000.parquet") + self.episode_rows = _load_parquet_rows(self.root / "meta" / "episodes" / "chunk-000" / "file-000.parquet") + self.task_rows = _load_parquet_rows(self.root / "meta" / "tasks.parquet") + + self.state_stats = _stack_stats(train_stats, self.state_keys) + self.action_stats = _stack_stats(train_stats, self.action_keys) + self.image_offsets = image_offsets + self.tolerance_s = tolerance_s + self.video_reader = TorchvisionVideoReaderCache(backend="pyav") + + @property + def num_episodes(self) -> int: + return len(self.episode_rows) + + @property + def physical_action_dim(self) -> int: + return 16 + + def episode_start_indices(self, max_episodes: int | None = None) -> list[tuple[int, list[int]]]: + rows = self.episode_rows if max_episodes is None else self.episode_rows[:max_episodes] + result = [] + for ep in rows: + start = int(ep["dataset_from_index"]) + end = int(ep["dataset_to_index"]) + result.append((int(ep["episode_index"]), list(range(start, end, self.config.chunk_size)))) + return result + + def _task_text(self, task_index: int) -> str: + return self.task_rows[task_index]["__index_level_0__"] + + def _episode_for_index(self, row: dict[str, Any]) -> dict[str, Any]: + return self.episode_rows[int(row["episode_index"])] + + def _state_vector(self, row: dict[str, Any]) -> torch.Tensor: + return torch.tensor(row[self.state_keys[0]] + row[self.state_keys[1]], dtype=torch.float32) + + def _action_vector(self, row: dict[str, Any]) -> torch.Tensor: + return torch.tensor(row[self.action_keys[0]] + row[self.action_keys[1]], dtype=torch.float32) + + def _query_rows(self, idx: int, deltas: list[int]) -> list[dict[str, Any]]: + row = self.data_rows[idx] + episode = self._episode_for_index(row) + start = int(episode["dataset_from_index"]) + end = int(episode["dataset_to_index"]) + return [self.data_rows[_clamp_index(idx + delta, start, end)] for delta in deltas] + + def _decode_camera_history( + self, episode: dict[str, Any], camera_key: str, rows: list[dict[str, Any]] + ) -> torch.Tensor: + timestamps = [float(r["timestamp"]) for r in rows] + shifted = [float(episode[f"videos/{camera_key}/from_timestamp"]) + ts for ts in timestamps] + chunk_idx = int(episode[f"videos/{camera_key}/chunk_index"]) + file_idx = int(episode[f"videos/{camera_key}/file_index"]) + path = self.root / self.info["video_path"].format( + video_key=camera_key, + chunk_index=chunk_idx, + file_index=file_idx, + ) + frames = self.video_reader.decode_frames(str(path), shifted, tolerance_s=self.tolerance_s) + return frames + + def get_sample(self, idx: int) -> A2DOpenLoopSample: + row = self.data_rows[idx] + episode = self._episode_for_index(row) + image_rows = self._query_rows(idx, list((-15, 0))) + action_rows = self._query_rows(idx, list(range(self.config.chunk_size))) + camera_images = [self._decode_camera_history(episode, camera_key, image_rows) for camera_key in self.image_keys] + state_raw = self._state_vector(row) + state_norm = normalize_vector(state_raw, self.state_stats) + action_raw = torch.stack([self._action_vector(action_row) for action_row in action_rows], dim=0) + task = self._task_text(int(row["task_index"])) + inputs = { + OBS_STATE: state_norm, + OBS_TASK: task, + f"{OBS_IMAGES}.image0": camera_images[0], + f"{OBS_IMAGES}.image1": camera_images[1], + f"{OBS_IMAGES}.image2": camera_images[2], + f"{OBS_IMAGES}.image0_mask": torch.tensor(True), + f"{OBS_IMAGES}.image1_mask": torch.tensor(True), + f"{OBS_IMAGES}.image2_mask": torch.tensor(True), + } + return A2DOpenLoopSample( + index=idx, + episode_index=int(row["episode_index"]), + task=task, + state_raw=state_raw, + action_raw=action_raw, + inputs=inputs, + ) + + +def collate_open_loop_samples( + samples: list[A2DOpenLoopSample], + *, + device: str, + dtype: torch.dtype, +) -> tuple[dict[str, torch.Tensor], dict[str, Any]]: + first = samples[0] + batch_inputs: dict[str, torch.Tensor] = {} + for key in first.inputs: + values = [sample.inputs[key] for sample in samples] + if isinstance(values[0], torch.Tensor): + tensor = torch.stack(values, dim=0) + if tensor.dtype in (torch.int64, torch.bool): + batch_inputs[key] = tensor.to(device=device) + else: + batch_inputs[key] = tensor.to(device=device, dtype=dtype) + else: + batch_inputs[key] = values + + metadata = { + "indices": [sample.index for sample in samples], + "episode_indices": [sample.episode_index for sample in samples], + "tasks": [sample.task for sample in samples], + "state_raw": torch.stack([sample.state_raw for sample in samples], dim=0), + "action_raw": torch.stack([sample.action_raw for sample in samples], dim=0), + } + return batch_inputs, metadata + + +def plot_prediction_series( + *, + series: dict[str, np.ndarray], + output_path: Path, + title: str, +) -> None: + import matplotlib.pyplot as plt + + first_series = next(iter(series.values())) + num_dims = int(first_series.shape[1]) + num_cols = 2 + num_rows = math.ceil(num_dims / num_cols) + fig, axs = plt.subplots(num_rows, num_cols, figsize=(16, max(6, 3 * num_rows))) + axs = np.array(axs).reshape(-1) + x_values = np.arange(first_series.shape[0]) + styles = [ + ("Ground Truth", "blue", "-", 1.5), + ("VLLM", "green", "-.", 1.4), + ("Predicted", "red", "--", 1.4), + ] + style_map = {name: (color, linestyle, linewidth) for name, color, linestyle, linewidth in styles} + + for dim in range(num_dims): + ax = axs[dim] + for name, values in series.items(): + color, linestyle, linewidth = style_map.get(name, (None, "-", 1.2)) + ax.plot( + x_values, + values[:, dim], + label=name, + color=color, + linestyle=linestyle, + linewidth=linewidth, + ) + ax.set_title(f"Dimension {dim + 1}") + ax.set_xlabel("Time Step / Sample Index") + ax.set_ylabel(f"Value Dim {dim + 1}") + ax.legend(loc="upper right") + ax.grid(True, linestyle="--", alpha=0.7) + + for dim in range(num_dims, len(axs)): + axs[dim].axis("off") + + fig.suptitle(title, fontsize=16) + plt.tight_layout(rect=(0, 0, 1, 0.97)) + output_path.parent.mkdir(parents=True, exist_ok=True) + plt.savefig(output_path) + plt.close(fig) + + +def summarize_prediction_metrics( + *, + gt_tensor: torch.Tensor, + pred_tensor: torch.Tensor, + joint_dims: int = 14, +) -> dict[str, float]: + return { + "mse": float(F.mse_loss(gt_tensor, pred_tensor, reduction="mean").item()), + "mae": float(F.l1_loss(gt_tensor, pred_tensor, reduction="mean").item()), + "mse_joint": float(F.mse_loss(gt_tensor[:, :joint_dims], pred_tensor[:, :joint_dims], reduction="mean").item()), + "mae_joint": float(F.l1_loss(gt_tensor[:, :joint_dims], pred_tensor[:, :joint_dims], reduction="mean").item()), + "mse_gripper": float( + F.mse_loss(gt_tensor[:, joint_dims:], pred_tensor[:, joint_dims:], reduction="mean").item() + ), + "mae_gripper": float( + F.l1_loss(gt_tensor[:, joint_dims:], pred_tensor[:, joint_dims:], reduction="mean").item() + ), + } + + +def summarize_metric_list(values: list[float]) -> float | None: + return float(np.mean(values)) if values else None + + +def run_open_loop_evaluation( + *, + mode: str, + policy, + config, + dataset, + train_meta, + collate_samples: Callable[..., tuple[dict[str, torch.Tensor], dict[str, Any]]], + run_sample_actions: Callable[[Any, dict[str, torch.Tensor], torch.Tensor], torch.Tensor], + num_episodes: int, + seed: int, + device: str, + dtype: torch.dtype, + output_dir: str | Path, + skip_plots: bool, +) -> dict[str, Any]: + output_dir = Path(output_dir) + output_dir.mkdir(parents=True, exist_ok=True) + plots_dir = output_dir / "plots" + plots_dir.mkdir(parents=True, exist_ok=True) + + mse_values: list[float] = [] + mae_values: list[float] = [] + mse_joint_values: list[float] = [] + mae_joint_values: list[float] = [] + mse_gripper_values: list[float] = [] + mae_gripper_values: list[float] = [] + per_episode: list[dict[str, Any]] = [] + + plotting_available = True + if not skip_plots: + try: + import matplotlib.pyplot as _plt # noqa: F401 + except ModuleNotFoundError: + plotting_available = False + print("[warning] matplotlib is not installed; plots will be skipped.") + + episode_specs = dataset.episode_start_indices(max_episodes=num_episodes) + for episode_index, indices in episode_specs: + print(f"[{mode}] episode: {episode_index}") + pred_chunks = [] + gt_chunks = [] + visited_indices = [] + task = None + + for index in indices: + sample = dataset.get_sample(index) + task = sample.task + batch_inputs, meta = collate_samples([sample], device=device, dtype=dtype) + noise = make_shared_noise( + seed, + index, + (batch_inputs[OBS_STATE].shape[0], config.chunk_size, config.max_action_dim), + device, + ) + with torch.no_grad(): + pred = run_sample_actions(policy, batch_inputs, noise) + + pred = pred[:, :, : dataset.physical_action_dim].to(torch.float32).cpu() + pred_phys = unnormalize_vector(pred, dataset.action_stats) + if train_meta.action_mode == "delta": + pred_phys[:, :, :14] += meta["state_raw"][:, None, :14] + + gt_phys = meta["action_raw"][:, :, : dataset.physical_action_dim].to(torch.float32).cpu() + pred_chunks.append(pred_phys[0]) + gt_chunks.append(gt_phys[0]) + visited_indices.append(index) + + pred_tensor = torch.cat(pred_chunks, dim=0) + gt_tensor = torch.cat(gt_chunks, dim=0) + metrics = summarize_prediction_metrics(gt_tensor=gt_tensor, pred_tensor=pred_tensor) + + mse_values.append(metrics["mse"]) + mae_values.append(metrics["mae"]) + mse_joint_values.append(metrics["mse_joint"]) + mae_joint_values.append(metrics["mae_joint"]) + mse_gripper_values.append(metrics["mse_gripper"]) + mae_gripper_values.append(metrics["mae_gripper"]) + + if not skip_plots and plotting_available: + plot_prediction_series( + series={ + "Ground Truth": gt_tensor.numpy(), + "VLLM": pred_tensor.numpy(), + }, + output_path=plots_dir / f"{mode}_open_loop_ep{episode_index}.jpg", + title=f"{mode} Ground Truth vs VLLM", + ) + + episode_log = { + "episode_id": int(episode_index), + "task": task, + "visited_indices": visited_indices, + "num_pred_steps": int(pred_tensor.shape[0]), + **metrics, + } + per_episode.append(episode_log) + print(json.dumps({"mode": mode, **episode_log}, ensure_ascii=False)) + + summary = { + "mode": mode, + "num_episodes": len(per_episode), + "mse": mse_values, + "mae": mae_values, + "average_mse": summarize_metric_list(mse_values), + "average_mae": summarize_metric_list(mae_values), + "mse_joint": mse_joint_values, + "mae_joint": mae_joint_values, + "average_mse_joint": summarize_metric_list(mse_joint_values), + "average_mae_joint": summarize_metric_list(mae_joint_values), + "mse_gripper": mse_gripper_values, + "mae_gripper": mae_gripper_values, + "average_mse_gripper": summarize_metric_list(mse_gripper_values), + "average_mae_gripper": summarize_metric_list(mae_gripper_values), + "episodes": per_episode, + } + with open(output_dir / "log.json", "w", encoding="utf-8") as f: + json.dump(summary, f, indent=2, ensure_ascii=False) + return summary diff --git a/examples/offline_inference/internvla_a1/run.sh b/examples/offline_inference/internvla_a1/run.sh new file mode 100755 index 000000000..3ae9f776f --- /dev/null +++ b/examples/offline_inference/internvla_a1/run.sh @@ -0,0 +1,23 @@ +#!/usr/bin/env bash +set -euo pipefail + +ROOT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)" +REPO_ROOT="$(cd "$ROOT_DIR/../../.." && pwd)" + : "${INTERNVLA_A1_MODEL_DIR:?Please export INTERNVLA_A1_MODEL_DIR=/path/to/InternVLA-A1-3B-ft-pen}" + : "${INTERNVLA_A1_DATASET_DIR:?Please export INTERNVLA_A1_DATASET_DIR=/path/to/Genie1-Place_Markpen}" + : "${INTERNVLA_A1_PROCESSOR_DIR:?Please export INTERNVLA_A1_PROCESSOR_DIR=/path/to/Qwen3-VL-2B-Instruct}" + : "${INTERNVLA_A1_COSMOS_DIR:?Please export INTERNVLA_A1_COSMOS_DIR=/path/to/Cosmos-Tokenizer-CI8x8-SafeTensor}" +INTERNVLA_A1_OUTPUT_DIR="${INTERNVLA_A1_OUTPUT_DIR:-$REPO_ROOT/outputs/internvla_a1/vllm_infer}" + +export INTERNVLA_A1_MODEL_DIR +export INTERNVLA_A1_DATASET_DIR +export INTERNVLA_A1_PROCESSOR_DIR +export INTERNVLA_A1_COSMOS_DIR + +python "$ROOT_DIR/end2end.py" \ + --model-dir "$INTERNVLA_A1_MODEL_DIR" \ + --dataset-dir "$INTERNVLA_A1_DATASET_DIR" \ + --output-dir "$INTERNVLA_A1_OUTPUT_DIR" \ + --num-episodes "${INTERNVLA_A1_NUM_EPISODES:-1}" \ + --attn-implementation eager \ + "$@" diff --git a/tests/e2e/offline_inference/test_internvla_a1.py b/tests/e2e/offline_inference/test_internvla_a1.py new file mode 100644 index 000000000..38104dc66 --- /dev/null +++ b/tests/e2e/offline_inference/test_internvla_a1.py @@ -0,0 +1,90 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project + +"""Offline E2E test for InternVLA-A1 open-loop inference.""" + +from __future__ import annotations + +import json +import os +import subprocess +import sys +import tempfile +from pathlib import Path + +import pytest + +from tests.helpers.mark import hardware_test + +os.environ["VLLM_TEST_CLEAN_GPU_MEMORY"] = "1" + +REPO_ROOT = Path(__file__).resolve().parents[3] +EXAMPLE_SCRIPT = REPO_ROOT / "examples" / "offline_inference" / "internvla_a1" / "end2end.py" + + +def _required_env(name: str) -> str: + value = os.getenv(name) + if not value: + pytest.skip(f"{name} is not set") + return value + + +@pytest.mark.core_model +@pytest.mark.advanced_model +@pytest.mark.diffusion +@hardware_test(res={"cuda": "H100", "rocm": "MI325"}) +def test_internvla_a1_offline_open_loop(run_level: str) -> None: + if run_level != "advanced_model": + pytest.skip("InternVLA-A1 offline evaluation requires real local checkpoints and dataset.") + + model_dir = _required_env("INTERNVLA_A1_MODEL_DIR") + dataset_dir = _required_env("INTERNVLA_A1_DATASET_DIR") + processor_dir = _required_env("INTERNVLA_A1_PROCESSOR_DIR") + cosmos_dir = _required_env("INTERNVLA_A1_COSMOS_DIR") + + env = os.environ.copy() + env["INTERNVLA_A1_PROCESSOR_DIR"] = processor_dir + env["INTERNVLA_A1_COSMOS_DIR"] = cosmos_dir + + with tempfile.TemporaryDirectory(prefix="internvla_a1_e2e_") as tmpdir: + output_dir = Path(tmpdir) / "outputs" + cmd = [ + sys.executable, + str(EXAMPLE_SCRIPT), + "--model-dir", + model_dir, + "--dataset-dir", + dataset_dir, + "--output-dir", + str(output_dir), + "--num-samples", + "1", + "--num-episodes", + "1", + "--dtype", + "float32", + "--attn-implementation", + "eager", + "--skip-plots", + ] + subprocess.run(cmd, check=True, env=env) + + summary_path = output_dir / "summary.json" + assert summary_path.exists(), f"Missing summary file: {summary_path}" + + with open(summary_path, encoding="utf-8") as f: + summary = json.load(f) + + assert summary["mode"] == "registry" + assert summary["attn_implementation"] == "eager" + assert summary["dtype"] == "float32" + assert len(summary["results"]) == 1 + assert "registry" in summary["eval_summaries"] + + eval_summary = summary["eval_summaries"]["registry"] + assert eval_summary["num_episodes"] == 1 + assert eval_summary["average_mse"] is not None + assert eval_summary["average_mae"] is not None + + log_path = output_dir / "registry" / "log.json" + assert log_path.exists(), f"Missing evaluation log: {log_path}" diff --git a/vllm_omni/diffusion/models/internvla_a1/__init__.py b/vllm_omni/diffusion/models/internvla_a1/__init__.py new file mode 100644 index 000000000..39de3012b --- /dev/null +++ b/vllm_omni/diffusion/models/internvla_a1/__init__.py @@ -0,0 +1,25 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +"""InternVLA-A1 diffusion model components.""" + +from vllm_omni.diffusion.models.internvla_a1.config import ( + InternVLAA1Config, + InternVLAA1TrainMetadata, +) +from vllm_omni.diffusion.models.internvla_a1.model_internvla_a1 import ( + InternVLAA1, + InternVLAA1Policy, +) +from vllm_omni.diffusion.models.internvla_a1.pipeline_internvla_a1 import ( + InternVLAA1Pipeline, + get_internvla_a1_post_process_func, +) + +__all__ = [ + "InternVLAA1", + "InternVLAA1Config", + "InternVLAA1Pipeline", + "InternVLAA1TrainMetadata", + "InternVLAA1Policy", + "get_internvla_a1_post_process_func", +] diff --git a/vllm_omni/diffusion/models/internvla_a1/adapter_qwen3_vl.py b/vllm_omni/diffusion/models/internvla_a1/adapter_qwen3_vl.py new file mode 100644 index 000000000..a820fb56a --- /dev/null +++ b/vllm_omni/diffusion/models/internvla_a1/adapter_qwen3_vl.py @@ -0,0 +1,252 @@ +from __future__ import annotations + +import torch +import torch.nn as nn +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + ALL_ATTENTION_FUNCTIONS, + BaseModelOutputWithPast, + Cache, + DynamicCache, + FlashAttentionKwargs, + Qwen3VLConfig, + Qwen3VLPreTrainedModel, + Qwen3VLTextConfig, + Qwen3VLTextRotaryEmbedding, + Qwen3VLVisionModel, + Unpack, + apply_rotary_pos_emb, + check_model_inputs, + create_causal_mask, + deprecate_kwarg, + eager_attention_forward, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLForConditionalGeneration as HFQwen3VLForConditionalGeneration, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLModel as HFQwen3VLModel, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLTextAttention as HFQwen3VLTextAttention, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLTextDecoderLayer as HFQwen3VLTextDecoderLayer, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLTextModel as HFQwen3VLTextModel, +) +from transformers.models.qwen3_vl.modeling_qwen3_vl import ( + Qwen3VLTextRMSNorm as HFQwen3VLTextRMSNorm, +) + + +class Qwen3VLTextRMSNorm(HFQwen3VLTextRMSNorm): + def forward(self, hidden_states: torch.Tensor) -> torch.Tensor: + input_dtype = hidden_states.dtype + hidden_states = hidden_states.to(torch.float32) + variance = hidden_states.pow(2).mean(-1, keepdim=True) + hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon) + return (self.weight * hidden_states).to(input_dtype) + + +class Qwen3VLTextAttention(HFQwen3VLTextAttention): + def __init__(self, config: Qwen3VLTextConfig, layer_idx: int): + super().__init__(config, layer_idx) + self.q_norm = Qwen3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) + self.k_norm = Qwen3VLTextRMSNorm(self.head_dim, eps=config.rms_norm_eps) + + @deprecate_kwarg("past_key_value", new_name="past_key_values", version="4.58") + def forward( + self, + hidden_states: torch.Tensor, + position_embeddings: tuple[torch.Tensor, torch.Tensor], + attention_mask: torch.Tensor | None, + past_key_values: Cache | None = None, + cache_position: torch.LongTensor | None = None, + use_cache: bool = False, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple[torch.Tensor, torch.Tensor | None]: + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, self.head_dim) + + query_states = self.q_norm(self.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + key_states = self.k_norm(self.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + + cos, sin = position_embeddings + query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin) + + if past_key_values is not None: + if use_cache: + cache_kwargs = {"sin": sin, "cos": cos, "cache_position": cache_position} + key_states, value_states = past_key_values.update( + key_states, value_states, self.layer_idx, cache_kwargs + ) + else: + key_states = torch.cat([past_key_values[self.layer_idx][0], key_states], dim=2) + value_states = torch.cat([past_key_values[self.layer_idx][1], value_states], dim=2) + + attention_interface = eager_attention_forward + if self.config._attn_implementation != "eager": + attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation] + + attn_output, attn_weights = attention_interface( + self, + query_states, + key_states, + value_states, + attention_mask, + dropout=0.0 if not self.training else self.attention_dropout, + scaling=self.scaling, + **kwargs, + ) + + attn_output = attn_output.reshape(*input_shape, -1).contiguous() + attn_output = self.o_proj(attn_output) + return attn_output, attn_weights + + +class Qwen3VLTextDecoderLayer(HFQwen3VLTextDecoderLayer): + def __init__(self, config: Qwen3VLTextConfig, layer_idx: int): + super().__init__(config, layer_idx) + self.self_attn = Qwen3VLTextAttention(config=config, layer_idx=layer_idx) + self.input_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.post_attention_layernorm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + + +class Qwen3VLTextModel(HFQwen3VLTextModel): + config: Qwen3VLTextConfig + _no_split_modules = ["Qwen3VLTextDecoderLayer"] + _repeated_blocks = ["Qwen3VLTextDecoderLayer"] + + def __init__(self, config: Qwen3VLTextConfig): + Qwen3VLPreTrainedModel.__init__(self, config) + self.padding_idx = config.pad_token_id + self.vocab_size = config.vocab_size + + self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx) + self.layers = nn.ModuleList( + [Qwen3VLTextDecoderLayer(config, layer_idx) for layer_idx in range(config.num_hidden_layers)] + ) + self.norm = Qwen3VLTextRMSNorm(config.hidden_size, eps=config.rms_norm_eps) + self.rotary_emb = Qwen3VLTextRotaryEmbedding(config=config) + self.gradient_checkpointing = False + + self.post_init() + + @check_model_inputs + def forward( + self, + input_ids: torch.LongTensor = None, + attention_mask: torch.Tensor | None = None, + position_ids: torch.LongTensor | None = None, + past_key_values: Cache | None = None, + inputs_embeds: torch.FloatTensor | None = None, + use_cache: bool | None = None, + cache_position: torch.LongTensor | None = None, + visual_pos_masks: torch.Tensor | None = None, + deepstack_visual_embeds: list[torch.Tensor] | None = None, + **kwargs: Unpack[FlashAttentionKwargs], + ) -> tuple | BaseModelOutputWithPast: + if (input_ids is None) ^ (inputs_embeds is not None): + raise ValueError("You must specify exactly one of input_ids or inputs_embeds") + + if use_cache and past_key_values is None and not torch.jit.is_tracing(): + past_key_values = DynamicCache(config=self.config) + + if inputs_embeds is None: + inputs_embeds = self.embed_tokens(input_ids) + + if cache_position is None: + past_seen_tokens = past_key_values.get_seq_length() if past_key_values is not None else 0 + cache_position = torch.arange( + past_seen_tokens, past_seen_tokens + inputs_embeds.shape[1], device=inputs_embeds.device + ) + + if position_ids is None: + position_ids = cache_position.view(1, 1, -1).expand(3, inputs_embeds.shape[0], -1) + elif position_ids.ndim == 2: + position_ids = position_ids[None, ...].expand(3, position_ids.shape[0], -1) + + if position_ids.ndim == 3 and position_ids.shape[0] == 4: + text_position_ids = position_ids[0] + position_ids = position_ids[1:] + else: + text_position_ids = position_ids[0] + + attention_mask = create_causal_mask( + config=self.config, + input_embeds=inputs_embeds, + attention_mask=attention_mask, + cache_position=cache_position, + past_key_values=past_key_values, + position_ids=text_position_ids, + ) + + hidden_states = inputs_embeds + position_embeddings = self.rotary_emb(hidden_states, position_ids) + + for layer_idx, decoder_layer in enumerate(self.layers): + hidden_states = decoder_layer( + hidden_states, + attention_mask=attention_mask, + position_ids=text_position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + cache_position=cache_position, + position_embeddings=position_embeddings, + **kwargs, + ) + + if deepstack_visual_embeds is not None and layer_idx in range(len(deepstack_visual_embeds)): + hidden_states = self._deepstack_process( + hidden_states, + visual_pos_masks, + deepstack_visual_embeds[layer_idx], + ) + + hidden_states = self.norm(hidden_states) + + return BaseModelOutputWithPast( + last_hidden_state=hidden_states, + past_key_values=past_key_values, + ) + + +class Qwen3VLModel(HFQwen3VLModel): + base_model_prefix = "" + _checkpoint_conversion_mapping = {} + accepts_loss_kwargs = False + config: Qwen3VLConfig + _no_split_modules = ["Qwen3VLTextDecoderLayer", "Qwen3VLVisionBlock"] + + def __init__(self, config): + Qwen3VLPreTrainedModel.__init__(self, config) + self.visual = Qwen3VLVisionModel._from_config(config.vision_config) + self.language_model = Qwen3VLTextModel._from_config(config.text_config) + self.rope_deltas = None + self.post_init() + + +class Qwen3VLForConditionalGeneration(HFQwen3VLForConditionalGeneration): + _checkpoint_conversion_mapping = {} + _tied_weights_keys = ["lm_head.weight"] + accepts_loss_kwargs = False + config: Qwen3VLConfig + + def __init__(self, config): + Qwen3VLPreTrainedModel.__init__(self, config) + self.model = Qwen3VLModel(config) + self.lm_head = nn.Linear(config.text_config.hidden_size, config.text_config.vocab_size, bias=False) + self.post_init() + + +__all__ = [ + "Qwen3VLVisionModel", + "Qwen3VLForConditionalGeneration", + "Qwen3VLModel", + "Qwen3VLPreTrainedModel", + "Qwen3VLTextModel", + "apply_rotary_pos_emb", + "eager_attention_forward", +] diff --git a/vllm_omni/diffusion/models/internvla_a1/config.py b/vllm_omni/diffusion/models/internvla_a1/config.py new file mode 100644 index 000000000..b813b0391 --- /dev/null +++ b/vllm_omni/diffusion/models/internvla_a1/config.py @@ -0,0 +1,126 @@ +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, + ) diff --git a/vllm_omni/diffusion/models/internvla_a1/cosmos_ci_torch.py b/vllm_omni/diffusion/models/internvla_a1/cosmos_ci_torch.py new file mode 100644 index 000000000..a8237bfe6 --- /dev/null +++ b/vllm_omni/diffusion/models/internvla_a1/cosmos_ci_torch.py @@ -0,0 +1,416 @@ +# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES. All rights reserved. +# SPDX-License-Identifier: Apache-2.0 +# +# This file vendors the minimal Continuous Image (CI) tokenizer implementation +# needed by InternVLA-A1 from NVIDIA/Cosmos-Tokenizer. The code is intentionally +# reduced to the CI image path only so the runtime does not depend on the +# external `cosmos_tokenizer` package. + +from __future__ import annotations + +import math +from collections import OrderedDict + +import torch +import torch.nn as nn +import torch.nn.functional as F +from einops import rearrange + +_WAVELETS = { + "haar": torch.tensor([0.7071067811865476, 0.7071067811865476]), + "rearrange": torch.tensor([1.0, 1.0]), +} +_PERSISTENT = False + + +def nonlinearity(x: torch.Tensor) -> torch.Tensor: + return x * torch.sigmoid(x) + + +def normalize(in_channels: int, num_groups: int = 32) -> nn.GroupNorm: + return nn.GroupNorm(num_groups=num_groups, num_channels=in_channels, eps=1e-6, affine=True) + + +class IdentityDistribution(nn.Module): + def forward(self, parameters: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + zeros = torch.tensor([0.0], device=parameters.device, dtype=parameters.dtype) + return parameters, (zeros, zeros) + + +class Patcher(nn.Module): + def __init__(self, patch_size: int = 1, patch_method: str = "haar") -> None: + super().__init__() + self.patch_size = patch_size + self.patch_method = patch_method + self.register_buffer("wavelets", _WAVELETS[patch_method], persistent=_PERSISTENT) + self.range = range(int(torch.log2(torch.tensor(self.patch_size)).item())) + self.register_buffer( + "_arange", + torch.arange(_WAVELETS[patch_method].shape[0]), + persistent=_PERSISTENT, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.patch_method == "haar": + return self._haar(x) + if self.patch_method == "rearrange": + return self._arrange(x) + raise ValueError(f"Unknown patch method: {self.patch_method}") + + def _dwt(self, x: torch.Tensor, mode: str = "reflect", *, rescale: bool = False) -> torch.Tensor: + dtype = x.dtype + h = self.wavelets + n = h.shape[0] + g = x.shape[1] + hl = h.flip(0).reshape(1, 1, -1).repeat(g, 1, 1).to(dtype=dtype) + hh = (h * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1).to(dtype=dtype) + + x = F.pad(x, pad=(n - 2, n - 1, n - 2, n - 1), mode=mode).to(dtype) + xl = F.conv2d(x, hl.unsqueeze(2), groups=g, stride=(1, 2)) + xh = F.conv2d(x, hh.unsqueeze(2), groups=g, stride=(1, 2)) + xll = F.conv2d(xl, hl.unsqueeze(3), groups=g, stride=(2, 1)) + xlh = F.conv2d(xl, hh.unsqueeze(3), groups=g, stride=(2, 1)) + xhl = F.conv2d(xh, hl.unsqueeze(3), groups=g, stride=(2, 1)) + xhh = F.conv2d(xh, hh.unsqueeze(3), groups=g, stride=(2, 1)) + + out = torch.cat([xll, xlh, xhl, xhh], dim=1) + if rescale: + out = out / 2 + return out + + def _haar(self, x: torch.Tensor) -> torch.Tensor: + for _ in self.range: + x = self._dwt(x, rescale=True) + return x + + def _arrange(self, x: torch.Tensor) -> torch.Tensor: + return rearrange( + x, + "b c (h p1) (w p2) -> b (c p1 p2) h w", + p1=self.patch_size, + p2=self.patch_size, + ).contiguous() + + +class UnPatcher(nn.Module): + def __init__(self, patch_size: int = 1, patch_method: str = "haar") -> None: + super().__init__() + self.patch_size = patch_size + self.patch_method = patch_method + self.register_buffer("wavelets", _WAVELETS[patch_method], persistent=_PERSISTENT) + self.range = range(int(torch.log2(torch.tensor(self.patch_size)).item())) + self.register_buffer( + "_arange", + torch.arange(_WAVELETS[patch_method].shape[0]), + persistent=_PERSISTENT, + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + if self.patch_method == "haar": + return self._ihaar(x) + if self.patch_method == "rearrange": + return self._iarrange(x) + raise ValueError(f"Unknown patch method: {self.patch_method}") + + def _idwt(self, x: torch.Tensor, *, rescale: bool = False) -> torch.Tensor: + dtype = x.dtype + h = self.wavelets + n = h.shape[0] + g = x.shape[1] // 4 + hl = h.flip([0]).reshape(1, 1, -1).repeat([g, 1, 1]).to(dtype=dtype) + hh = (h * ((-1) ** self._arange)).reshape(1, 1, -1).repeat(g, 1, 1).to(dtype=dtype) + + xll, xlh, xhl, xhh = torch.chunk(x.to(dtype), 4, dim=1) + yl = F.conv_transpose2d(xll, hl.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0)) + yl += F.conv_transpose2d(xlh, hh.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0)) + yh = F.conv_transpose2d(xhl, hl.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0)) + yh += F.conv_transpose2d(xhh, hh.unsqueeze(3), groups=g, stride=(2, 1), padding=(n - 2, 0)) + y = F.conv_transpose2d(yl, hl.unsqueeze(2), groups=g, stride=(1, 2), padding=(0, n - 2)) + y += F.conv_transpose2d(yh, hh.unsqueeze(2), groups=g, stride=(1, 2), padding=(0, n - 2)) + + if rescale: + y = y * 2 + return y + + def _ihaar(self, x: torch.Tensor) -> torch.Tensor: + for _ in self.range: + x = self._idwt(x, rescale=True) + return x + + def _iarrange(self, x: torch.Tensor) -> torch.Tensor: + return rearrange( + x, + "b (c p1 p2) h w -> b c (h p1) (w p2)", + p1=self.patch_size, + p2=self.patch_size, + ) + + +class Upsample(nn.Module): + def __init__(self, in_channels: int) -> None: + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = x.repeat_interleave(2, dim=2).repeat_interleave(2, dim=3) + return self.conv(x) + + +class Downsample(nn.Module): + def __init__(self, in_channels: int) -> None: + super().__init__() + self.conv = nn.Conv2d(in_channels, in_channels, kernel_size=3, stride=2, padding=0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + return self.conv(F.pad(x, (0, 1, 0, 1), mode="constant", value=0)) + + +class ResnetBlock(nn.Module): + def __init__(self, *, in_channels: int, out_channels: int | None = None, dropout: float, **_: object) -> None: + super().__init__() + out_channels = in_channels if out_channels is None else out_channels + self.norm1 = normalize(in_channels) + self.conv1 = nn.Conv2d(in_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.norm2 = normalize(out_channels) + self.dropout = nn.Dropout(dropout) + self.conv2 = nn.Conv2d(out_channels, out_channels, kernel_size=3, stride=1, padding=1) + self.nin_shortcut = ( + nn.Conv2d(in_channels, out_channels, kernel_size=1, stride=1, padding=0) + if in_channels != out_channels + else nn.Identity() + ) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.conv1(nonlinearity(self.norm1(x))) + h = self.conv2(self.dropout(nonlinearity(self.norm2(h)))) + return self.nin_shortcut(x) + h + + +class AttnBlock(nn.Module): + def __init__(self, in_channels: int) -> None: + super().__init__() + self.norm = normalize(in_channels) + self.q = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.k = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.v = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + self.proj_out = nn.Conv2d(in_channels, in_channels, kernel_size=1, stride=1, padding=0) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + h = self.norm(x) + q = self.q(h) + k = self.k(h) + v = self.v(h) + b, c, height, width = q.shape + q = q.reshape(b, c, height * width).permute(0, 2, 1) + k = k.reshape(b, c, height * width) + attn = torch.bmm(q, k) * (int(c) ** -0.5) + attn = F.softmax(attn, dim=2) + v = v.reshape(b, c, height * width) + attn = attn.permute(0, 2, 1) + h = torch.bmm(v, attn).reshape(b, c, height, width) + return x + self.proj_out(h) + + +class Encoder(nn.Module): + def __init__( + self, + in_channels: int, + channels: int, + channels_mult: list[int], + num_res_blocks: int, + attn_resolutions: list[int], + dropout: float, + resolution: int, + z_channels: int, + spatial_compression: int, + **kwargs: object, + ) -> None: + super().__init__() + self.num_resolutions = len(channels_mult) + self.num_res_blocks = num_res_blocks + + patch_size = int(kwargs.get("patch_size", 1)) + self.patcher = Patcher(patch_size, str(kwargs.get("patch_method", "rearrange"))) + in_channels = in_channels * patch_size * patch_size + self.num_downsamples = int(math.log2(spatial_compression)) - int(math.log2(patch_size)) + assert self.num_downsamples <= self.num_resolutions + + self.conv_in = nn.Conv2d(in_channels, channels, kernel_size=3, stride=1, padding=1) + curr_res = resolution // patch_size + in_ch_mult = (1,) + tuple(channels_mult) + self.down = nn.ModuleList() + block_in = channels + for i_level in range(self.num_resolutions): + block = nn.ModuleList() + attn = nn.ModuleList() + block_in = channels * in_ch_mult[i_level] + block_out = channels * channels_mult[i_level] + for _ in range(self.num_res_blocks): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(AttnBlock(block_in)) + down = nn.Module() + down.block = block + down.attn = attn + if i_level < self.num_downsamples: + down.downsample = Downsample(block_in) + curr_res //= 2 + self.down.append(down) + + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout) + self.norm_out = normalize(block_in) + self.conv_out = nn.Conv2d(block_in, z_channels, kernel_size=3, stride=1, padding=1) + + def forward(self, x: torch.Tensor) -> torch.Tensor: + x = self.patcher(x) + hs = [self.conv_in(x)] + for i_level in range(self.num_resolutions): + for i_block in range(self.num_res_blocks): + h = self.down[i_level].block[i_block](hs[-1]) + if len(self.down[i_level].attn) > 0: + h = self.down[i_level].attn[i_block](h) + hs.append(h) + if i_level < self.num_downsamples: + hs.append(self.down[i_level].downsample(hs[-1])) + + h = hs[-1] + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + h = nonlinearity(self.norm_out(h)) + return self.conv_out(h) + + +class Decoder(nn.Module): + def __init__( + self, + out_channels: int, + channels: int, + channels_mult: list[int], + num_res_blocks: int, + attn_resolutions: list[int], + dropout: float, + resolution: int, + z_channels: int, + spatial_compression: int, + **kwargs: object, + ) -> None: + super().__init__() + self.num_resolutions = len(channels_mult) + self.num_res_blocks = num_res_blocks + + patch_size = int(kwargs.get("patch_size", 1)) + self.unpatcher = UnPatcher(patch_size, str(kwargs.get("patch_method", "rearrange"))) + out_ch = out_channels * patch_size * patch_size + self.num_upsamples = int(math.log2(spatial_compression)) - int(math.log2(patch_size)) + assert self.num_upsamples <= self.num_resolutions + + block_in = channels * channels_mult[self.num_resolutions - 1] + curr_res = (resolution // patch_size) // 2 ** (self.num_resolutions - 1) + self.z_shape = (1, z_channels, curr_res, curr_res) + self.conv_in = nn.Conv2d(z_channels, block_in, kernel_size=3, stride=1, padding=1) + + self.mid = nn.Module() + self.mid.block_1 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout) + self.mid.attn_1 = AttnBlock(block_in) + self.mid.block_2 = ResnetBlock(in_channels=block_in, out_channels=block_in, dropout=dropout) + + self.up = nn.ModuleList() + for i_level in reversed(range(self.num_resolutions)): + block = nn.ModuleList() + attn = nn.ModuleList() + block_out = channels * channels_mult[i_level] + for _ in range(self.num_res_blocks + 1): + block.append(ResnetBlock(in_channels=block_in, out_channels=block_out, dropout=dropout)) + block_in = block_out + if curr_res in attn_resolutions: + attn.append(AttnBlock(block_in)) + up = nn.Module() + up.block = block + up.attn = attn + if i_level >= (self.num_resolutions - self.num_upsamples): + up.upsample = Upsample(block_in) + curr_res *= 2 + self.up.insert(0, up) + + self.norm_out = normalize(block_in) + self.conv_out = nn.Conv2d(block_in, out_ch, kernel_size=3, stride=1, padding=1) + + def forward(self, z: torch.Tensor) -> torch.Tensor: + h = self.conv_in(z) + h = self.mid.block_1(h) + h = self.mid.attn_1(h) + h = self.mid.block_2(h) + for i_level in reversed(range(self.num_resolutions)): + for i_block in range(self.num_res_blocks + 1): + h = self.up[i_level].block[i_block](h) + if len(self.up[i_level].attn) > 0: + h = self.up[i_level].attn[i_block](h) + if i_level >= (self.num_resolutions - self.num_upsamples): + h = self.up[i_level].upsample(h) + h = self.conv_out(nonlinearity(self.norm_out(h))) + return self.unpatcher(h) + + +class ContinuousImageTokenizer(nn.Module): + def __init__(self, z_channels: int, z_factor: int, latent_channels: int, **kwargs: object) -> None: + super().__init__() + self.latent_channels = latent_channels + self.encoder = Encoder(z_channels=z_factor * z_channels, **kwargs) + self.decoder = Decoder(z_channels=z_channels, **kwargs) + self.quant_conv = nn.Conv2d(z_factor * z_channels, z_factor * latent_channels, 1) + self.post_quant_conv = nn.Conv2d(latent_channels, z_channels, 1) + self.distribution = IdentityDistribution() + + def encoder_module(self) -> nn.Sequential: + return nn.Sequential( + OrderedDict( + [ + ("encoder", self.encoder), + ("quant_conv", self.quant_conv), + ("distribution", self.distribution), + ] + ) + ) + + def decoder_module(self) -> nn.Sequential: + return nn.Sequential( + OrderedDict( + [ + ("post_quant_conv", self.post_quant_conv), + ("decoder", self.decoder), + ] + ) + ) + + def encode(self, x: torch.Tensor) -> tuple[torch.Tensor, tuple[torch.Tensor, torch.Tensor]]: + h = self.encoder(x) + moments = self.quant_conv(h) + return self.distribution(moments) + + def decode(self, z: torch.Tensor) -> torch.Tensor: + return self.decoder(self.post_quant_conv(z)) + + +def build_cosmos_ci_torch_model(spatial_compression: int) -> ContinuousImageTokenizer: + return ContinuousImageTokenizer( + attn_resolutions=[32], + channels=128, + channels_mult=[2, 4, 4], + dropout=0.0, + in_channels=3, + spatial_compression=spatial_compression, + num_res_blocks=2, + out_channels=3, + resolution=1024, + patch_size=4, + patch_method="haar", + latent_channels=16, + z_channels=16, + z_factor=1, + name="CI", + ) diff --git a/vllm_omni/diffusion/models/internvla_a1/model_cosmos.py b/vllm_omni/diffusion/models/internvla_a1/model_cosmos.py new file mode 100644 index 000000000..8439c7569 --- /dev/null +++ b/vllm_omni/diffusion/models/internvla_a1/model_cosmos.py @@ -0,0 +1,97 @@ +from __future__ import annotations + +from pathlib import Path + +import torch +from safetensors.torch import load_file + +from .cosmos_ci_torch import build_cosmos_ci_torch_model + +_SAFETENSORS_EXTENSION = ".safetensors" + + +def is_safetensors_checkpoint(path: str | Path) -> bool: + return Path(path).suffix.lower() == _SAFETENSORS_EXTENSION + + +def infer_cosmos_ci_spatial_compression(path: str | Path) -> int: + normalized = str(path).upper() + if "CI8X8" in normalized or "8X8" in normalized: + return 8 + if "CI16X16" in normalized or "16X16" in normalized: + return 16 + raise ValueError(f"Unable to infer Cosmos CI spatial compression from checkpoint path: {path}") + + +def _load_safetensors_checkpoint(path: str | Path, device: str) -> dict[str, torch.Tensor]: + return load_file(str(path), device=device) + + +def _load_native_cosmos_component( + checkpoint_filepath: str | Path, + *, + component: str, + device: str, +) -> torch.nn.Module: + full_model = build_cosmos_ci_torch_model( + spatial_compression=infer_cosmos_ci_spatial_compression(checkpoint_filepath) + ) + if component == "encoder": + model = full_model.encoder_module() + elif component == "decoder": + model = full_model.decoder_module() + else: + raise ValueError(f"Unsupported Cosmos component: {component}") + + if not is_safetensors_checkpoint(checkpoint_filepath): + raise ValueError(f"Native Cosmos torch backend expects `.safetensors` checkpoints. Got: {checkpoint_filepath}") + state_dict = _load_safetensors_checkpoint(checkpoint_filepath, device) + model.load_state_dict(state_dict, strict=False) + return model.eval().to(device) + + +def load_cosmos_component( + checkpoint_filepath: str | Path, + *, + component: str, + device: str = "cuda", +) -> torch.nn.Module: + return _load_native_cosmos_component( + checkpoint_filepath, + component=component, + device=device, + ) + + +class ImageTokenizer(torch.nn.Module): + def __init__( + self, + checkpoint_enc: str, + checkpoint_dec: str, + device: str = "cuda", + ) -> None: + super().__init__() + self._enc_model = load_cosmos_component( + checkpoint_enc, + component="encoder", + device=device, + ) + self._dec_model = load_cosmos_component( + checkpoint_dec, + component="decoder", + device=device, + ) + + def decode(self, input_latent: torch.Tensor) -> torch.Tensor: + original_dtype = input_latent.dtype + model_dtype = next(self.parameters()).dtype + output_tensor = self._dec_model(input_latent.to(model_dtype)) + return output_tensor.to(original_dtype) + + def encode(self, input_tensor: torch.Tensor) -> torch.Tensor: + original_dtype = input_tensor.dtype + model_dtype = next(self.parameters()).dtype + output_latent = self._enc_model(input_tensor.to(model_dtype)) + if isinstance(output_latent, torch.Tensor): + return output_latent.to(original_dtype) + return output_latent[0].to(original_dtype) diff --git a/vllm_omni/diffusion/models/internvla_a1/model_internvla_a1.py b/vllm_omni/diffusion/models/internvla_a1/model_internvla_a1.py new file mode 100644 index 000000000..7cd064366 --- /dev/null +++ b/vllm_omni/diffusion/models/internvla_a1/model_internvla_a1.py @@ -0,0 +1,962 @@ +from __future__ import annotations + +import math +import os +from dataclasses import dataclass +from pathlib import Path +from typing import Any, Literal + +import torch +import torch._dynamo as dynamo +import torch.nn.functional as F +from einops import rearrange +from safetensors import safe_open +from torch import Tensor, nn +from transformers import Qwen3VLProcessor +from transformers.models.auto import CONFIG_MAPPING +from vllm.model_executor.models.utils import AutoWeightsLoader + +from .adapter_qwen3_vl import ( + Qwen3VLForConditionalGeneration, + Qwen3VLTextModel, + apply_rotary_pos_emb, + eager_attention_forward, +) +from .config import ( + DEFAULT_COSMOS_DIR, + DEFAULT_COSMOS_REPO, + DEFAULT_QWEN3_VL_MODEL, + OBS_IMAGES, + OBS_PREFIX, + OBS_STATE, + OBS_TASK, + OPENPI_ATTENTION_MASK_VALUE, + InternVLAA1Config, +) +from .model_cosmos import ImageTokenizer + + +def get_safe_dtype(target_dtype: torch.dtype, device_type: str) -> torch.dtype: + if device_type == "mps" and target_dtype == torch.float64: + return torch.float32 + if device_type == "cpu" and target_dtype == torch.bfloat16: + return torch.float32 + return target_dtype + + +def create_sinusoidal_pos_embedding( + time: torch.Tensor, + dimension: int, + min_period: float, + max_period: float, + device: torch.device, +) -> Tensor: + if dimension % 2 != 0: + raise ValueError(f"dimension ({dimension}) must be divisible by 2") + dtype = get_safe_dtype(torch.float64, device.type) + fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device) + period = min_period * (max_period / min_period) ** fraction + scaling_factor = 1.0 / period * 2 * math.pi + sin_input = scaling_factor[None, :] * time[:, None] + return torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1) + + +def make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor: + cumsum = torch.cumsum(att_masks, dim=1) + att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None] + pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None] + return att_2d_masks & pad_2d_masks + + +def pad_vector(vector: torch.Tensor, new_dim: int) -> torch.Tensor: + if vector.shape[-1] >= new_dim: + return vector + return F.pad(vector, (0, new_dim - vector.shape[-1])) + + +def resize_with_pad(images: torch.Tensor, size: tuple[int, int]) -> torch.Tensor: + target_h, target_w = size + if images.ndim != 4: + raise ValueError(f"Expected [T, C, H, W], got {tuple(images.shape)}") + _, _, src_h, src_w = images.shape + scale = min(target_h / src_h, target_w / src_w) + resized_h = max(1, int(round(src_h * scale))) + resized_w = max(1, int(round(src_w * scale))) + resized = F.interpolate(images, size=(resized_h, resized_w), mode="bilinear", align_corners=False) + pad_h = target_h - resized_h + pad_w = target_w - resized_w + pad_top = pad_h // 2 + pad_bottom = pad_h - pad_top + pad_left = pad_w // 2 + pad_right = pad_w - pad_left + return F.pad(resized, (pad_left, pad_right, pad_top, pad_bottom)) + + +class Qwen3VLInputBuilder: + def __init__( + self, + processor_model_name: str = DEFAULT_QWEN3_VL_MODEL, + max_length: int = 48, + spatial_merge_size: int = 2, + ) -> None: + processor_model_name = os.getenv("INTERNVLA_A1_PROCESSOR_DIR", processor_model_name) + local_path = Path(processor_model_name).expanduser() + if local_path.exists(): + self.processor = Qwen3VLProcessor.from_pretrained(str(local_path), local_files_only=True) + else: + if os.getenv("INTERNVLA_A1_PROCESSOR_DIR"): + raise FileNotFoundError(f"INTERNVLA_A1_PROCESSOR_DIR points to a missing path: {local_path}") + self.processor = Qwen3VLProcessor.from_pretrained(processor_model_name) + self.max_length = max_length + self.spatial_merge_size = spatial_merge_size + + def build(self, camera_images: list[torch.Tensor], task: str) -> dict[str, torch.Tensor]: + input_ids: list[int] = [] + attention_mask: list[int] = [] + pixel_values: list[torch.Tensor] = [] + image_grid_thw: list[torch.Tensor] = [] + + for image_history in camera_images: + current_image = image_history[1] + image_inputs = self.processor.image_processor(current_image, do_rescale=False) + pixel_values.append(image_inputs.pixel_values) + image_grid_thw.append(image_inputs.image_grid_thw) + num_img_token = int(torch.prod(image_inputs.image_grid_thw[0]).item()) // (self.spatial_merge_size**2) + input_ids.extend( + [self.processor.vision_start_token_id] + + [self.processor.image_token_id] * num_img_token + + [self.processor.vision_end_token_id] + ) + attention_mask.extend([1] * (num_img_token + 2)) + + text_inputs = self.processor.tokenizer( + task, + max_length=self.max_length, + padding="max_length", + padding_side="right", + truncation=True, + ) + input_ids.extend(text_inputs.input_ids) + attention_mask.extend(text_inputs.attention_mask) + + return { + f"{OBS_PREFIX}pixel_values": torch.cat(pixel_values, dim=0), + f"{OBS_PREFIX}image_grid_thw": torch.cat(image_grid_thw, dim=0), + f"{OBS_PREFIX}input_ids": torch.tensor(input_ids, dtype=torch.long), + f"{OBS_PREFIX}attention_mask": torch.tensor(attention_mask, dtype=torch.long), + } + + +@dataclass +class SuffixStaticContext: + state_emb: torch.Tensor + full_att_2d_masks_4d: torch.Tensor + position_ids: torch.Tensor + + +def compute_layer_complete( + layer_idx: int, + inputs_embeds: list[torch.Tensor], + attention_mask: torch.Tensor, + position_ids: torch.Tensor, + und_expert: nn.Module, + gen_expert: nn.Module, + act_expert: nn.Module, +) -> list[torch.Tensor]: + models = [und_expert.language_model, gen_expert, act_expert] + query_states = [] + key_states = [] + value_states = [] + for i, hidden_states in enumerate(inputs_embeds): + layer = models[i].layers[layer_idx] + hidden_states = layer.input_layernorm(hidden_states) + input_shape = hidden_states.shape[:-1] + hidden_shape = (*input_shape, -1, layer.self_attn.head_dim) + if layer.self_attn.q_proj.weight.dtype == torch.bfloat16: + hidden_states = hidden_states.to(dtype=torch.bfloat16) + query_state = layer.self_attn.q_norm(layer.self_attn.q_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + key_state = layer.self_attn.k_norm(layer.self_attn.k_proj(hidden_states).view(hidden_shape)).transpose(1, 2) + value_state = layer.self_attn.v_proj(hidden_states).view(hidden_shape).transpose(1, 2) + query_states.append(query_state) + key_states.append(key_state) + value_states.append(value_state) + + query_states = torch.cat(query_states, dim=2) + key_states = torch.cat(key_states, dim=2) + value_states = torch.cat(value_states, dim=2) + dummy_tensor = torch.zeros( + query_states.shape[0], + query_states.shape[2], + query_states.shape[-1], + device=query_states.device, + dtype=query_states.dtype, + ) + cos, sin = und_expert.model.language_model.rotary_emb(dummy_tensor, position_ids) + query_states, key_states = apply_rotary_pos_emb( + query_states, + key_states, + cos, + sin, + unsqueeze_dim=1, + ) + scaling = und_expert.language_model.layers[layer_idx].self_attn.scaling + att_output, _ = eager_attention_forward( + und_expert.language_model.layers[layer_idx].self_attn, + query_states, + key_states, + value_states, + attention_mask, + scaling, + ) + + head_dim = und_expert.language_model.layers[layer_idx].self_attn.head_dim + num_attention_heads = und_expert.language_model.layers[layer_idx].self_attn.config.num_attention_heads + batch_size = query_states.shape[0] + att_output = att_output.reshape(batch_size, -1, num_attention_heads * head_dim) + + outputs_embeds = [] + start_pos = 0 + for i, hidden_states in enumerate(inputs_embeds): + layer = models[i].layers[layer_idx] + end_pos = start_pos + hidden_states.shape[1] + att_chunk = att_output[:, start_pos:end_pos] + if att_chunk.dtype != layer.self_attn.o_proj.weight.dtype: + att_chunk = att_chunk.to(layer.self_attn.o_proj.weight.dtype) + out_emb = layer.self_attn.o_proj(att_chunk) + out_emb = out_emb + hidden_states + residual = out_emb + out_emb = layer.post_attention_layernorm(out_emb) + if layer.mlp.up_proj.weight.dtype == torch.bfloat16: + out_emb = out_emb.to(dtype=torch.bfloat16) + out_emb = layer.mlp(out_emb) + out_emb = out_emb + residual + outputs_embeds.append(out_emb) + start_pos = end_pos + return outputs_embeds + + +class QwenConfig: + def __init__( + self, + head_dim: int, + hidden_size: int, + intermediate_size: int, + num_attention_heads: int, + num_hidden_layers: int, + num_key_value_heads: int, + ) -> None: + self.head_dim = head_dim + self.hidden_size = hidden_size + self.intermediate_size = intermediate_size + self.num_attention_heads = num_attention_heads + self.num_hidden_layers = num_hidden_layers + self.num_key_value_heads = num_key_value_heads + + +def get_qwen_config(variant: str) -> QwenConfig: + num_hidden_layers = int(variant.split("_")[-1][:-1]) + if variant.startswith("qwen3_vl"): + return QwenConfig(128, 2048, 6144, 16, num_hidden_layers, 8) + if variant.startswith("qwen3"): + return QwenConfig(128, 1024, 3072, 16, num_hidden_layers, 8) + raise ValueError(f"Unknown variant: {variant}") + + +def resolve_cosmos_checkpoint_paths( + *, + encoder_path: str | Path | None = None, + decoder_path: str | Path | None = None, +) -> tuple[Path, Path]: + encoder_override = encoder_path or os.getenv("INTERNVLA_A1_COSMOS_ENCODER_PATH") + decoder_override = decoder_path or os.getenv("INTERNVLA_A1_COSMOS_DECODER_PATH") + default_encoder_name = "encoder.safetensors" + default_decoder_name = "decoder.safetensors" + checkpoint_enc = ( + Path(encoder_override).expanduser() + if encoder_override is not None + else DEFAULT_COSMOS_DIR / default_encoder_name + ) + checkpoint_dec = ( + Path(decoder_override).expanduser() + if decoder_override is not None + else DEFAULT_COSMOS_DIR / default_decoder_name + ) + + missing = [str(path) for path in (checkpoint_enc, checkpoint_dec) if not path.exists()] + if missing: + missing_display = ", ".join(missing) + raise FileNotFoundError( + "InternVLA-A1 requires local Cosmos tokenizer checkpoints. " + f"Missing: {missing_display}. " + f"Download {DEFAULT_COSMOS_REPO} and set INTERNVLA_A1_COSMOS_DIR to that directory, " + "or point INTERNVLA_A1_COSMOS_ENCODER_PATH / INTERNVLA_A1_COSMOS_DECODER_PATH " + "to explicit checkpoint files." + ) + + return checkpoint_enc, checkpoint_dec + + +class Qwen3VLWithExpertModel(nn.Module): + def __init__( + self, + vlm_config: QwenConfig, + action_expert_config: QwenConfig, + precision: Literal["bfloat16", "float32"] = "bfloat16", + ) -> None: + super().__init__() + + vlm_config_hf = CONFIG_MAPPING["qwen3_vl"]() + vlm_config_hf.text_config.hidden_size = vlm_config.hidden_size + vlm_config_hf.text_config.intermediate_size = vlm_config.intermediate_size + vlm_config_hf.text_config.num_attention_heads = vlm_config.num_attention_heads + vlm_config_hf.text_config.head_dim = vlm_config.head_dim + vlm_config_hf.text_config.num_hidden_layers = vlm_config.num_hidden_layers + vlm_config_hf.text_config.num_key_value_heads = vlm_config.num_key_value_heads + vlm_config_hf.text_config.max_position_embeddings = 262144 + vlm_config_hf.text_config.rope_scaling = { + "mrope_interleaved": True, + "mrope_section": [24, 20, 20], + "rope_type": "default", + } + vlm_config_hf.text_config.tie_word_embeddings = True + vlm_config_hf.tie_word_embeddings = True + vlm_config_hf.vision_config.deepstack_visual_indexes = [5, 11, 17] + vlm_config_hf.vision_config.depth = 24 + vlm_config_hf.vision_config.hidden_size = 1024 + vlm_config_hf.vision_config.intermediate_size = 4096 + vlm_config_hf.vision_config.out_hidden_size = 2048 + + self.und_expert = Qwen3VLForConditionalGeneration(config=vlm_config_hf) + + gen_expert_config_hf = CONFIG_MAPPING["qwen3_vl_text"]() + gen_expert_config_hf.head_dim = action_expert_config.head_dim + gen_expert_config_hf.hidden_size = action_expert_config.hidden_size + gen_expert_config_hf.intermediate_size = action_expert_config.intermediate_size + gen_expert_config_hf.num_attention_heads = action_expert_config.num_attention_heads + gen_expert_config_hf.num_hidden_layers = action_expert_config.num_hidden_layers + gen_expert_config_hf.num_key_value_heads = action_expert_config.num_key_value_heads + gen_expert_config_hf.max_position_embeddings = self.und_expert.config.text_config.max_position_embeddings + gen_expert_config_hf.rope_scaling = self.und_expert.config.text_config.rope_scaling + self.gen_expert = Qwen3VLTextModel(config=gen_expert_config_hf) + self.gen_expert.embed_tokens = None + self.gen_expert.lm_head = None + + act_expert_config_hf = CONFIG_MAPPING["qwen3_vl_text"]() + act_expert_config_hf.head_dim = action_expert_config.head_dim + act_expert_config_hf.hidden_size = action_expert_config.hidden_size + act_expert_config_hf.intermediate_size = action_expert_config.intermediate_size + act_expert_config_hf.num_attention_heads = action_expert_config.num_attention_heads + act_expert_config_hf.num_hidden_layers = action_expert_config.num_hidden_layers + act_expert_config_hf.num_key_value_heads = action_expert_config.num_key_value_heads + act_expert_config_hf.max_position_embeddings = self.und_expert.config.text_config.max_position_embeddings + act_expert_config_hf.rope_scaling = self.und_expert.config.text_config.rope_scaling + self.act_expert = Qwen3VLTextModel(config=act_expert_config_hf) + self.act_expert.embed_tokens = None + self.act_expert.lm_head = None + + self.to_selected_precision(precision) + + def to_selected_precision(self, precision: Literal["bfloat16", "float32"]) -> None: + if precision == "bfloat16": + self.to(dtype=torch.bfloat16) + elif precision == "float32": + # The fp32 path keeps the whole module in full precision, so no + # selective post-conversion fixups are required. + self.to(dtype=torch.float32) + return + else: + raise ValueError(f"Unsupported precision: {precision}") + + # Match the original runtime numerics by keeping normalization layers + # in fp32 even when the rest of the module runs in bf16. + keep_fp32 = ["input_layernorm", "post_attention_layernorm", "model.norm"] + for name, param in self.named_parameters(): + if any(key in name for key in keep_fp32): + param.data = param.data.to(dtype=torch.float32) + + def forward( + self, + attention_mask: torch.Tensor | None, + position_ids: torch.LongTensor | None, + past_key_values: Any, + inputs_embeds: list[torch.Tensor | None], + use_cache: bool, + ) -> tuple[list[torch.Tensor | None], Any]: + if inputs_embeds[1] is None and inputs_embeds[2] is None: + prefix_output = self.und_expert.language_model.forward( + inputs_embeds=inputs_embeds[0], + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + ) + return [prefix_output.last_hidden_state, None, None], prefix_output.past_key_values + + if inputs_embeds[0] is None and inputs_embeds[2] is None: + middle_output = self.gen_expert.forward( + inputs_embeds=inputs_embeds[1], + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + ) + return [None, middle_output.last_hidden_state, None], middle_output.past_key_values + + if inputs_embeds[0] is None and inputs_embeds[1] is None: + suffix_output = self.act_expert.forward( + inputs_embeds=inputs_embeds[2], + attention_mask=attention_mask, + position_ids=position_ids, + past_key_values=past_key_values, + use_cache=use_cache, + ) + return [None, None, suffix_output.last_hidden_state], None + + models = [self.und_expert.language_model, self.gen_expert, self.act_expert] + stacked_inputs = [inputs_embeds[0], inputs_embeds[1], inputs_embeds[2]] + for layer_idx in range(self.und_expert.config.text_config.num_hidden_layers): + stacked_inputs = compute_layer_complete( + layer_idx, + stacked_inputs, + attention_mask, + position_ids, + und_expert=self.und_expert, + gen_expert=self.gen_expert, + act_expert=self.act_expert, + ) + outputs = [models[i].norm(stacked_inputs[i]) for i in range(3)] + return outputs, None + + +class InternVLAA1(nn.Module): + def __init__( + self, + config: InternVLAA1Config, + *, + cosmos_encoder_path: str | Path | None = None, + cosmos_decoder_path: str | Path | None = None, + ) -> None: + super().__init__() + self.config = config + + vlm_config = get_qwen_config(config.qwen3_vl_variant) + action_expert_config = get_qwen_config(config.action_expert_variant) + self.qwen3_vl_with_expert = Qwen3VLWithExpertModel(vlm_config, action_expert_config, precision=config.dtype) + + cosmos_encoder_path, cosmos_decoder_path = resolve_cosmos_checkpoint_paths( + encoder_path=cosmos_encoder_path, + decoder_path=cosmos_decoder_path, + ) + self.cosmos = ImageTokenizer( + checkpoint_enc=str(cosmos_encoder_path), + checkpoint_dec=str(cosmos_decoder_path), + device=config.device, + ) + + hidden_size = action_expert_config.hidden_size + vae_dim = 16 + ds = config.scale_factor + self.cosmos_in_proj = nn.Conv2d(vae_dim, hidden_size, kernel_size=1, stride=1, padding=0) + self.downsample_conv = nn.Conv2d(hidden_size, hidden_size, kernel_size=ds, stride=ds, padding=0) + self.upsample_conv = nn.ConvTranspose2d(hidden_size, hidden_size, kernel_size=ds, stride=ds, padding=0) + self.cosmos_out_proj = nn.Linear(hidden_size, vae_dim) + self.cosmos_out_layer_norm = nn.LayerNorm(hidden_size) + + self.action_in_proj = nn.Linear(config.max_action_dim, hidden_size) + self.action_out_proj = nn.Linear(hidden_size, config.max_action_dim) + self.state_proj = nn.Linear(config.max_state_dim, hidden_size) + self.action_time_mlp_in = nn.Linear(2 * hidden_size, hidden_size) + self.action_time_mlp_out = nn.Linear(hidden_size, hidden_size) + + if config.compile_model: + torch.set_float32_matmul_precision("high") + self.sample_actions = torch.compile(self.sample_actions, mode=config.compile_mode) + + self.cosmos.eval() + for param in self.cosmos.parameters(): + param.requires_grad = False + + self.set_attention_implementation(config.attn_implementation) + + def set_attention_implementation(self, attn_implementation: str) -> None: + self.config.attn_implementation = attn_implementation + self.qwen3_vl_with_expert.und_expert.config.text_config._attn_implementation = attn_implementation + self.qwen3_vl_with_expert.und_expert.language_model.config._attn_implementation = attn_implementation + self.qwen3_vl_with_expert.gen_expert.config._attn_implementation = attn_implementation + self.qwen3_vl_with_expert.act_expert.config._attn_implementation = attn_implementation + + def _prepare_attention_masks_4d( + self, + att_2d_masks: torch.Tensor, + *, + dtype: torch.dtype, + ) -> torch.Tensor: + att_2d_masks_4d = att_2d_masks[:, None, :, :] + bias_dtype = torch.float32 if self.config.attn_implementation == "eager" else dtype + valid_bias = torch.zeros((), dtype=bias_dtype, device=att_2d_masks.device) + invalid_bias = torch.tensor( + OPENPI_ATTENTION_MASK_VALUE, + dtype=bias_dtype, + device=att_2d_masks.device, + ) + return torch.where(att_2d_masks_4d, valid_bias, invalid_bias) + + def sample_noise(self, shape: tuple[int, ...], device: torch.device) -> torch.Tensor: + return torch.normal(0.0, 1.0, shape, dtype=torch.float32, device=device) + + @dynamo.disable + def embed_prefix( + self, + pixel_values: torch.Tensor, + image_grid_thw: torch.Tensor, + lang_tokens: torch.Tensor, + lang_masks: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + image_token_id = self.qwen3_vl_with_expert.und_expert.config.image_token_id + pixel_values = pixel_values.view(-1, pixel_values.shape[-1]) + image_grid_thw = image_grid_thw.view(-1, 3) + image_embs, _ = self.qwen3_vl_with_expert.und_expert.visual(pixel_values, image_grid_thw) + + embs = self.qwen3_vl_with_expert.und_expert.get_input_embeddings()(lang_tokens) + batch_size, seq_len, hidden_dim = embs.shape + embs = embs.view(-1, hidden_dim) + lang_tokens_flat = lang_tokens.view(-1) + embs[lang_tokens_flat == image_token_id] = image_embs + embs = embs.view(batch_size, seq_len, hidden_dim) + + pad_masks = lang_masks.to(torch.bool) + att_masks = torch.zeros_like(pad_masks, dtype=torch.bool, device=pad_masks.device) + return embs, pad_masks, att_masks + + def get_cosmos_features(self, images: torch.Tensor) -> torch.Tensor: + shape = images.shape[:-3] + channels, height, width = images.shape[-3:] + images = images.reshape(-1, channels, height, width) + if (height, width) != (256, 256): + images = F.interpolate(images, size=(256, 256), mode="bilinear", align_corners=False) + images = images * 2 - 1 + features = self.cosmos.encode(images) + channels, height, width = features.shape[-3:] + return features.view(*shape, channels, height, width) + + def embed_middle( + self, + images: torch.Tensor, + img_masks: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + device = images.device + batch_size, n_view, timesteps = images.shape[:3] + features = self.get_cosmos_features(images) + features = rearrange(features, "b n t c h w -> (b n t) c h w") + features = self.cosmos_in_proj(features) + features = self.downsample_conv(features) + features = rearrange(features, "(b n t) c h w -> b n t c h w", b=batch_size, n=n_view, t=timesteps) + + _, _, _, _, height, width = features.shape + embs = rearrange(features, "b n t c h w -> b (n t h w) c") + pad_masks = torch.zeros((batch_size, n_view, timesteps, height, width), dtype=torch.bool, device=device) + pad_masks[img_masks] = True + pad_masks = rearrange(pad_masks, "b n t h w -> b (n t h w)") + att_masks = torch.tensor([1] + [0] * (embs.shape[1] - 1), dtype=torch.bool, device=device) + att_masks = att_masks[None, :].expand(batch_size, att_masks.shape[0]) + return embs, pad_masks, att_masks + + def prepare_suffix_static_context( + self, + state: torch.Tensor, + prefix_pad_masks: torch.Tensor, + max_prefix_position_ids: torch.Tensor, + ) -> SuffixStaticContext: + if self.state_proj.weight.dtype == torch.float32: + state = state.to(torch.float32) + + state_emb = self.state_proj(state) + device = state_emb.device + batch_size = state_emb.shape[0] + suffix_len = self.config.chunk_size + 1 + prefix_len = prefix_pad_masks.shape[1] + + suffix_pad_masks = torch.ones(batch_size, suffix_len, dtype=torch.bool, device=device) + suffix_att_masks = torch.tensor([1] + [1] + [0] * (self.config.chunk_size - 1), dtype=torch.bool, device=device) + suffix_att_masks = suffix_att_masks[None, :].expand(batch_size, suffix_len) + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len) + suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) + full_att_2d_masks_4d = self._prepare_attention_masks_4d( + torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2), + dtype=state_emb.dtype, + ) + position_ids = ( + torch.arange(1, suffix_len + 1, device=max_prefix_position_ids.device) + .repeat(3, 1, 1) + .to(max_prefix_position_ids) + + max_prefix_position_ids + ) + return SuffixStaticContext( + state_emb=state_emb, + full_att_2d_masks_4d=full_att_2d_masks_4d, + position_ids=position_ids, + ) + + def embed_suffix( + self, + state: torch.Tensor, + noisy_actions: torch.Tensor, + timestep: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: + if self.state_proj.weight.dtype == torch.float32: + state = state.to(torch.float32) + + state_emb = self.state_proj(state) + batch_size = state_emb.shape[0] + device = state_emb.device + state_mask = torch.ones(batch_size, 1, dtype=torch.bool, device=device) + + time_emb = create_sinusoidal_pos_embedding( + timestep, + self.action_in_proj.out_features, + min_period=self.config.min_period, + max_period=self.config.max_period, + device=timestep.device, + ).to(dtype=timestep.dtype) + action_emb = self.action_in_proj(noisy_actions) + time_emb = time_emb[:, None, :].expand_as(action_emb) + action_time_emb = torch.cat([action_emb, time_emb], dim=2) + action_time_emb = self.action_time_mlp_out(F.silu(self.action_time_mlp_in(action_time_emb))) + action_mask = torch.ones(batch_size, action_time_emb.shape[1], dtype=torch.bool, device=device) + + embs = torch.cat([state_emb[:, None, :], action_time_emb], dim=1) + pad_masks = torch.cat([state_mask, action_mask], dim=1) + att_masks = torch.tensor([1] + [1] + [0] * (self.config.chunk_size - 1), dtype=embs.dtype, device=device) + att_masks = att_masks[None, :].expand(batch_size, self.config.chunk_size + 1) + return embs, pad_masks, att_masks + + def get_position_ids( + self, + lang_tokens: torch.Tensor, + image_grid_thw: torch.Tensor | None, + pad_masks: torch.Tensor, + ) -> tuple[torch.Tensor, Any]: + seq_len = lang_tokens.shape[1] + padded_lang_tokens = torch.ones_like(pad_masks).to(lang_tokens) * 777 + padded_lang_tokens[:, :seq_len] = lang_tokens + attention_mask = pad_masks.to(lang_tokens) + if image_grid_thw is not None: + image_grid_thw = image_grid_thw.view(-1, 3) + return self.qwen3_vl_with_expert.und_expert.model.get_rope_index( + padded_lang_tokens, + image_grid_thw, + attention_mask=attention_mask, + ) + + @torch.no_grad() + def sample_actions( + self, + images: torch.Tensor, + img_masks: torch.Tensor, + pixel_values: torch.Tensor, + image_grid_thw: torch.Tensor, + lang_tokens: torch.Tensor, + lang_masks: torch.Tensor, + state: torch.Tensor, + *, + noise: torch.Tensor | None = None, + num_steps: int | None = None, + decode_image: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + if num_steps is None: + num_steps = self.config.num_inference_steps + batch_size = state.shape[0] + device = state.device + dtype = state.dtype + if noise is None: + noise = self.sample_noise((batch_size, self.config.chunk_size, self.config.max_action_dim), device) + + prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix( + pixel_values, + image_grid_thw, + lang_tokens, + lang_masks, + ) + prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks) + prefix_position_ids, _ = self.get_position_ids(lang_tokens, image_grid_thw, prefix_pad_masks) + prefix_att_2d_masks_4d = self._prepare_attention_masks_4d( + prefix_att_2d_masks, + dtype=prefix_embs.dtype, + ) + _, past_key_values = self.qwen3_vl_with_expert.forward( + attention_mask=prefix_att_2d_masks_4d, + position_ids=prefix_position_ids, + past_key_values=None, + inputs_embeds=[prefix_embs, None, None], + use_cache=True, + ) + max_prefix_position_ids = prefix_position_ids.max(dim=-1, keepdim=True).values + + middle_embs, middle_pad_masks, middle_att_masks = self.embed_middle(images[:, :, :2], img_masks) + middle_len = middle_pad_masks.shape[1] + prefix_len = prefix_pad_masks.shape[1] + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, middle_len, prefix_len) + middle_att_2d_masks = make_att_2d_masks(middle_pad_masks, middle_att_masks) + full_att_2d_masks = torch.cat([prefix_pad_2d_masks, middle_att_2d_masks], dim=2) + middle_position_ids = ( + torch.arange(1, middle_len + 1, device=max_prefix_position_ids.device) + .repeat(3, 1, 1) + .to(max_prefix_position_ids) + + max_prefix_position_ids + ) + full_att_2d_masks_4d = self._prepare_attention_masks_4d( + full_att_2d_masks, + dtype=middle_embs.dtype, + ) + (_, middle_out, _), past_key_values = self.qwen3_vl_with_expert.forward( + attention_mask=full_att_2d_masks_4d, + position_ids=middle_position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, middle_embs, None], + use_cache=True, + ) + + max_position_ids = middle_position_ids.max(dim=-1, keepdim=True).values + curr_pad_masks = torch.cat([prefix_pad_masks, middle_pad_masks], dim=1) + dt = torch.tensor(-1.0 / num_steps, dtype=torch.float32, device=device) + x_t = noise + time = torch.tensor(1.0, dtype=torch.float32, device=device) + suffix_static = None + if self.config.enable_suffix_static_context_optimization: + suffix_static = self.prepare_suffix_static_context(state, curr_pad_masks, max_position_ids) + while time >= -dt / 2: + expanded_time = time.expand(batch_size) + if suffix_static is None: + v_t = self.denoise_step( + state, + curr_pad_masks, + past_key_values, + max_position_ids, + x_t.to(dtype), + expanded_time.to(dtype), + ) + else: + v_t = self.denoise_step_optimized( + suffix_static, past_key_values, x_t.to(dtype), expanded_time.to(dtype) + ) + x_t = x_t + dt * v_t + time += dt + + return x_t, None if not decode_image else middle_out + + def denoise_step( + self, + state: torch.Tensor, + prefix_pad_masks: torch.Tensor, + past_key_values: Any, + max_prefix_position_ids: torch.Tensor, + x_t: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + suffix_embs, suffix_pad_masks, suffix_att_masks = self.embed_suffix(state, x_t, timestep) + suffix_len = suffix_pad_masks.shape[1] + batch_size = prefix_pad_masks.shape[0] + prefix_len = prefix_pad_masks.shape[1] + prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len) + suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks) + full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2) + position_ids = ( + torch.arange(1, suffix_len + 1, device=max_prefix_position_ids.device) + .repeat(3, 1, 1) + .to(max_prefix_position_ids) + + max_prefix_position_ids + ) + full_att_2d_masks_4d = self._prepare_attention_masks_4d( + full_att_2d_masks, + dtype=suffix_embs.dtype, + ) + outputs_embeds, _ = self.qwen3_vl_with_expert.forward( + attention_mask=full_att_2d_masks_4d, + position_ids=position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, None, suffix_embs], + use_cache=False, + ) + suffix_out = outputs_embeds[2][:, -self.config.chunk_size :].to(dtype=torch.float32) + return self.action_out_proj(suffix_out) + + def denoise_step_optimized( + self, + suffix_static: SuffixStaticContext, + past_key_values: Any, + x_t: torch.Tensor, + timestep: torch.Tensor, + ) -> torch.Tensor: + state_emb = suffix_static.state_emb + time_emb = create_sinusoidal_pos_embedding( + timestep, + self.action_in_proj.out_features, + min_period=self.config.min_period, + max_period=self.config.max_period, + device=timestep.device, + ).to(dtype=timestep.dtype) + action_emb = self.action_in_proj(x_t) + time_emb = time_emb[:, None, :].expand_as(action_emb) + action_time_emb = torch.cat([action_emb, time_emb], dim=2) + action_time_emb = self.action_time_mlp_out(F.silu(self.action_time_mlp_in(action_time_emb))) + suffix_embs = torch.cat([state_emb[:, None, :], action_time_emb], dim=1) + outputs_embeds, _ = self.qwen3_vl_with_expert.forward( + attention_mask=suffix_static.full_att_2d_masks_4d, + position_ids=suffix_static.position_ids, + past_key_values=past_key_values, + inputs_embeds=[None, None, suffix_embs], + use_cache=False, + ) + suffix_out = outputs_embeds[2][:, -self.config.chunk_size :].to(dtype=torch.float32) + return self.action_out_proj(suffix_out) + + +class InternVLAA1Policy(nn.Module): + _AUTO_WEIGHTS_IGNORE_PREFIXES = ["model.cosmos."] + + def __init__( + self, + config: InternVLAA1Config, + *, + processor_model_name: str = DEFAULT_QWEN3_VL_MODEL, + ) -> None: + super().__init__() + self.config = config + self.input_builder = Qwen3VLInputBuilder( + processor_model_name=processor_model_name, + max_length=config.tokenizer_max_length, + ) + cosmos_encoder_path, cosmos_decoder_path = resolve_cosmos_checkpoint_paths() + self.model = InternVLAA1( + config, + cosmos_encoder_path=cosmos_encoder_path, + cosmos_decoder_path=cosmos_decoder_path, + ) + self.model.to(config.device) + + @classmethod + def from_pretrained( + cls, + checkpoint_dir: str | Path, + *, + config: InternVLAA1Config | None = None, + processor_model_name: str = DEFAULT_QWEN3_VL_MODEL, + strict: bool = False, + ) -> InternVLAA1Policy: + checkpoint_dir = Path(checkpoint_dir) + if config is None: + config = InternVLAA1Config.from_pretrained(checkpoint_dir) + instance = cls(config, processor_model_name=processor_model_name) + checkpoint_path = checkpoint_dir / "model.safetensors" + if not checkpoint_path.exists(): + raise FileNotFoundError(f"InternVLAA1 checkpoint not found: {checkpoint_path}") + + loader = AutoWeightsLoader( + instance, + ignore_unexpected_prefixes=cls._AUTO_WEIGHTS_IGNORE_PREFIXES, + ) + with safe_open(str(checkpoint_path), framework="pt", device=str(config.device)) as f: + file_keys = list(f.keys()) + loaded = loader.load_weights((name, f.get_tensor(name)) for name in file_keys) + + if strict: + expected = { + name + for name, _ in instance.state_dict().items() + if not any(name.startswith(prefix) for prefix in cls._AUTO_WEIGHTS_IGNORE_PREFIXES) + } + missing = sorted(expected - loaded) + if missing: + preview = ", ".join(missing[:10]) + suffix = " ..." if len(missing) > 10 else "" + raise RuntimeError(f"Missing weights after AutoWeightsLoader load: {preview}{suffix}") + + instance.eval() + return instance + + def to(self, *args, **kwargs): # type: ignore[override] + super().to(*args, **kwargs) + self.model.cosmos.to(torch.bfloat16) + self.model.action_out_proj.to(torch.float32) + return self + + def _prepare_resized_histories(self, batch: dict[str, Any]) -> list[torch.Tensor]: + resized_images: list[torch.Tensor] = [] + for i in range(3): + image_history = batch[f"{OBS_IMAGES}.image{i}"] + if image_history.shape[0] != 1: + raise ValueError(f"InternVLAA1Policy only supports bs=1, got image batch {tuple(image_history.shape)}") + resized_images.append(resize_with_pad(image_history[0], self.config.image_resolution)) + return resized_images + + def _preprocess_images( + self, + batch: dict[str, torch.Tensor], + *, + resized_histories: list[torch.Tensor] | None = None, + ) -> tuple[torch.Tensor, torch.Tensor]: + if resized_histories is None: + resized_histories = self._prepare_resized_histories(batch) + images = torch.stack(resized_histories, dim=0).unsqueeze(0) + img_masks = torch.stack([batch[f"{OBS_IMAGES}.image{i}_mask"] for i in range(3)], dim=1) + return images, img_masks + + def prepare_state(self, batch: dict[str, torch.Tensor]) -> torch.Tensor: + return pad_vector(batch[OBS_STATE], self.config.max_state_dim) + + def _get_task(self, batch: dict[str, Any]) -> str: + task = batch[OBS_TASK] + if isinstance(task, str): + return task + if isinstance(task, (list, tuple)): + if len(task) != 1: + raise ValueError(f"InternVLAA1Policy only supports bs=1, got {len(task)} tasks") + item = task[0] + if not isinstance(item, str): + raise TypeError(f"Expected task string, got {type(item)!r}") + return item + raise TypeError(f"Unsupported task payload type: {type(task)!r}") + + def _prepare_qwen_prefix_inputs( + self, + resized_histories: list[torch.Tensor], + task: str, + ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: + qwen_inputs = self.input_builder.build(resized_histories, task) + device = self.model.action_in_proj.weight.device + return ( + qwen_inputs[f"{OBS_PREFIX}pixel_values"] + .unsqueeze(0) + .to(device=device, dtype=self.model.action_in_proj.weight.dtype), + qwen_inputs[f"{OBS_PREFIX}image_grid_thw"].unsqueeze(0).to(device=device), + qwen_inputs[f"{OBS_PREFIX}input_ids"].unsqueeze(0).to(device=device), + qwen_inputs[f"{OBS_PREFIX}attention_mask"].unsqueeze(0).to(device=device), + ) + + @torch.inference_mode() + def forward( + self, + batch: dict[str, Any], + *, + noise: torch.Tensor | None = None, + decode_image: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + task = self._get_task(batch) + resized_histories = self._prepare_resized_histories(batch) + pixel_values, image_grid_thw, lang_tokens, lang_masks = self._prepare_qwen_prefix_inputs( + resized_histories, + task, + ) + state = self.prepare_state(batch) + images, img_masks = self._preprocess_images(batch, resized_histories=resized_histories) + return self.model.sample_actions( + images, + img_masks, + pixel_values, + image_grid_thw, + lang_tokens, + lang_masks, + state, + noise=noise, + decode_image=decode_image, + ) diff --git a/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py b/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py new file mode 100644 index 000000000..b761ae8c6 --- /dev/null +++ b/vllm_omni/diffusion/models/internvla_a1/pipeline_internvla_a1.py @@ -0,0 +1,257 @@ +# SPDX-License-Identifier: Apache-2.0 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project +from __future__ import annotations + +import json +import os +from typing import Any + +import torch +from torch import nn +from vllm.logger import init_logger + +from vllm_omni.diffusion.compile import regionally_compile +from vllm_omni.diffusion.data import DiffusionOutput, OmniDiffusionConfig +from vllm_omni.diffusion.profiler.diffusion_pipeline_profiler import ( + DiffusionPipelineProfilerMixin, + wrap_methods_by_paths, +) +from vllm_omni.diffusion.request import OmniDiffusionRequest + +from .config import ( + DEFAULT_QWEN3_VL_MODEL, + OBS_IMAGES, + OBS_STATE, + OBS_TASK, + InternVLAA1Config, +) +from .model_internvla_a1 import InternVLAA1Policy, resolve_cosmos_checkpoint_paths + +logger = init_logger(__name__) + + +def get_internvla_a1_post_process_func(od_config: OmniDiffusionConfig): + del od_config + + def post_process_func(x): + return x + + return post_process_func + + +class InternVLAA1Pipeline(nn.Module, DiffusionPipelineProfilerMixin): + """InternVLA-A1 pipeline wrapper for the policy implementation.""" + + def __init__(self, *, od_config: OmniDiffusionConfig, prefix: str = ""): + super().__init__() + self.od_config = od_config + self.prefix = prefix + self.model_dir = od_config.model + self.config = self._build_config(od_config) + custom_args = od_config.custom_pipeline_args or {} + self.strict_load = bool(custom_args.get("strict_load", False)) + self.processor_model_name = str(custom_args.get("processor_model_name", DEFAULT_QWEN3_VL_MODEL)) + enable_warmup = custom_args.get("enable_warmup") + self.enable_warmup = bool(enable_warmup) if isinstance(enable_warmup, bool) else False + + self.setup_diffusion_pipeline_profiler( + profiler_targets=["forward"], + enable_diffusion_pipeline_profiler=self.od_config.enable_diffusion_pipeline_profiler, + ) + self.policy = self._initialize_policy() + self._setup_policy_profiler_targets() + if self.enable_warmup: + self._warmup() + + def _build_config(self, od_config: OmniDiffusionConfig) -> InternVLAA1Config: + config_dict = self._load_config_dict(od_config) + config = InternVLAA1Config.from_model_config(config_dict) + + custom_args = od_config.custom_pipeline_args or {} + device = custom_args.get("device") + if isinstance(device, str): + config.device = device + + dtype = custom_args.get("dtype") + if isinstance(dtype, str): + config.dtype = dtype + elif od_config.dtype is not None: + config.dtype = str(od_config.dtype).split(".")[-1] + + compile_model = custom_args.get("compile_model") + if isinstance(compile_model, bool): + config.compile_model = compile_model + + attn_implementation = custom_args.get("attn_implementation") + if isinstance(attn_implementation, str): + config.attn_implementation = attn_implementation + + enable_regional_compile = custom_args.get("enable_regional_compile") + if isinstance(enable_regional_compile, bool): + config.enable_regional_compile = enable_regional_compile + + regional_compile_dynamic = custom_args.get("regional_compile_dynamic") + if isinstance(regional_compile_dynamic, bool): + config.regional_compile_dynamic = regional_compile_dynamic + + return config + + def _load_config_dict(self, od_config: OmniDiffusionConfig) -> dict[str, Any]: + if od_config.model_config: + return dict(od_config.model_config) + + model_path = od_config.model + if not model_path: + return {} + + config_path = os.path.join(model_path, "config.json") + if not os.path.exists(config_path): + logger.info("InternVLAA1Pipeline config.json not found under %s; using defaults.", model_path) + return {} + + with open(config_path, encoding="utf-8") as f: + return json.load(f) + + def has_real_checkpoint(self) -> bool: + return bool(self.model_dir) and os.path.exists(os.path.join(self.model_dir, "model.safetensors")) + + def runtime_mode(self) -> str: + return "real_checkpoint_loaded" if self.has_real_checkpoint() else "no_checkpoint_policy" + + def _setup_policy_profiler_targets(self) -> None: + if not self.od_config.enable_diffusion_pipeline_profiler: + return + + wrap_methods_by_paths( + self, + [ + "policy.model.sample_actions", + "policy.model.embed_prefix", + "policy.model.embed_middle", + "policy.model.denoise_step", + ], + ) + + def _apply_policy_optimizations(self, policy: InternVLAA1Policy) -> None: + policy.model.set_attention_implementation(self.config.attn_implementation) + if not self.config.enable_regional_compile: + return + + compile_targets = [ + "qwen3_vl_with_expert.und_expert.visual", + "qwen3_vl_with_expert.und_expert.language_model", + "qwen3_vl_with_expert.gen_expert", + "qwen3_vl_with_expert.act_expert", + ] + + for path in compile_targets: + current = policy.model + for part in path.split("."): + current = getattr(current, part, None) + if current is None: + break + if current is None: + continue + try: + regionally_compile(current, dynamic=self.config.regional_compile_dynamic) + logger.info("InternVLAA1Pipeline regional compile applied to %s", path) + except Exception as exc: + logger.warning("InternVLAA1Pipeline regional compile failed for %s: %s", path, exc) + + def _initialize_policy(self) -> InternVLAA1Policy: + resolve_cosmos_checkpoint_paths() + if self.has_real_checkpoint(): + logger.info("Loading InternVLA-A1 weights from %s", self.model_dir) + policy = InternVLAA1Policy.from_pretrained( + self.model_dir, + config=self.config, + processor_model_name=self.processor_model_name, + strict=self.strict_load, + ) + else: + logger.info("Initializing InternVLA-A1 policy without checkpoint weights.") + policy = InternVLAA1Policy(self.config, processor_model_name=self.processor_model_name) + + policy.to(self.config.device) + policy.to(getattr(torch, self.config.dtype)) + policy.eval() + self._apply_policy_optimizations(policy) + return policy + + def _build_fake_batch_inputs(self) -> dict[str, torch.Tensor]: + device = torch.device(self.config.device) + dtype = getattr(torch, self.config.dtype) + history = 2 + channels = 3 + image_size = self.config.image_resolution + fake_image = torch.zeros((1, history, channels, image_size[0], image_size[1]), device=device, dtype=dtype) + fake_mask = torch.ones((1,), device=device, dtype=torch.bool) + return { + OBS_STATE: torch.zeros((1, self.config.max_state_dim), device=device, dtype=dtype), + OBS_TASK: [""], + f"{OBS_IMAGES}.image0": fake_image.clone(), + f"{OBS_IMAGES}.image1": fake_image.clone(), + f"{OBS_IMAGES}.image2": fake_image.clone(), + f"{OBS_IMAGES}.image0_mask": fake_mask.clone(), + f"{OBS_IMAGES}.image1_mask": fake_mask.clone(), + f"{OBS_IMAGES}.image2_mask": fake_mask.clone(), + } + + def _warmup(self) -> None: + logger.info("InternVLAA1Pipeline warmup started") + try: + batch_inputs = self._build_fake_batch_inputs() + noise = torch.zeros( + (1, self.config.chunk_size, self.config.max_action_dim), + device=self.config.device, + dtype=torch.float32, + ) + with torch.inference_mode(): + self.policy.forward(batch_inputs, noise=noise, decode_image=False) + except Exception as exc: + logger.warning("InternVLAA1Pipeline warmup failed: %s", exc) + return + logger.info("InternVLAA1Pipeline warmup finished") + + def _predict_actions( + self, + batch_inputs: dict[str, Any], + *, + noise: torch.Tensor | None = None, + decode_image: bool = False, + ) -> tuple[torch.Tensor, torch.Tensor | None]: + logger.debug("InternVLAA1Pipeline forward mode=%s", self.runtime_mode()) + return self.policy.forward( + batch_inputs, + noise=noise, + decode_image=decode_image, + ) + + @torch.inference_mode() + def forward(self, req: OmniDiffusionRequest) -> DiffusionOutput: + if len(req.prompts) > 1: + logger.warning("InternVLAA1Pipeline only supports a single prompt/request; taking the first sample.") + extra_args = getattr(req.sampling_params, "extra_args", {}) or {} + batch_inputs = extra_args.get("batch_inputs") + if batch_inputs is None: + return DiffusionOutput( + error=( + "InternVLAA1Pipeline.forward expects sampling_params.extra_args['batch_inputs'] " + "with pre-built repo-side inputs." + ), + post_process_func=get_internvla_a1_post_process_func(self.od_config), + ) + + output, decoded = self._predict_actions( + batch_inputs, + noise=extra_args.get("noise"), + decode_image=bool(extra_args.get("decode_image", False)), + ) + custom_output: dict[str, Any] = {} + if decoded is not None: + custom_output["decoded"] = decoded + return DiffusionOutput( + output=output, + custom_output=custom_output, + post_process_func=get_internvla_a1_post_process_func(self.od_config), + ) diff --git a/vllm_omni/diffusion/registry.py b/vllm_omni/diffusion/registry.py index ebcd263c1..37f519944 100644 --- a/vllm_omni/diffusion/registry.py +++ b/vllm_omni/diffusion/registry.py @@ -135,6 +135,11 @@ _DIFFUSION_MODELS = { "pipeline_bagel", "BagelPipeline", ), + "InternVLAA1Pipeline": ( + "internvla_a1", + "pipeline_internvla_a1", + "InternVLAA1Pipeline", + ), "LongCatImageEditPipeline": ( "longcat_image", "pipeline_longcat_image_edit", @@ -421,6 +426,7 @@ _DIFFUSION_POST_PROCESS_FUNCS = { "WanI2VDMD2Pipeline": "get_wan22_i2v_post_process_func", "LongCatImagePipeline": "get_longcat_image_post_process_func", "BagelPipeline": "get_bagel_post_process_func", + "InternVLAA1Pipeline": "get_internvla_a1_post_process_func", "LongCatImageEditPipeline": "get_longcat_image_post_process_func", "StableDiffusion3Pipeline": "get_sd3_image_post_process_func", "FluxKontextPipeline": "get_flux_kontext_post_process_func",